from odoo import api, fields, models
from odoo.exceptions import UserError
from odoo.fields import Domain
import logging

_logger = logging.getLogger(__name__)

# Log when this module is loaded
_logger.info("="*80)
_logger.info("STOCK_RESTRICT_SOURCE_LOCATION: stock_location.py module is being loaded!")
_logger.info("="*80)

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

    @api.model
    def _get_gather_domain(self, product_id, location_id, lot_id=None, package_id=None, owner_id=None, strict=False):
        """Override to apply location restrictions during reservation (gather)"""
        ctx = self.env.context
        domain = super()._get_gather_domain(product_id, location_id, lot_id, package_id, owner_id, strict)
        
        # Try to find picking type from context or active_move_id
        picking_type_id = ctx.get('default_picking_type_id')
        
        if not picking_type_id:
            active_move_id = ctx.get('active_move_id')
            if active_move_id:
                move = self.env['stock.move'].browse(active_move_id)
                if move.exists():
                    picking_type_id = move.picking_type_id.id

        if not picking_type_id:
            active_mo_id = ctx.get('active_mo_id')
            if active_mo_id:
                mo = self.env['mrp.production'].browse(active_mo_id)
                if mo.exists():
                    picking_type_id = mo.picking_type_id.id
        
        if picking_type_id:
            picking_type = self.env['stock.picking.type'].browse(picking_type_id)
            if picking_type.default_location_src_ids:
                allowed_location_ids = picking_type.default_location_src_ids.ids
                # Restrict the domain to allowed source locations
                domain = ['&'] + list(domain) + [('location_id', 'in', allowed_location_ids)]
        
        return domain

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

    @api.model
    def _name_search(self, name='', domain=None, operator='ilike', limit=None, order=None):
        """Override to restrict locations based on picking_type_id in context"""
        _logger.info(f"LOCATION SEARCH CALLED: name={name}, context keys={list(self.env.context.keys())}, domain={domain}")
        domain = domain or []
        
        # Check if we have a picking_type_id in context (from MO or picking)
        picking_type_id = self.env.context.get('default_picking_type_id')
        
        if picking_type_id:
            picking_type = self.env['stock.picking.type'].browse(picking_type_id)
            if picking_type.default_location_src_ids:
                allowed_location_ids = picking_type.default_location_src_ids.ids
                _logger.info(f"LOCATION_RESTRICT LOCATION SEARCH: Filtering locations to {allowed_location_ids} for picking type {picking_type.name}")
                # Add restriction to domain
                domain = domain + [('id', 'in', allowed_location_ids)]
        else:
            _logger.info(f"LOCATION SEARCH: No picking_type_id in context")
        
        return super()._name_search(name=name, domain=domain, operator=operator, limit=limit, order=order)

_logger.info("="*80)
_logger.info("STOCK_RESTRICT_SOURCE_LOCATION: stock_location.py module loaded successfully!")
_logger.info("="*80)