from typing import Dict, Optional, Any, Union
from langchain_core.language_models import BaseChatModel
from langchain_openai import ChatOpenAI
from langchain_community.document_compressors import JinaRerank


from alayaflow.utils.singleton import SingletonMeta
from alayaflow.component.model.schemas import ModelProfile, GenerationConfig


class ModelManager(metaclass=SingletonMeta):
    def __init__(self):
        self._profiles: Dict[str, ModelProfile] = {}

    def register_profile(self, profile: ModelProfile, override: bool = False) -> None:
        if profile.model_id in self._profiles and not override:
            raise ValueError(f"Profile with model ID {profile.model_id} already exists.")
        self._profiles[profile.model_id] = profile

    def get_model(
        self, 
        model_id: str, 
        runtime_config: Optional[Union[GenerationConfig, Dict[str, Any]]] = None
    ) -> BaseChatModel:
        if model_id not in self._profiles:
            raise ValueError(f"Model ID '{model_id}' not found. Please register it first.")

        profile = self._profiles[model_id]
        
        # Merge config

        final_params = profile.default_config.model_dump(exclude={"extra_kwargs"})
        extra_kwargs = profile.default_config.extra_kwargs.copy()

        if runtime_config:
            if isinstance(runtime_config, GenerationConfig):
                override_dict = runtime_config.model_dump(exclude_unset=True, exclude={"extra_kwargs"})
                final_params.update(override_dict)
                extra_kwargs.update(runtime_config.extra_kwargs)
            elif isinstance(runtime_config, dict):
                final_params.update(runtime_config)

        # Instantiate model

        if profile.model_type == "OpenAI":
            return ChatOpenAI(
                model=profile.model_name,
                openai_api_key=profile.api_key.get_secret_value(),
                openai_api_base=profile.base_url,
                # Generation parameters
                temperature=final_params.get("temperature"),
                max_tokens=final_params.get("max_tokens"),
                top_p=final_params.get("top_p"),
                frequency_penalty=final_params.get("frequency_penalty"),
                presence_penalty=final_params.get("presence_penalty"),
                model_kwargs=extra_kwargs
            )
        elif profile.model_type == "JinaRerank":
            return JinaRerank(
                model=profile.model_name,
                jina_api_key=profile.api_key.get_secret_value(),
                # Generation parameters
                top_n=final_params.get("top_n")
            )
        else:
            raise ValueError(f"Unsupported model type: {profile.model_type}")

    def get_profile(self, model_id: str) -> Optional[ModelProfile]:
        return self._profiles.get(model_id)