"""
LangChain callback handler for automatic usage tracking with Paygent.
This callback intercepts LangChain LLM calls and automatically sends usage data to Paygent.
"""

from typing import Any, Dict, List, Optional
from uuid import UUID

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
except ImportError:
    raise ImportError(
        "LangChain is required for this integration. "
        "Install it with: pip install langchain-core"
    )

from ..client import Client
from ..models import UsageData, UsageDataWithStrings, CustomerCreateOrGetRequest
from ..constants import ServiceProvider
import threading
import logging

logger = logging.getLogger("paygent_sdk.langchain")


class PaygentLangChainCallback(BaseCallbackHandler):
    """
    LangChain callback handler that automatically tracks usage with Paygent.
    
    Usage example:
    ```python
    from paygent_sdk import Client, PaygentLangChainCallback
    from langchain_openai import ChatOpenAI
    
    paygent_client = Client("your-api-key")
    callback = PaygentLangChainCallback(
        paygent_client,
        indicator="chat_completion",
        external_agent_id="agent-123",
        external_customer_id="customer-456"
    )
    
    llm = ChatOpenAI(callbacks=[callback])
    response = llm.invoke("Hello!")
    # Usage automatically tracked!
    ```
    """
    
    def __init__(
        self,
        paygent_client: Client,
        indicator: str,
        external_agent_id: str,
        external_customer_id: str,
        external_customer_name: Optional[str] = None
    ):
        """
        Create a new PaygentLangChainCallback.
        
        Args:
            paygent_client: The Paygent client instance for usage tracking
            indicator: Indicator for the usage event (e.g., "chat_completion")
            external_customer_id: External customer identifier
            external_customer_name: Optional External customer name (defaults to ID if not provided)
        """
        super().__init__()
        self.paygent_client = paygent_client
        self.indicator = indicator
        self.external_agent_id = external_agent_id
        self.external_customer_id = external_customer_id
        self.external_customer_name = external_customer_name
        self.run_info: Dict[UUID, Dict[str, str]] = {}
        
        # Trigger background customer creation
        if self.external_customer_id:
            self._background_create_customer()
            
    def _background_create_customer(self):
        """Fire-and-forget customer creation in the background."""
        def create_task():
            try:
                name = self.external_customer_name or self.external_customer_id
                req = CustomerCreateOrGetRequest(
                    external_id=self.external_customer_id,
                    name=name
                )
                self.paygent_client.create_or_get_customer(req)
            except Exception:
                pass
                pass
                
        thread = threading.Thread(target=create_task, daemon=True)
        thread.start()
    
    # ------------------------------------------------------------------ #
    # Provider lookup table — keyed on LangChain's stable module package  #
    # prefix (the first element of the serialized `id` list).             #
    # Order matters: more-specific entries must come before general ones. #
    # To add a new provider: add ONE line here. Nothing else to change.   #
    # ------------------------------------------------------------------ #
    _MODULE_TO_PROVIDER: List[tuple] = [
        # ── OpenAI family ────────────────────────────────────────────────
        ("langchain_openai",            ServiceProvider.OPENAI),
        ("langchain_azure_openai",      ServiceProvider.AZURE_OPENAI),

        # ── Anthropic ────────────────────────────────────────────────────
        ("langchain_anthropic",         ServiceProvider.ANTHROPIC),

        # ── Google — Vertex AI must come BEFORE generic google/genai ─────
        ("langchain_google_vertexai",   ServiceProvider.VERTEX_AI),
        ("langchain_google_genai",      ServiceProvider.VERTEX_AI),   # Gemini API → same DB provider

        # ── AWS ──────────────────────────────────────────────────────────
        ("langchain_aws",               ServiceProvider.BEDROCK),

        # ── Mistral ──────────────────────────────────────────────────────
        ("langchain_mistralai",         ServiceProvider.MISTRAL_AI),

        # ── Cohere ───────────────────────────────────────────────────────
        ("langchain_cohere",            ServiceProvider.COHERE),

        # ── Groq ─────────────────────────────────────────────────────────
        ("langchain_groq",              ServiceProvider.GROQ),

        # ── Together AI ──────────────────────────────────────────────────
        ("langchain_together",          ServiceProvider.TOGETHER_AI),

        # ── Fireworks AI ─────────────────────────────────────────────────
        ("langchain_fireworks",         ServiceProvider.FIREWORKS_AI),

        # ── Cerebras ─────────────────────────────────────────────────────
        ("langchain_cerebras",          ServiceProvider.CEREBRAS),

        # ── Hugging Face ─────────────────────────────────────────────────
        ("langchain_huggingface",       ServiceProvider.HUGGINGFACE),

        # ── DeepSeek ─────────────────────────────────────────────────────
        ("langchain_deepseek",          ServiceProvider.DEEPSEEK),

        # ── Perplexity ───────────────────────────────────────────────────
        ("langchain_perplexity",        ServiceProvider.PERPLEXITY_AI),

        # ── DeepInfra ────────────────────────────────────────────────────
        ("langchain_deepinfra",         ServiceProvider.DEEPINFRA),

        # ── Nvidia / NIM ─────────────────────────────────────────────────
        ("langchain_nvidia_ai_endpoints", ServiceProvider.NEBIUS),

        # ── OpenRouter ───────────────────────────────────────────────────
        ("langchain_openrouter",        ServiceProvider.OPENROUTER),

        # ── AI21 ─────────────────────────────────────────────────────────
        ("langchain_ai21",              ServiceProvider.AI21),

        # ── Ollama (self-hosted, no pricing — keep last before fallback) ──
        ("langchain_ollama",            "Ollama"),
    ]

    # Name-based heuristic table — used only when module prefix lookup fails.
    # Same ordered approach; avoids substring collision bugs.
    _NAME_HEURISTICS: List[tuple] = [
        ("vertexai",     ServiceProvider.VERTEX_AI),
        ("vertex_ai",    ServiceProvider.VERTEX_AI),
        ("vertex",       ServiceProvider.VERTEX_AI),
        ("bedrock",      ServiceProvider.BEDROCK),
        ("anthropic",    ServiceProvider.ANTHROPIC),
        ("mistral",      ServiceProvider.MISTRAL_AI),
        ("cohere",       ServiceProvider.COHERE),
        ("groq",         ServiceProvider.GROQ),
        ("fireworks",    ServiceProvider.FIREWORKS_AI),
        ("cerebras",     ServiceProvider.CEREBRAS),
        ("together",     ServiceProvider.TOGETHER_AI),
        ("huggingface",  ServiceProvider.HUGGINGFACE),
        ("deepseek",     ServiceProvider.DEEPSEEK),
        ("perplexity",   ServiceProvider.PERPLEXITY_AI),
        ("gemini",       ServiceProvider.VERTEX_AI),     # Gemini → maps to Vertex AI provider in DB
        ("google",       ServiceProvider.GOOGLE),
        ("azure_openai", ServiceProvider.AZURE_OPENAI),  # must precede "openai"
        ("azure",        ServiceProvider.AZURE_AI),
        ("openai",       ServiceProvider.OPENAI),
    ]

    def _extract_provider(self, serialized: Dict[str, Any]) -> str:
        """
        Determine the service provider from LangChain's serialized LLM object.

        Strategy (two-tier):
        1. Module-path lookup: LangChain always serializes the class's fully
           qualified module path in the `id` list. We match the root package
           name against a definitive lookup table — exact, fast, zero false
           positives.
        2. Name heuristics: fallback for edge cases (custom wrappers, future
           LangChain refactors). Uses an ordered table — no raw if/elif chains.
        """
        if not serialized:
            return "unknown"

        # ── Tier 1: module-path lookup (primary, deterministic) ────────────
        id_field = serialized.get("id")
        if id_field:
            # id is always a list: ["langchain_openai", "chat_models", "base", "ChatOpenAI"]
            package = (id_field[0] if isinstance(id_field, list) else str(id_field)).lower()
            for prefix, provider in self._MODULE_TO_PROVIDER:
                if package.startswith(prefix):
                    return provider

        # ── Tier 2: name heuristic (fallback for custom wrappers) ──────────
        name = str(serialized.get("name", "") or "").lower()
        if name:
            for keyword, provider in self._NAME_HEURISTICS:
                if keyword in name:
                    return provider

        return "unknown"


    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a legacy BaseLLM starts running."""
        self._handle_llm_start(serialized, run_id, metadata, kwargs)

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a BaseChatModel starts running (ChatVertexAI, ChatBedrock, etc.)."""
        self._handle_llm_start(serialized, run_id, metadata, kwargs)

    def _handle_llm_start(
        self,
        serialized: Dict[str, Any],
        run_id: UUID,
        metadata: Optional[Dict[str, Any]],
        kwargs: Dict[str, Any],
    ) -> None:
        """Shared logic for both on_llm_start and on_chat_model_start."""
        try:
            provider = self._extract_provider(serialized)
            model_name = metadata.get("ls_model_name", "unknown") if metadata else "unknown"
            
            # For Bedrock, try to get the raw model_id from invocation_params or kwargs
            # since ls_model_name may not be set correctly
            if provider == "AWS Bedrock" and model_name == "unknown":
                invocation_params = kwargs.get("invocation_params", {})
                if isinstance(invocation_params, dict):
                    model_name = (
                        invocation_params.get("model_id")
                        or invocation_params.get("model")
                        or "unknown"
                    )
                # Also check serialized kwargs
                if model_name == "unknown" and "kwargs" in serialized:
                    sk = serialized["kwargs"]
                    model_name = (
                        sk.get("model_id")
                        or sk.get("model")
                        or "unknown"
                    )
            
            # For Vertex AI, try to extract model from serialized kwargs
            if provider == "Google Vertex AI" and model_name == "unknown":
                if "kwargs" in serialized:
                    sk = serialized["kwargs"]
                    model_name = sk.get("model_name") or sk.get("model") or "unknown"

            # Store the run info for use in on_llm_end
            self.run_info[run_id] = {
                "provider": provider,
                "model_name": model_name
            }
        except Exception:
            pass
    
    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM ends running."""
        try:
            # Get the stored run info
            info = self.run_info.get(run_id, {})
            provider = info.get("provider", "unknown")
            model_name = info.get("model_name", "unknown")
            
            # Extract usage information from LangChain's LLMResult
            llm_output = response.llm_output or {}
            
            prompt_tokens = 0
            completion_tokens = 0
            total_tokens = 0
            cached_tokens = 0
            
            # OpenAI-style usage (in llm_output["token_usage"])
            if "token_usage" in llm_output:
                token_usage = llm_output["token_usage"]
                prompt_tokens = token_usage.get("prompt_tokens", 0)
                completion_tokens = token_usage.get("completion_tokens", 0)
                total_tokens = token_usage.get("total_tokens", 0)
                
                # Extract cached tokens if available
                prompt_tokens_details = token_usage.get("prompt_tokens_details", {})
                if isinstance(prompt_tokens_details, dict):
                    cached_tokens = prompt_tokens_details.get("cached_tokens", 0)
            
            # Anthropic-style usage (in llm_output["usage"]) — also used by Bedrock
            elif "usage" in llm_output:
                usage = llm_output["usage"]
                prompt_tokens = usage.get("input_tokens", usage.get("prompt_tokens", 0))
                completion_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
                total_tokens = prompt_tokens + completion_tokens
                
                # Extract Anthropic / Bedrock cached tokens
                cached_tokens = usage.get("cache_read_input_tokens", 0)
            
            # Also check generations[0][0].message for Bedrock and Vertex usage_metadata
            elif (response.generations and 
                  len(response.generations) > 0 and 
                  len(response.generations[0]) > 0):
                gen = response.generations[0][0]
                msg = getattr(gen, "message", None)
                
                # Vertex AI — usage_metadata attribute on the ChatMessage
                if msg is not None and hasattr(msg, "usage_metadata"):
                    usage_metadata = msg.usage_metadata
                    prompt_tokens = getattr(usage_metadata, "input_tokens", 0)
                    completion_tokens = getattr(usage_metadata, "output_tokens", 0)
                    total_tokens = getattr(usage_metadata, "total_tokens", 0)
                
                # Bedrock — usage stored in response_metadata on the ChatMessage
                elif msg is not None and hasattr(msg, "response_metadata"):
                    rmd = msg.response_metadata or {}
                    usage = rmd.get("usage", {})
                    if usage:
                        prompt_tokens = usage.get("input_tokens", usage.get("prompt_tokens", 0))
                        completion_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
                        total_tokens = prompt_tokens + completion_tokens
            
            # Try to extract model name from response if not already set
            if model_name == "unknown":
                if "model_name" in llm_output:
                    model_name = llm_output["model_name"]
                elif (response.generations and 
                      len(response.generations) > 0 and 
                      len(response.generations[0]) > 0):
                    gen = response.generations[0][0]
                    if hasattr(gen, "message") and hasattr(gen.message, "response_metadata"):
                        rmd = gen.message.response_metadata or {}
                        model_name = rmd.get("model_name", rmd.get("model_id", "unknown"))
            
            # Send usage data using a fire-and-forget background thread
            def send_task():
                try:
                    if total_tokens > 0:
                        usage_data = UsageData(
                            service_provider=provider,
                            model=model_name,
                            prompt_tokens=prompt_tokens,
                            completion_tokens=completion_tokens,
                            total_tokens=total_tokens,
                            cached_tokens=cached_tokens
                        )
                        
                        self.paygent_client.send_usage(
                            self.external_agent_id,
                            self.external_customer_id,
                            self.indicator,
                            usage_data
                        )
                    else:
                        # Fallback: use string-based tracking
                        # Extract prompt and output strings
                        prompt_string = ""
                        if response.generations and len(response.generations) > 0:
                            # Get the first generation's text
                            if len(response.generations[0]) > 0:
                                prompt_string = str(response.generations[0][0].text)
                        
                        output_string = ""
                        if response.generations and len(response.generations) > 0:
                            outputs = [gen.text for gen in response.generations[0]]
                            output_string = " ".join(outputs)
                        
                        usage_data_with_strings = UsageDataWithStrings(
                            service_provider=provider,
                            model=model_name,
                            prompt_string=prompt_string,
                            output_string=output_string
                        )
                        
                        self.paygent_client.send_usage_with_token_string(
                            self.external_agent_id,
                            self.external_customer_id,
                            self.indicator,
                            usage_data_with_strings
                        )
                except Exception:
                    pass

            thread = threading.Thread(target=send_task, daemon=True)
            thread.start()
            
            # Clean up the run info
            if run_id in self.run_info:
                del self.run_info[run_id]
                
        except Exception:
            if run_id in self.run_info:
                del self.run_info[run_id]
    
    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM encounters an error."""
        try:
            # Clean up the run info
            if run_id in self.run_info:
                del self.run_info[run_id]
        except Exception:
            pass
