"""
CXR Model Registry

Central registry for managing and routing to different CXR kernel models.
Supports model variants: k1, k2, sim1, lite1, audit1
"""

from typing import Dict, Optional, Type, Protocol
from enum import Enum
import structlog

logger = structlog.get_logger(__name__)


class ModelVariant(str, Enum):
    """Available CXR model variants"""
    K1 = "cxr-k1"  # Flagship: Full 7-plane fusion, all contexts
    K2 = "cxr-k2"  # Premium: Enhanced precision, slower
    SIM1 = "cxr-k-sim1"  # Simulation: Monte Carlo integration
    LITE1 = "cxr-k-lite1"  # Speed: Fast, reduced contexts
    AUDIT1 = "cxr-k-audit1"  # Audit: Compliance-focused


class ModelCapabilities:
    """Describes capabilities of a model variant"""
    def __init__(
        self,
        fusion_planes: int = 3,
        uses_monte_carlo: bool = False,
        uses_all_contexts: bool = True,
        speed_profile: str = "balanced",  # "fast", "balanced", "thorough"
        precision_level: str = "standard",  # "standard", "enhanced", "maximum"
        compliance_focus: bool = False,
    ):
        self.fusion_planes = fusion_planes
        self.uses_monte_carlo = uses_monte_carlo
        self.uses_all_contexts = uses_all_contexts
        self.speed_profile = speed_profile
        self.precision_level = precision_level
        self.compliance_focus = compliance_focus


class CXRModel(Protocol):
    """Protocol defining the interface all CXR models must implement"""
    
    def analyze(self, claim: Dict, tenant_id: str) -> Dict:
        """
        Analyze a claim and return results.
        
        Args:
            claim: Claim payload dictionary
            tenant_id: Tenant identifier for multi-tenancy
            
        Returns:
            Dictionary with analysis results including:
            - verdict: str (approved/denied/pending_review)
            - explanation: Dict with summary and drivers
            - financialImpact: Dict with projectedRecovery, currency
            - reportId: Optional[str]
        """
        ...


class ModelRegistry:
    """
    Central registry for CXR kernel models.
    
    Usage:
        registry = ModelRegistry()
        model = registry.get_model("cxr-k1")
        result = model.analyze(claim, tenant_id)
    """
    
    def __init__(self):
        self._models: Dict[str, Type[CXRModel]] = {}
        self._instances: Dict[str, CXRModel] = {}
        self._capabilities: Dict[str, ModelCapabilities] = {}
        self._default_model = ModelVariant.K1.value
        
    def register(
        self,
        model_id: str,
        model_class: Type[CXRModel],
        capabilities: ModelCapabilities,
        singleton: bool = True,
    ) -> None:
        """
        Register a model variant.
        
        Args:
            model_id: Unique identifier (e.g., "cxr-k1")
            model_class: Class implementing CXRModel protocol
            capabilities: Model capabilities description
            singleton: If True, create single instance and reuse it
        """
        if model_id in self._models:
            logger.warning("model.reregistered", model_id=model_id)
        
        self._models[model_id] = model_class
        self._capabilities[model_id] = capabilities
        
        if singleton:
            # Create instance immediately for singleton models
            try:
                self._instances[model_id] = model_class()
                logger.info("model.registered", model_id=model_id, singleton=True)
            except Exception as e:
                logger.error("model.instantiation_failed", model_id=model_id, error=str(e))
                raise
        
    def get_model(self, model_id: Optional[str] = None) -> CXRModel:
        """
        Get a model instance by ID.
        
        Args:
            model_id: Model identifier (defaults to default model)
            
        Returns:
            Model instance
            
        Raises:
            ValueError: If model_id is not registered
        """
        model_id = model_id or self._default_model
        
        if model_id not in self._models:
            available = ", ".join(self._models.keys())
            raise ValueError(
                f"Model '{model_id}' not registered. Available: {available}"
            )
        
        # Return singleton instance if available
        if model_id in self._instances:
            return self._instances[model_id]
        
        # Otherwise create new instance
        return self._models[model_id]()
    
    def get_capabilities(self, model_id: Optional[str] = None) -> ModelCapabilities:
        """Get capabilities for a model variant"""
        model_id = model_id or self._default_model
        if model_id not in self._capabilities:
            raise ValueError(f"Model '{model_id}' not registered")
        return self._capabilities[model_id]
    
    def list_models(self) -> Dict[str, ModelCapabilities]:
        """List all registered models and their capabilities"""
        return self._capabilities.copy()
    
    def set_default(self, model_id: str) -> None:
        """Set the default model variant"""
        if model_id not in self._models:
            raise ValueError(f"Model '{model_id}' not registered")
        self._default_model = model_id
        logger.info("model.default_changed", model_id=model_id)


# Global registry instance
_registry: Optional[ModelRegistry] = None


def get_registry() -> ModelRegistry:
    """Get the global model registry instance"""
    global _registry
    if _registry is None:
        _registry = ModelRegistry()
    return _registry


def register_model(
    model_id: str,
    model_class: Type[CXRModel],
    capabilities: ModelCapabilities,
    singleton: bool = True,
) -> None:
    """Convenience function to register a model"""
    get_registry().register(model_id, model_class, capabilities, singleton)

