from odoo import models, api
from odoo.tools.float_utils import float_round

class MrpProductionSchedule(models.Model):
    _inherit = 'mrp.production.schedule'

    def set_replenish_qty(self, date_index, quantity, period_scale=False):
        """ Save the replenish quantity and mark the cells as manually updated.
        We override this to provide two-way (Trace up and Drill down) updating 
        of linked components in the BOM hierarchy.
        """
        self.ensure_one()
        
        # Calculate the difference between new quantity and old quantity
        date_start, date_stop = self.company_id._get_date_range(force_period=period_scale)[date_index]
        existing_forecast = self.forecast_ids.filtered(lambda f:
            f.date >= date_start and f.date <= date_stop)
            
        old_qty = sum(existing_forecast.mapped('replenish_qty')) if existing_forecast else 0.0
        
        # Call super first to update this component
        res = super().set_replenish_qty(date_index, quantity, period_scale)
        
        quantity = float_round(float(quantity), precision_rounding=self.product_uom_id.rounding)
        diff_qty = quantity - old_qty
        
        # If the quantity changed and we're not already propagating an update to avoid infinite loops
        if diff_qty != 0 and not self.env.context.get('skip_two_way_mps'):
            self._propagate_replenish_qty(diff_qty, date_index, period_scale)

        return res

    def _propagate_replenish_qty(self, diff_qty, date_index, period_scale=False):
        """Propagate replenishment difference to parents (bottom-up) and children (top-down)"""
        
        # We need to trace up (bottom-up)
        parent_schedules = self._get_impacted_parent_schedules()
        # We need to drill down (top-down) 
        child_schedules = self._get_impacted_child_schedules()
        
        # Fetch the BOM hierarchy relationships
        schedules_to_compute = parent_schedules | child_schedules | self
        indirect_demand_trees = schedules_to_compute._get_indirect_demand_tree()
        indirect_ratio_mps = schedules_to_compute._get_indirect_demand_ratio_mps(indirect_demand_trees)
        
        # TRACE UP: Update parents
        for parent_schedule in parent_schedules:
            # Find the ratio: how many of 'self.product_id' are needed to make 1 'parent_schedule.product_id'
            ratios = indirect_ratio_mps.get((parent_schedule.warehouse_id, parent_schedule.product_id), {})
            ratio = ratios.get(self.product_id, 0.0)
            
            if ratio > 0:
                parent_diff = diff_qty / ratio
                
                # Retrieve the existing parent quantity for this period
                p_date_start, p_date_stop = parent_schedule.company_id._get_date_range(force_period=period_scale)[date_index]
                p_exist = parent_schedule.forecast_ids.filtered(lambda f: f.date >= p_date_start and f.date <= p_date_stop)
                p_old_qty = sum(p_exist.mapped('replenish_qty')) if p_exist else 0.0
                
                new_parent_qty = p_old_qty + parent_diff
                if new_parent_qty < 0:
                    new_parent_qty = 0.0
                
                # Update parent with skip context to prevent loops
                parent_schedule.with_context(skip_two_way_mps=True).set_replenish_qty(date_index, new_parent_qty, period_scale)
        
        # DRILL DOWN: Update children
        for child_schedule in child_schedules:
            # Find the ratio: how many of 'child.product_id' base units are needed to make 1 'self.product_id' base unit
            ratios = indirect_ratio_mps.get((self.warehouse_id, self.product_id), {})
            ratio = ratios.get(child_schedule.product_id, 0.0)
            
            if ratio > 0:
                # child_diff is in base units
                child_diff = diff_qty * ratio
                
                # If the packaging module is used, we must convert the base unit diff_qty 
                # into the packaging unit diff_qty.
                # Because child_schedule.set_replenish_qty will do: quantity = float(quantity) * self.packaging_id.qty
                # Therefore, we must pass the quantity *in packaging units*.
                if hasattr(child_schedule, 'packaging_id') and child_schedule.packaging_id:
                    packaging_qty = child_schedule.packaging_id.qty
                    if packaging_qty:
                        child_diff = child_diff / packaging_qty
                
                c_date_start, c_date_stop = child_schedule.company_id._get_date_range(force_period=period_scale)[date_index]
                c_exist = child_schedule.forecast_ids.filtered(lambda f: f.date >= c_date_start and f.date <= c_date_stop)
                
                # The existing replenish_qty in the database is actually stored in base units.
                # BUT we need to pass the new value to set_replenish_qty in the *packaging* units if packaging is used.
                c_old_qty_base = sum(c_exist.mapped('replenish_qty')) if c_exist else 0.0
                
                if hasattr(child_schedule, 'packaging_id') and child_schedule.packaging_id and child_schedule.packaging_id.qty:
                    c_old_qty_pack = c_old_qty_base / child_schedule.packaging_id.qty
                else:
                    c_old_qty_pack = c_old_qty_base
                
                new_child_qty_pack = c_old_qty_pack + child_diff
                if new_child_qty_pack < 0:
                    new_child_qty_pack = 0.0
                
                child_schedule.with_context(skip_two_way_mps=True).set_replenish_qty(date_index, new_child_qty_pack, period_scale)