"""
Inventory Client for direct Redis operations with atomic LUA scripts.
Provides sub-millisecond inventory reservation and release operations.
"""

from typing import Literal, Optional, List, Dict, Any, Callable, Awaitable
import json
import logging

from .lua_scripts import (
    RESERVE_SCRIPT,
    RELEASE_SCRIPT,
    SET_INVENTORY_SCRIPT,
    CHECK_AVAILABLE_SCRIPT,
    GET_INVENTORY_STATE_SCRIPT,
    BULK_SET_INVENTORY_SCRIPT,
)
from .models import (
    ReserveResult,
    ReleaseResult,
    SetInventoryResult,
    InventoryState,
    InsufficientInventoryError,
)

logger = logging.getLogger(__name__)

ChannelSource = Literal["medusa", "shopify"]


class InventoryClient:
    """
    Direct Redis inventory operations with atomic LUA scripts.
    Sub-millisecond latency for reservation operations.
    
    Usage:
        from helpr.inventory import InventoryClient
        from helpr.cache import AsyncRedisHelper
        
        redis_helper = AsyncRedisHelper(redis_client)
        inventory = InventoryClient(redis_helper)
        
        # Reserve inventory
        result = await inventory.reserve("SKU123", "WH001", 5, "medusa")
        if result.success:
            print(f"Reserved! Available: {result.available}")
        else:
            print(f"Failed: {result.reason}")
        
        # Release on fulfillment
        await inventory.release("SKU123", "WH001", 5, "medusa")
    """
    
    KEY_PREFIX = "inv"
    
    def __init__(self, redis_helper, event_callback: Callable[..., Awaitable[None]] = None):
        """
        Initialize with an AsyncRedisHelper instance.
        
        Args:
            redis_helper: AsyncRedisHelper instance with redis_read_client and redis_write_client
            event_callback: Optional async callback for persisting events to database.
                           Signature: async def callback(event_type, event_data) -> None
        """
        self.redis = redis_helper
        self.event_callback = event_callback
        self._script_shas: Dict[str, str] = {}
        self._scripts_registered = False
    
    def _build_key(self, warehouse_id: str, sku: str) -> str:
        """Build Redis key for inventory hash."""
        return f"{self.KEY_PREFIX}:{warehouse_id}:{sku}"
    
    async def _ensure_scripts_registered(self) -> None:
        """Register LUA scripts with Redis on first use. Scripts are cached by SHA."""
        if self._scripts_registered:
            return
        
        try:
            write_client = self.redis.redis_write_client
            
            self._script_shas["reserve"] = await write_client.script_load(RESERVE_SCRIPT)
            self._script_shas["release"] = await write_client.script_load(RELEASE_SCRIPT)
            self._script_shas["set_inventory"] = await write_client.script_load(SET_INVENTORY_SCRIPT)
            self._script_shas["check_available"] = await write_client.script_load(CHECK_AVAILABLE_SCRIPT)
            self._script_shas["get_state"] = await write_client.script_load(GET_INVENTORY_STATE_SCRIPT)
            self._script_shas["bulk_set"] = await write_client.script_load(BULK_SET_INVENTORY_SCRIPT)
            
            self._scripts_registered = True
            logger.debug("Inventory LUA scripts registered with Redis")
        except Exception as e:
            logger.error(f"Failed to register LUA scripts: {e}")
            raise
    
    async def reserve(
        self,
        sku: str,
        warehouse_id: str,
        quantity: int,
        source: ChannelSource,
        order_id: Optional[str] = None,
        order_display_id: Optional[str] = None,
        platform_order_id: Optional[str] = None,
        line_item_id: Optional[str] = None,
        raise_on_failure: bool = False,
    ) -> ReserveResult:
        """
        Atomically reserve inventory for an order.
        
        Args:
            sku: Product SKU
            warehouse_id: ERP warehouse ID
            quantity: Amount to reserve (must be positive)
            source: Channel making the reservation ("medusa" or "shopify")
            order_id: Optional order ID for logging
            raise_on_failure: If True, raises InsufficientInventoryError on failure
            
        Returns:
            ReserveResult with success status and current inventory state
            
        Raises:
            InsufficientInventoryError: If raise_on_failure=True and reservation fails
        """
        await self._ensure_scripts_registered()
        
        key = self._build_key(warehouse_id, sku)
        
        try:
            result_json = await self.redis.redis_write_client.evalsha(
                self._script_shas["reserve"],
                1,
                key,
                quantity,
                source
            )
            result = ReserveResult.from_json(result_json)
            
            if result.success:
                logger.info(
                    f"Reserved {quantity} of {sku} in {warehouse_id} for {source}",
                    extra={
                        "sku": sku,
                        "warehouse_id": warehouse_id,
                        "quantity": quantity,
                        "source": source,
                        "order_id": order_id,
                        "available_after": result.available,
                    }
                )
            else:
                logger.warning(
                    f"Reservation failed for {sku}: {result.reason}",
                    extra={
                        "sku": sku,
                        "warehouse_id": warehouse_id,
                        "quantity": quantity,
                        "source": source,
                        "order_id": order_id,
                        "available": result.available,
                        "reason": result.reason,
                    }
                )
                
                if raise_on_failure:
                    raise InsufficientInventoryError(
                        sku=sku,
                        warehouse_id=warehouse_id,
                        requested=quantity,
                        available=result.available,
                    )
            
            # Fire event callback if configured (non-blocking)
            if self.event_callback and result.success:
                try:
                    await self.event_callback(
                        event_type="reserved",
                        sku=sku,
                        warehouse_id=warehouse_id,
                        quantity=quantity,
                        source=source,
                        order_id=order_id,
                        order_display_id=order_display_id,
                        platform_order_id=platform_order_id,
                        line_item_id=line_item_id,
                        available_after=result.available,
                        reserved_after=result.reserved,
                    )
                except Exception as cb_err:
                    logger.warning(f"Event callback failed: {cb_err}")
            
            return result
            
        except InsufficientInventoryError:
            raise
        except Exception as e:
            logger.error(f"Reserve operation failed: {e}", extra={"sku": sku, "warehouse_id": warehouse_id})
            raise
    
    async def release(
        self,
        sku: str,
        warehouse_id: str,
        quantity: int,
        source: ChannelSource,
        reason: str = "fulfilled",
        order_id: Optional[str] = None,
        order_display_id: Optional[str] = None,
        platform_order_id: Optional[str] = None,
        line_item_id: Optional[str] = None,
    ) -> ReleaseResult:
        """
        Release reserved inventory back to available.
        
        Args:
            sku: Product SKU
            warehouse_id: ERP warehouse ID
            quantity: Amount to release (must be positive)
            source: Channel releasing the reservation ("medusa" or "shopify")
            reason: Why releasing ("fulfilled", "cancelled", "expired")
            order_id: Optional order ID for logging
            
        Returns:
            ReleaseResult with release details and current inventory state
        """
        await self._ensure_scripts_registered()
        
        key = self._build_key(warehouse_id, sku)
        
        try:
            result_json = await self.redis.redis_write_client.evalsha(
                self._script_shas["release"],
                1,
                key,
                quantity,
                source,
                reason  # Pass reason to distinguish fulfillment vs cancellation
            )
            result = ReleaseResult.from_json(result_json)
            
            logger.info(
                f"Released {result.released} of {sku} in {warehouse_id} for {source} ({reason})",
                extra={
                    "sku": sku,
                    "warehouse_id": warehouse_id,
                    "quantity_requested": quantity,
                    "quantity_released": result.released,
                    "source": source,
                    "reason": reason,
                    "order_id": order_id,
                    "available_after": result.available,
                }
            )
            
            # Fire event callback if configured (non-blocking)
            if self.event_callback and result.released > 0:
                try:
                    await self.event_callback(
                        event_type="released",
                        sku=sku,
                        warehouse_id=warehouse_id,
                        quantity=result.released,
                        source=source,
                        reason=reason,
                        order_id=order_id,
                        order_display_id=order_display_id,
                        platform_order_id=platform_order_id,
                        line_item_id=line_item_id,
                        available_after=result.available,
                        reserved_after=result.reserved,
                    )
                except Exception as cb_err:
                    logger.warning(f"Event callback failed: {cb_err}")
            
            return result
            
        except Exception as e:
            logger.error(f"Release operation failed: {e}", extra={"sku": sku, "warehouse_id": warehouse_id})
            raise
    
    async def check_available(
        self,
        sku: str,
        warehouse_id: str,
    ) -> int:
        """
        Get current available quantity for a SKU/warehouse.
        This is a non-blocking read operation.
        
        Args:
            sku: Product SKU
            warehouse_id: ERP warehouse ID
            
        Returns:
            Available quantity (0 if not found)
        """
        key = self._build_key(warehouse_id, sku)
        result = await self.redis.redis_read_client.hget(key, "available")
        return int(result) if result else 0
    
    async def check_available_aggregate(
        self,
        sku: str,
        warehouse_ids: List[str],
    ) -> int:
        """
        Get total available quantity across multiple warehouses.
        
        Args:
            sku: Product SKU
            warehouse_ids: List of ERP warehouse IDs
            
        Returns:
            Total available quantity across all warehouses
        """
        total = 0
        for wh_id in warehouse_ids:
            total += await self.check_available(sku, wh_id)
        return total
    
    async def set_inventory(
        self,
        sku: str,
        warehouse_id: str,
        erp_quantity: int,
        event_id: Optional[str] = None,
    ) -> SetInventoryResult:
        """
        Set ERP inventory quantity and recalculate available.
        Called by Storehouse when ERP sends inventory updates.
        
        Available is calculated as: erp_quantity - total_reserved
        This preserves existing reservations while updating the base quantity.
        
        Args:
            sku: Product SKU
            warehouse_id: ERP warehouse ID
            erp_quantity: Raw quantity from ERP (must be non-negative)
            
        Returns:
            SetInventoryResult with new inventory state
        """
        await self._ensure_scripts_registered()
        
        key = self._build_key(warehouse_id, sku)
        
        try:
            result_json = await self.redis.redis_write_client.evalsha(
                self._script_shas["set_inventory"],
                1,
                key,
                erp_quantity
            )
            result = SetInventoryResult.from_json(result_json)
            
            logger.info(
                f"Set inventory for {sku} in {warehouse_id}: erp={erp_quantity}, available={result.available}",
                extra={
                    "sku": sku,
                    "warehouse_id": warehouse_id,
                    "erp_quantity": erp_quantity,
                    "available": result.available,
                    "reserved_medusa": result.reserved_medusa,
                    "reserved_shopify": result.reserved_shopify,
                }
            )
            
            # Fire event callback if configured (non-blocking)
            if self.event_callback:
                try:
                    await self.event_callback(
                        event_type="erp_update",
                        sku=sku,
                        warehouse_id=warehouse_id,
                        erp_quantity=erp_quantity,
                        available_after=result.available,
                        reserved_medusa_after=result.reserved_medusa,
                        reserved_shopify_after=result.reserved_shopify,
                        event_id=event_id,
                    )
                except Exception as cb_err:
                    logger.warning(f"Event callback failed: {cb_err}")
            
            return result
            
        except Exception as e:
            logger.error(f"Set inventory failed: {e}", extra={"sku": sku, "warehouse_id": warehouse_id})
            raise
    
    async def get_state(
        self,
        sku: str,
        warehouse_id: str,
    ) -> InventoryState:
        """
        Get full inventory state for debugging and monitoring.
        
        Args:
            sku: Product SKU
            warehouse_id: ERP warehouse ID
            
        Returns:
            InventoryState with all inventory fields
        """
        await self._ensure_scripts_registered()
        
        key = self._build_key(warehouse_id, sku)
        
        result_json = await self.redis.redis_write_client.evalsha(
            self._script_shas["get_state"],
            1,
            key
        )
        
        return InventoryState.from_json(sku, warehouse_id, result_json)
    
    async def bulk_set_inventory(
        self,
        items: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """
        Set inventory for multiple SKUs in a single atomic operation.
        More efficient than individual set_inventory calls for bulk ERP updates.
        
        Args:
            items: List of {"sku": str, "warehouse_id": str, "quantity": int}
            
        Returns:
            List of results with success status per item
        """
        await self._ensure_scripts_registered()
        
        if not items:
            return []
        
        keys = [self._build_key(item["warehouse_id"], item["sku"]) for item in items]
        quantities = [item["quantity"] for item in items]
        
        try:
            result_json = await self.redis.redis_write_client.evalsha(
                self._script_shas["bulk_set"],
                len(keys),
                *keys,
                *quantities
            )
            results = json.loads(result_json)
            
            success_count = sum(1 for r in results if r.get("success"))
            logger.info(
                f"Bulk set inventory: {success_count}/{len(items)} successful",
                extra={"total_items": len(items), "success_count": success_count}
            )
            
            return results
            
        except Exception as e:
            logger.error(f"Bulk set inventory failed: {e}")
            raise
