"""
LLM Factory - Multi-provider LLM creation with caching.
Supports OpenAI, Groq, and Google Gemini models.
"""

import contextlib
import logging

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

from config import GOOGLE_API_KEY, GROQ_API_KEY, OPENAI_API_KEY

logger = logging.getLogger(__name__)


class LLMFactory:
    """Singleton factory for managing LLM instances with caching (OpenAI, Groq, Gemini)."""

    COMMON_MODELS: list[str] = [
        "gemini-3.1-flash-lite-preview",
    ]

    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: {model_name}")
        return self._create_instance(model_name, 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 LLM instance (auto-detect provider)."""
        try:
            # Auto-detect: Gemini models use Google API
            is_gemini = "gemini" in model_name.lower()

            if is_gemini:
                llm = self._create_gemini(model_name, streaming, json_mode, api_key)
            else:
                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_gemini(self, model_name: str, streaming: bool, json_mode: bool, api_key: str | None) -> BaseChatModel:
        """Create Google Gemini model instance. Always uses GOOGLE_API_KEY (ignores api_key param which may be OpenAI key)."""
        key = GOOGLE_API_KEY
        if not key:
            raise ValueError("GOOGLE_API_KEY is required for Gemini models. Set it in .env")

        llm_kwargs = {
            "model": model_name,
            "google_api_key": key,
            "temperature": 0,
            "max_output_tokens": 1500,
            "streaming": streaming,
        }

        if json_mode:
            llm_kwargs["response_mime_type"] = "application/json"
            logger.info(f"⚙️ Initializing Gemini in JSON mode: {model_name}")

        logger.warning(f"🔍 DEBUG: provider=Gemini | model={model_name} | key={key[:10]}...")
        llm = ChatGoogleGenerativeAI(**llm_kwargs)
        logger.info(f"✅ Created Gemini: {model_name} | Streaming: {streaming}")
        return llm

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

        # --- Auto-detect provider ---
        is_groq = any(kw in model_name.lower() for kw in ("gpt-oss", "llama", "mixtral", "gemma", "qwen", "deepseek"))
        # Also detect openai/ prefix used by Groq (e.g. "openai/gpt-oss-120b")
        if model_name.startswith("openai/"):
            is_groq = True

        if is_groq:
            # Always use GROQ_API_KEY for Groq models (ignore api_key param which may be OpenAI key)
            key = GROQ_API_KEY
            base_url = "https://api.groq.com/openai/v1"
            if not key:
                raise ValueError("GROQ_API_KEY is required for Groq models")
        else:
            key = api_key or OPENAI_API_KEY
            base_url = None  # default OpenAI
            if not key:
                raise ValueError("OPENAI_API_KEY is required")

        # Models that require /v1/responses API instead of /v1/chat/completions
        needs_responses_api = "codex" in model_name.lower()

        llm_kwargs = {
            "model": model_name,
            "streaming": streaming,
            "api_key": key,
            "temperature": 0,
            "max_tokens": 1500,
        }

        if base_url:
            llm_kwargs["base_url"] = base_url

        if needs_responses_api:
            llm_kwargs["use_responses_api"] = True
            logger.info(f"🔄 Using Responses API for model: {model_name}")

        if json_mode:
            llm_kwargs["model_kwargs"] = {"response_format": {"type": "json_object"}}
            logger.info(f"⚙️ Initializing OpenAI in JSON mode: {model_name}")

        provider = "Groq" if is_groq else "OpenAI"
        logger.warning(f"🔍 DEBUG: provider={provider} | model={model_name} | base_url={base_url} | key={key[:10]}... | is_groq={is_groq}")
        llm = ChatOpenAI(**llm_kwargs)
        logger.info(f"✅ Created {provider}: {model_name} | Streaming: {streaming}")
        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)
