import logging
from odoo import api, fields, models, _
from odoo.osv import expression
from odoo.exceptions import UserError, ValidationError

_logger = logging.getLogger(__name__)

class StockLocation(models.Model):
    _inherit = 'stock.location'

    def _get_allowed_locations(self):
        """
        Helper to retrieve allowed locations based on the current context (picking type).
        Used by UI-level overrides in StockQuant and StockLot.
        Returns:
            list: IDs of allowed locations.
            False: If no restriction should be applied (e.g. no picking type found).
        """
        ctx = self.env.context
        # 1. Identify the Picking Type (Operation Type)
        # Check multiple common context keys used by stock/mrp models
        picking_type_id = (ctx.get('default_picking_type_id') or 
                          ctx.get('picking_type_id') or 
                          ctx.get('active_picking_type_id'))
        
        # 2. Advanced MO Support: find picking type from MO ID if needed
        mo_id = (ctx.get('active_mo_id') or 
                 ctx.get('default_mo_id') or 
                 ctx.get('default_production_id') or 
                 ctx.get('mo_id') or 
                 ctx.get('production_id'))
                 
        if not picking_type_id and mo_id:
            mo = self.env['mrp.production'].browse(mo_id)
            if mo.exists():
                picking_type_id = mo.picking_type_id.id

        if not picking_type_id:
            # NO identification of operation type = NO restriction (safe fallback)
            return False

        # 3. Retrieve allowed locations from the picking type
        picking_type = self.env['stock.picking.type'].browse(picking_type_id)
        if picking_type.exists():
            # Priority 1: Use the M2M "Allowed Source Locations" if set
            if picking_type.default_location_src_ids:
                return picking_type.default_location_src_ids.ids
            
            # Priority 2: Fallback to the primary "Default Source Location" (M2O)
            if picking_type.default_location_src_id:
                return [picking_type.default_location_src_id.id]
            
        return False

class StockQuant(models.Model):
    _inherit = 'stock.quant'

    @api.model
    def web_search_read(self, domain, specification, offset=0, limit=None, order=None, count_limit=None):
        """
        UI-SURFACE OVERRIDE: Applies location filtering ONLY for the web interface.
        """
        ctx = self.env.context
        if not ctx.get('skip_location_restriction') and ctx.get('uid'):
            allowed_location_ids = self.env['stock.location']._get_allowed_locations()
            if allowed_location_ids:
                # Add location filter to the domain ONLY if we have an explicit list
                domain = expression.AND([domain, [('location_id', 'in', allowed_location_ids)]])
                
        return super().web_search_read(domain, specification, offset=offset, limit=limit, order=order, count_limit=count_limit)

class StockLot(models.Model):
    _inherit = 'stock.lot'

    @api.model
    def name_search(self, name='', args=None, operator='ilike', limit=100):
        """
        UI-SURFACE OVERRIDE: Filters the many2one lot selection dropdown.
        """
        ctx = self.env.context
        if not ctx.get('skip_location_restriction') and ctx.get('uid'):
            allowed_location_ids = self.env['stock.location']._get_allowed_locations()
            
            # If our custom helper failed to find a picking-type restriction,
            # we check if a simple default_location_id filter was passed directly.
            if not allowed_location_ids and ctx.get('default_location_id'):
                allowed_location_ids = [ctx.get('default_location_id')]

            if allowed_location_ids:
                quant_domain = [
                    ('location_id', 'in', allowed_location_ids),
                    ('quantity', '>', 0),
                    ('lot_id', '!=', False)
                ]
                product_id = ctx.get('default_product_id')
                if product_id:
                    quant_domain.append(('product_id', '=', product_id))
                
                quants = self.env['stock.quant'].with_context(skip_location_restriction=True).sudo().search(quant_domain)
                lot_ids = quants.mapped('lot_id').ids
                
                args = expression.AND([args or [], [('id', 'in', lot_ids)]])

        return super().name_search(name, args=args, operator=operator, limit=limit)

    @api.model
    def web_search_read(self, domain, specification, offset=0, limit=None, order=None, count_limit=None):
        """
        UI-SURFACE OVERRIDE: Applies filtering for the Lot Catalog and list views.
        """
        ctx = self.env.context
        if not ctx.get('skip_location_restriction') and ctx.get('uid'):
            allowed_location_ids = self.env['stock.location']._get_allowed_locations()
            
            if not allowed_location_ids and ctx.get('default_location_id'):
                allowed_location_ids = [ctx.get('default_location_id')]

            if allowed_location_ids:
                quant_domain = [
                    ('location_id', 'in', allowed_location_ids),
                    ('quantity', '>', 0),
                    ('lot_id', '!=', False)
                ]
                product_id = ctx.get('default_product_id')
                if product_id:
                    quant_domain.append(('product_id', '=', product_id))

                quants = self.env['stock.quant'].with_context(skip_location_restriction=True).sudo().search(quant_domain)
                lot_ids = quants.mapped('lot_id').ids
                
                domain = expression.AND([domain, [('id', 'in', lot_ids)]])

        return super().web_search_read(domain, specification, offset=offset, limit=limit, order=order, count_limit=count_limit)