"""
LLM Factory - OpenAI LLM creation with caching.
Manages initialization and caching of OpenAI models.
"""

import contextlib
import logging

from langchain_core.language_models import BaseChatModel
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from config import OPENAI_API_KEY

logger = logging.getLogger(__name__)


class LLMFactory:
    """Singleton factory for managing OpenAI LLM instances with caching."""

    COMMON_MODELS: list[str] = [
        "gpt-4o-mini",
        "gpt-4o",
        "gpt-5-nano",
        "gpt-5-mini",
    ]

    def __init__(self):
        """Initialize LLM factory with empty cache."""
        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:
        """
        Get or create an LLM instance from cache.

        Args:
            model_name: Model identifier (e.g., "gpt-4o-mini", "gemini-2.0-flash-lite-preview-02-05")
            streaming: Enable streaming responses
            json_mode: Enable JSON output format
            api_key: Optional API key override

        Returns:
            Configured LLM instance
        """
        clean_model = model_name.split("/")[-1] if "/" in model_name else model_name
        cache_key = (clean_model, streaming, json_mode, api_key)

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

        logger.info(f"Creating new LLM instance: {clean_model}")
        return self._create_instance(clean_model, streaming, json_mode, api_key)

    def _create_instance(
        self,
        model_name: str,
        streaming: bool = False,
        json_mode: bool = False,
        api_key: str | None = None,
    ) -> BaseChatModel:
        """Create and cache a new OpenAI LLM instance."""
        try:
            llm = self._create_openai(model_name, streaming, json_mode, api_key)
            
            cache_key = (model_name, streaming, json_mode, api_key)
            self._cache[cache_key] = llm
            return llm

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

    def _create_openai(self, model_name: str, streaming: bool, json_mode: bool, api_key: str | None) -> BaseChatModel:
        """Create OpenAI model instance."""
        key = api_key or OPENAI_API_KEY

        if not key:
            raise ValueError("OPENAI_API_KEY is required")

        llm_kwargs = {
            "model": model_name,
            "streaming": streaming,
            "api_key": key,
            "temperature": 0,
            "max_tokens": 1000,
        }
        
        # Nếu bật json_mode, tiêm trực tiếp vào constructor
        if json_mode:
            llm_kwargs["model_kwargs"] = {"response_format": {"type": "json_object"}}
            logger.info(f"⚙️ Initializing OpenAI in JSON mode: {model_name}")

        llm = ChatOpenAI(**llm_kwargs)
        logger.info(f"✅ Created OpenAI: {model_name}")
        return llm

    def _enable_json_mode(self, llm: BaseChatModel, model_name: str) -> BaseChatModel:
        """Enable JSON mode for the LLM."""
        try:
            llm = llm.bind(response_format={"type": "json_object"})
            logger.debug(f"⚙️ JSON mode enabled for {model_name}")
        except Exception as e:
            logger.warning(f"⚠️ JSON mode not supported: {e}")
        return llm

    def initialize(self, skip_warmup: bool = True) -> None:
        """
        Pre-initialize common models.

        Args:
            skip_warmup: Skip initialization if True
        """
        if skip_warmup or self._cache:
            return

        logger.info("🔥 Warming up LLM Factory...")
        for model_name in self.COMMON_MODELS:
            with contextlib.suppress(Exception):
                self.get_model(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,
) -> BaseChatModel:
    """Create or get cached LLM instance."""
    return _factory.get_model(model_name, streaming=streaming, json_mode=json_mode, api_key=api_key)


def init_llm_factory(skip_warmup: bool = True) -> None:
    """Initialize the LLM factory."""
    _factory.initialize(skip_warmup)


def create_embedding_model() -> OpenAIEmbeddings:
    """Create OpenAI embeddings model."""
    return OpenAIEmbeddings(model="text-embedding-3-small", api_key=OPENAI_API_KEY)
