"""
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_core.language_models import BaseChatModel
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from config import CLAUDE_BASE_URL, DS2API_API_KEY, DS2API_BASE_URL, 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: Claude models
            is_claude = "claude" in model_name.lower()

            if is_claude and CLAUDE_BASE_URL:
                llm = self._create_claude(model_name, streaming, json_mode, api_key)
            else:
                # 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. Uses GOOGLE_API_KEY if available, otherwise falls back to Gemini CLI (OAuth)."""
        # Override: use DS2API if configured
        if DS2API_BASE_URL:
            logger.info(f"🔄 Gemini model {model_name} -> using DS2API wrapper (OpenAI-compatible)")
            # Treat as OpenAI-compatible via DS2API
            key = DS2API_API_KEY or api_key or "dummy-key"
            llm_kwargs = {
                "model": model_name,
                "streaming": streaming,
                "api_key": key,
                "temperature": 0,
                "max_tokens": 1500,
                "base_url": DS2API_BASE_URL,
            }
            if json_mode:
                llm_kwargs["model_kwargs"] = {"response_mime_type": "application/json"}
            llm = ChatOpenAI(**llm_kwargs)
            logger.info(f"✅ Created DS2API (from Gemini): {model_name} | Streaming: {streaming}")
            return llm

        key = GOOGLE_API_KEY
        if not key:
            # Fallback: dùng Gemini CLI (OAuth local) thay vì crash
            logger.info("🔄 No GOOGLE_API_KEY — falling back to ChatGeminiCLI (OAuth)")
            from common.gemini_cli_wrapper import ChatGeminiCLI
            llm = ChatGeminiCLI(
                model=model_name,
                temperature=0,
                max_output_tokens=1500,
                streaming=streaming,
            )
            return llm

        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, Groq, or DS2API)."""

        # --- Check DS2API first (override) ---
        if DS2API_BASE_URL:
            # Use DS2API as OpenAI-compatible endpoint
            key = DS2API_API_KEY or api_key or OPENAI_API_KEY or "dummy-key"
            base_url = DS2API_BASE_URL
            logger.info(f"🔄 Using DS2API wrapper: {base_url} | model={model_name}")
            llm_kwargs = {
                "model": model_name,
                "streaming": streaming,
                "api_key": key,
                "temperature": 0,
                "max_tokens": 1500,
                "base_url": base_url,
            }
            if json_mode:
                llm_kwargs["model_kwargs"] = {"response_format": {"type": "json_object"}}
                logger.info(f"⚙️ DS2API in JSON mode: {model_name}")
            if "codex" in model_name.lower():
                llm_kwargs["use_responses_api"] = True
            llm = ChatOpenAI(**llm_kwargs)
            logger.info(f"✅ DS2API created: {model_name} | Streaming: {streaming}")
            return llm

        # --- Auto-detect provider ---
        is_groq = any(kw in model_name.lower() for kw in ("gpt-oss", "llama", "mixtral", "gemma", "qwen"))
        # 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 _create_claude(self, model_name: str, streaming: bool, json_mode: bool, api_key: str | None) -> BaseChatModel:
        """Create Claude model via requests library (bypassing LangChain/anthropic SDK to avoid Cloudflare blocks)."""
        import requests, json as json_lib

        key = api_key or OPENAI_API_KEY
        base_url = CLAUDE_BASE_URL

        logger.info(f"🔄 Using Claude API (requests): {base_url} | model={model_name}")

        # Return a wrapper that mimics LangChain's BaseChatModel interface
        class AnthropicWrapper:
            def __init__(self, key, base_url, model_name, streaming):
                self.key = key
                self.base_url = base_url
                self.model_name = model_name
                self.streaming = streaming
                self._output_schema = None

            def bind_tools(self, tools, **kwargs):
                """Store tools for potential use. Returns self for chaining."""
                self._bound_tools = tools
                return self

            def with_structured_output(self, output_schema, **kwargs):
                self._output_schema = output_schema
                return self

            def _convert_messages(self, messages):
                anthropic_messages = []
                system_prompt = None
                for msg in (messages if isinstance(messages, list) else [messages]):
                    if hasattr(msg, 'content') and hasattr(msg, 'type'):
                        if msg.type == 'system':
                            system_prompt = msg.content
                        else:
                            role = 'user' if msg.type == 'human' else 'assistant'
                            anthropic_messages.append({'role': role, 'content': msg.content})
                    elif isinstance(msg, dict):
                        anthropic_messages.append(msg)
                    elif isinstance(msg, str):
                        anthropic_messages.append({'role': 'user', 'content': msg})
                return anthropic_messages, system_prompt

            def _call_api(self, messages_list, system_prompt=None):
                url = f"{self.base_url}/messages"
                headers = {
                    'x-api-key': self.key,
                    'Content-Type': 'application/json'
                }
                payload = {
                    'model': self.model_name,
                    'max_tokens': 1500,
                    'temperature': 0,
                    'messages': messages_list
                }

                # Handle system prompt - ALWAYS inject JSON format requirement
                json_instruction = "\n\nIMPORTANT: You MUST output ONLY valid JSON. No text before or after the JSON object."
                if system_prompt:
                    payload['system'] = system_prompt + json_instruction
                    logger.info(f"[CLAUDE API] System prompt set (first 200 chars): {payload['system'][:200]}")
                elif self._output_schema:
                    schema_str = str(self._output_schema.model_json_schema() if hasattr(self._output_schema, 'model_json_schema') else self._output_schema)
                    payload['system'] = f"You must output valid JSON matching this schema: {schema_str}. Output ONLY the JSON, no other text."
                    logger.info(f"[CLAUDE API] Output schema prompt set (first 200 chars): {payload['system'][:200]}")

                logger.info(f"[CLAUDE API] Payload messages count: {len(messages_list)}")
                resp = requests.post(url, headers=headers, json=payload, timeout=30, stream=True)
                resp.raise_for_status()

                # Parse SSE response
                full_text = ""
                for line in resp.iter_lines():
                    if line:
                        line_str = line.decode('utf-8') if isinstance(line, bytes) else line
                        if line_str.startswith('data: '):
                            data_str = line_str[6:].strip()
                            if data_str == '[DONE]':
                                break
                            try:
                                data = json_lib.loads(data_str)
                                if data.get('type') == 'content_block_delta':
                                    delta = data.get('delta', {})
                                    if delta.get('type') == 'text_delta':
                                        full_text += delta.get('text', '')
                            except Exception:
                                pass
                return full_text

            def invoke(self, messages, **kwargs):
                result = self._convert_messages(messages)
                messages_list = result[0]
                system_prompt = result[1]
                text = self._call_api(messages_list, system_prompt)
                from langchain_core.messages import AIMessage
                return AIMessage(content=text)

            async def ainvoke(self, messages, **kwargs):
                import asyncio
                loop = asyncio.get_event_loop()
                result = self._convert_messages(messages)
                messages_list = result[0]
                system_prompt = result[1]
                text = await loop.run_in_executor(None, self._call_api, messages_list, system_prompt)
                from langchain_core.messages import AIMessage
                return AIMessage(content=text)

            async def astream(self, messages, **kwargs):
                """Yield the full response as a single AIMessageChunk for streaming compat."""
                from langchain_core.messages import AIMessageChunk
                result = await self.ainvoke(messages, **kwargs)
                yield AIMessageChunk(content=result.content)

        llm = AnthropicWrapper(key, base_url, model_name, streaming)
        logger.info(f"✅ Claude API (requests) created: {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)
