from typing import Dict, Any, Optional, Type, List, Union, Callable
from enum import Enum
import logging
from threading import RLock
from pydantic import BaseModel, Field, ConfigDict

logger = logging.getLogger(__name__)

class ModelType(str, Enum):
    LLM = "llm"
    IMAGE = "image"
    VIDEO = "video"
    AUDIO = "audio"
    LYRICS = "lyrics"
    MUSIC = "music"
    UPSCALE = "upscale"

class ServiceConfig(BaseModel):
    """ 服务 configuration schema with validation"""
    id: str
    module: str
    class_name: str = Field(alias="class")
    name: str
    args: List[Any] = Field(default_factory=list)
    kwargs: Dict[str, Any] = Field(default_factory=dict)
    type: str
    provider: str
    provider_name: Optional[str] = None
    model_key: Optional[str] = None  # 源al model key without provider prefix (e.g., 'qwen-plus')
    enabled: bool = True
    is_default: bool = False
    variants: Optional[Dict[str, str]] = None
    capabilities: Optional[Dict[str, Any]] = None
    resolutions: Optional[Dict[str, Any]] = None
    durations: Optional[Union[Dict[str, Any], List[Any]]] = None
    counts: Optional[Dict[str, int]] = None  # {"min": 1, "max": 4}
    modes: Optional[List[str]] = None  # ["std", "pro"]
    voices: Optional[List[Dict[str, str]]] = None  # [{id: "alloy", name: "Alloy", gender: "female"}, ...]
    fallback_models: Optional[List[str]] = None  # Fallback model IDs for automatic failover
    
    model_config = ConfigDict(populate_by_name=True)

class ServiceFactory:
    """ 工厂 for creating service instances"""
    def __init__(self, config: ServiceConfig, service_class: Type):
        self.config = config
        self.service_class = service_class
        
    def create(self, **overrides) -> Any:
        """ 创建 a new service instance with optional parameter overrides"""
        args = self.config.args
        kwargs = {**self.config.kwargs, **overrides}
        instance = self.service_class(*args, **kwargs)
        # Attach config as dict for metadata access
        setattr(instance, "config", self.config.model_dump())
        return instance

class ModelRegistry:
    """
    Central registry for AI model services.
    Thread-safe registry with factory pattern to avoid shared state issues.
    """
    _factories: Dict[str, ServiceFactory] = {}
    _defaults: Dict[ModelType, str] = {}
    _provider_metadata: Dict[str, Dict[str, Any]] = {}
    _lock = RLock()  # 线程-safe operations

    @classmethod
    def register_provider_metadata(cls, provider_id: str, metadata: Dict[str, Any]):
        """Register metadata for a provider."""
        with cls._lock:
            cls._provider_metadata[provider_id] = metadata
            logger.info(f"Registered metadata for provider: {provider_id}")

    @classmethod
    def get_provider_metadata(cls, provider_id: str) -> Optional[Dict[str, Any]]:
        """ 获取 metadata for a provider."""
        with cls._lock:
            return cls._provider_metadata.get(provider_id)

    @classmethod
    def register_factory(cls, name: str, factory: ServiceFactory, model_type: Optional[ModelType] = None, is_default: bool = False):
        """
        Register a service factory (new approach).
        
        Args:
            name: Unique identifier for the model
            factory: ServiceFactory instance
            model_type: Type of the model
            is_default: Whether this should be the default model for its type
        """
        with cls._lock:
            if name in cls._factories:
                logger.warning(f"Overwriting existing factory registration: {name}")
            
            cls._factories[name] = factory
            logger.info(f"Registered service factory: {name}")

            if model_type and is_default:
                cls._defaults[model_type] = name
                logger.info(f"Set {name} as default for {model_type}")
    
    @classmethod
    def _resolve_name(cls, name: str) -> Optional[str]:
        """
        Resolve a model name to its registered ID.
        
        Supports:
        - Direct ID lookup (e.g., 'dashscope/qwen-plus')
        - Short name lookup via model_key (e.g., 'qwen-plus')
        
        Returns the resolved ID or None if not found.
        """
        # Direct lookup first
        if name in cls._factories:
            return name
        
        # If not found, try to find by model_key
        for registered_id, factory in cls._factories.items():
            config = factory.config
            if hasattr(config, 'model_dump'):
                config = config.model_dump()
            if isinstance(config, dict):
                # Check if model_key matches the short name
                if config.get('model_key') == name:
                    return registered_id
        
        return None

    @classmethod
    def set_default_by_id(cls, model_type: ModelType, model_id: str):
        """ 集合 a default model for a specific type using its ID.
        Supports both composite IDs and short names (via model_key lookup).
        """
        with cls._lock:
            resolved_id = cls._resolve_name(model_id)
            if not resolved_id:
                logger.warning(f"Cannot set default: Model {model_id} not found.")
                return
                
            cls._defaults[model_type] = resolved_id
            logger.info(f"Set {resolved_id} as default for {model_type}")

    @classmethod
    def get_default_id(cls, model_type: ModelType) -> Optional[str]:
        """ 获取 the default model ID for a specific type."""
        with cls._lock:
            return cls._defaults.get(model_type)

    @classmethod
    def get_config(cls, name: str) -> Optional[Dict[str, Any]]:
        """Get the configuration dictionary for a specific model service."""
        with cls._lock:
            resolved_name = cls._resolve_name(name)
            if not resolved_name:
                return None
            
            factory = cls._factories.get(resolved_name)
            if factory:
                config = factory.config
                if hasattr(config, 'model_dump'):
                    return config.model_dump()
                return config
            return None

    @classmethod
    def get_all_configs(cls) -> Dict[str, Dict[str, Any]]:
        """Get all registered service configurations."""
        configs = {}
        with cls._lock:
            for name, factory in cls._factories.items():
                config = factory.config
                if hasattr(config, 'model_dump'):
                    configs[name] = config.model_dump()
                else:
                    configs[name] = config
        return configs

    @classmethod
    def get(cls, name: str, **overrides) -> Optional[Any]:
        """ 获取 a specific model service by name.
        Creates a new instance each time to avoid shared state issues.
        
        Supports both composite IDs (e.g., 'dashscope/qwen-plus') and 
        short names (e.g., 'qwen-plus') via model_key lookup.
        
        Args:
            name: Service identifier (can be model ID, short name, or variant name)
            **overrides: Optional parameter overrides for service creation
        """
        with cls._lock:
            # 1. Direct lookup or resolve by model_key
            resolved_name = cls._resolve_name(name)
            if resolved_name:
                factory = cls._factories.get(resolved_name)
                if factory:
                    return factory.create(**overrides)
            
            # 2. Variant resolution - search all models for matching variant
            for model_id, factory in cls._factories.items():
                if hasattr(factory, 'config'):
                    config = factory.config
                    # Handle both dict and ServiceConfig object
                    variants = None
                    if hasattr(config, 'variants'):
                        variants = config.variants
                    elif isinstance(config, dict):
                        variants = config.get('variants')
                    
                    if variants and isinstance(variants, dict):
                        # Check if name matches any variant
                        for variant_type, variant_name in variants.items():
                            if variant_name == name:
                                logger.info(f"Resolved variant '{name}' to model '{model_id}' (variant type: {variant_type})")
                                return factory.create(**overrides)
            
            return None

    @classmethod
    def get_default(cls, model_type: ModelType, **overrides) -> Optional[Any]:
        """ 获取 the default model service for a specific type."""
        name = cls.get_default_id(model_type)
        if name:
            return cls.get(name, **overrides)
        return None

    @classmethod
    def list_providers(cls) -> List[Dict[str, Any]]:
        """ 列表 all unique providers and their capabilities from registered models."""
        with cls._lock:
            providers_map = {}
            
            for factory in cls._factories.values():
                if not hasattr(factory, 'config'):
                    continue
                
                config = factory.config
                
                # Handle both dict and ServiceConfig object
                if isinstance(config, dict):
                    p_id = config.get('provider')
                    p_name = config.get('provider_name')
                    p_type = config.get('type')
                else:
                    p_id = getattr(config, 'provider', None)
                    p_name = getattr(config, 'provider_name', None)
                    p_type = getattr(config, 'type', None)
                
                if not p_id:
                    continue
                
                if p_id not in providers_map:
                    providers_map[p_id] = {
                        "id": p_id,
                        "name": p_name or p_id,
                        "types": set()
                    }
                
                # Prioritize human-readable name if available
                if p_name and providers_map[p_id]["name"] == p_id:
                    providers_map[p_id]["name"] = p_name
                    
                if p_type:
                    providers_map[p_id]["types"].add(p_type)
            
            # 合并 provider metadata and add providers that only have metadata
            for p_id, metadata in cls._provider_metadata.items():
                if p_id not in providers_map:
                    providers_map[p_id] = {
                        "id": p_id,
                        "name": metadata.get("name", p_id),
                        "types": set()
                    }
                
                # 合并 metadata, preserving id and types
                for k, v in metadata.items():
                    if k not in ["id", "types"]:
                        providers_map[p_id][k] = v

            return [
                {**p, "types": list(p["types"])}
                for p in providers_map.values()
            ]

    @classmethod
    def list_models(cls) -> Dict[str, Dict[str, Any]]:
        """ 列表 all registered models with their configurations."""
        with cls._lock:
            result = {}
            for name, factory in cls._factories.items():
                if hasattr(factory, 'config'):
                    config = factory.config.model_dump() if hasattr(factory.config, 'model_dump') else factory.config
                    
                    # Check if this model is the default for its type
                    model_type_str = config.get('type')
                    if model_type_str:
                        try:
                            model_type = ModelType(model_type_str.lower())
                            config['is_default'] = (cls._defaults.get(model_type) == name)
                        except ValueError:
                            config['is_default'] = False
                    else:
                        config['is_default'] = False
                    
                    result[name] = config
            return result
    
    @classmethod
    def find_services(
        cls,
        provider: Optional[str] = None,
        model_type: Optional[ModelType] = None,
        capabilities: Optional[Dict[str, bool]] = None,
        enabled_only: bool = True
    ) -> List[Dict[str, Any]]:
        """ 查找 services matching specified criteria.
        
        Args:
            provider: Filter by provider name
            model_type: Filter by model type
            capabilities: Filter by capabilities (e.g., {"supportsRefImage": True})
            enabled_only: Only return enabled services
            
        Returns:
            List of matching service configurations
        """
        with cls._lock:
            results = []
            for name, factory in cls._factories.items():
                if not hasattr(factory, 'config'):
                    continue
                
                # 转换 config to dict
                config = factory.config.model_dump() if hasattr(factory.config, 'model_dump') else factory.config
                
                # Apply filters
                if enabled_only and not config.get("enabled", True):
                    continue
                    
                if provider and config.get("provider") != provider:
                    continue
                    
                if model_type and config.get("type") != model_type.value:
                    continue
                    
                if capabilities:
                    service_caps = config.get("capabilities", {})
                    if not all(service_caps.get(k) == v for k, v in capabilities.items()):
                        continue
                
                # Check if this model is the default for its type
                model_type_str = config.get('type')
                if model_type_str:
                    try:
                        mt = ModelType(model_type_str.lower())
                        config['is_default'] = (cls._defaults.get(mt) == name)
                    except ValueError:
                        config['is_default'] = False
                else:
                    config['is_default'] = False
                
                results.append({
                    "id": name,
                    **config
                })
            
            return results
    
    @classmethod
    def get_config(cls, name: str) -> Optional[Dict[str, Any]]:
        """ 获取 configuration for a specific service.
        
        Supports both composite IDs (e.g., 'dashscope/qwen-plus') and 
        short names (e.g., 'qwen-plus') via model_key lookup.
        """
        with cls._lock:
            # Resolve short name to composite ID if needed
            resolved_name = cls._resolve_name(name)
            if not resolved_name:
                return None
                
            factory = cls._factories.get(resolved_name)
            if factory and hasattr(factory, 'config'):
                config = factory.config.model_dump() if hasattr(factory.config, 'model_dump') else factory.config
                
                # Check if this model is the default for its type
                model_type_str = config.get('type')
                if model_type_str:
                    try:
                        model_type = ModelType(model_type_str.lower())
                        config['is_default'] = (cls._defaults.get(model_type) == resolved_name)
                    except ValueError:
                        config['is_default'] = False
                else:
                    config['is_default'] = False
                    
                return config
            return None
