"""
LLM Factory - Centralized LLM creation for OpenAI & Gemini.
Quản lý việc khởi tạo và caching các LLM models, tự động nhận diện provider.
"""

import contextlib
import logging
from typing import cast

from langchain_core.language_models import BaseChatModel
from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI

from config import GOOGLE_API_KEY, OPENAI_API_KEY

logger = logging.getLogger(__name__)


class LLMFactory:
    """
    Singleton Class quản lý việc khởi tạo và caching các LLM Models.
    """

    COMMON_MODELS: list[str] = [
        "gpt-4o-mini",
        "gemini-2.0-flash-lite-preview-02-05",
        "gemini-1.5-flash",
    ]

    def __init__(self):
        # Cache dict: Key=(model_name, streaming, json_mode, api_key), Value=LLM Instance
        self._cache: dict[tuple[str, bool, bool, str | None], BaseChatModel] = {}

    def get_model(
        self,
        model_name: str,
        streaming: bool = True,
        json_mode: bool = False,
        api_key: str | None = None,
    ) -> BaseChatModel:
        """
        Lấy LLM instance từ cache hoặc tạo mới.
        """
        # Clean model name
        clean_model = model_name.split("/")[-1] if "/" in model_name else model_name

        cache_key = (clean_model, streaming, json_mode, api_key)

        # 1. Hit Cache
        if cache_key in self._cache:
            logger.debug(f"♻️ Using cached model: {clean_model}")
            return self._cache[cache_key]

        # 2. Miss Cache -> Create New
        return self._create_new_instance(clean_model, streaming, json_mode, api_key)

    def _create_new_instance(
        self,
        model_name: str,
        streaming: bool = False,
        json_mode: bool = False,
        api_key: str | None = None,
    ) -> BaseChatModel:
        """Khởi tạo LLM instance dựa trên tên model"""
        try:
            # 1. Nhận diện Gemini
            if "gemini" in model_name.lower():
                effective_key = api_key or GOOGLE_API_KEY
                if not effective_key:
                    raise ValueError("GOOGLE_API_KEY is missing for Gemini model")
                
                llm = ChatGoogleGenerativeAI(
                    model=model_name,
                    streaming=streaming,
                    google_api_key=effective_key,
                    temperature=0
                )
                logger.info(f"✨ Created Gemini model: {model_name}")

            # 2. Nhận diện OpenAI (hoặc mặc định)
            else:
                effective_key = api_key or OPENAI_API_KEY
                if not effective_key:
                    # Nếu không có OpenAI key, thử dùng Gemini làm fallback cuối cùng
                    if GOOGLE_API_KEY:
                        logger.warning(f"⚠️ No OpenAI key found, falling back to Gemini for {model_name}")
                        llm = ChatGoogleGenerativeAI(
                            model="gemini-1.5-flash",
                            streaming=streaming,
                            google_api_key=GOOGLE_API_KEY,
                            temperature=0
                        )
                    else:
                        raise ValueError("Neither OPENAI_API_KEY nor GOOGLE_API_KEY is available.")
                else:
                    # Khởi tạo OpenAI
                    # Lưu ý: gpt-5-nano nếu không tồn tại sẽ bị lỗi từ phía OpenAI API
                    llm = ChatOpenAI(
                        model=model_name, 
                        streaming=streaming, 
                        api_key=effective_key, 
                        temperature=0
                    )
                    logger.info(f"✅ Created OpenAI model: {model_name}")

            # Apply JSON mode nếu cần
            if json_mode:
                try:
                    llm = llm.bind(response_format={"type": "json_object"})
                    logger.debug(f"⚙️ Enabled JSON Mode for {model_name}")
                except Exception as ex:
                    logger.warning(f"⚠️ Failed to bind JSON mode: {ex}")

            # Lưu vào cache
            cache_key = (model_name, streaming, json_mode, api_key)
            self._cache[cache_key] = cast(BaseChatModel, llm)
            return self._cache[cache_key]

        except Exception as e:
            logger.error(f"❌ Failed to create model {model_name}: {e}")
            raise

    def initialize(self, skip_warmup: bool = True):
        """Pre-initialize common models"""
        if skip_warmup or self._cache:
            return

        logger.info("🔥 Warming up LLM Factory...")
        for name in self.COMMON_MODELS:
            with contextlib.suppress(Exception):
                self.get_model(name, streaming=True)


# --- Singleton Instance & Public API ---
_factory = LLMFactory()

def create_llm(model_name: str, streaming: bool = True, json_mode: bool = False, api_key: str | None = None):
    return _factory.get_model(model_name, streaming=streaming, json_mode=json_mode, api_key=api_key)

def init_llm_factory(skip_warmup: bool = True):
    _factory.initialize(skip_warmup)

def create_embedding_model():
    """Helper để tạo embedding model (OpenAI focus)"""
    from langchain_openai import OpenAIEmbeddings
    from config import OPENAI_API_KEY
    return OpenAIEmbeddings(model="text-embedding-3-small", api_key=OPENAI_API_KEY)
