"""
LangChain integration for Amplitude AI SDK.

Provides callback handlers to automatically track LangChain operations in Amplitude.
"""
from __future__ import annotations

import time
from typing import Any
import uuid

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.messages import BaseMessage
    from langchain_core.outputs import GenerationChunk, LLMResult

    LANGCHAIN_AVAILABLE = True
except ImportError:
    try:
        # Fallback for older LangChain versions
        from langchain.callbacks.base import BaseCallbackHandler
        from langchain.schema import GenerationChunk, LLMResult
        from langchain.schema.messages import BaseMessage

        LANGCHAIN_AVAILABLE = True
    except ImportError:
        LANGCHAIN_AVAILABLE = False
        BaseCallbackHandler = object  # Fallback for type hints
        LLMResult = object  # Fallback for type hints
        GenerationChunk = object  # Fallback for type hints
        BaseMessage = object  # Fallback for type hints

from amplitude import Amplitude

from ..context import get_active_context
from ..core.privacy import PrivacyConfig
from ..core.tracking import track_ai_message, track_tool_call, track_user_message
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger


class AmplitudeCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler that tracks LLM usage in Amplitude.

    Automatically captures LangChain operations and emits the standard
    3-event pattern: user_message, ai_message, and tool_call.
    """

    def __init__(
        self,
        amplitude: Amplitude,
        user_id: str,
        privacy_config: PrivacyConfig | None = None,
        conversation_id: str | None = None,
        trace_id: str | None = None,
        agent_id: str | None = None,
        parent_agent_id: str | None = None,
        customer_org_id: str | None = None,
        agent_version: str | None = None,
        context: dict[str, Any] | None = None,
        env: str | None = None,
        event_properties: dict[str, Any] | None = None,
        user_properties: dict[str, Any] | None = None,
        groups: dict[str, Any] | None = None,
    ) -> None:
        """
        Initialize the Amplitude callback handler.

        Args:
            amplitude: Amplitude client or tracker instance
            user_id: User ID for all tracked events
            privacy_config: Privacy configuration for content sanitization
            conversation_id: Conversation ID for linking events
            trace_id: Trace ID for linking related events
            agent_id: Agent identifier
            parent_agent_id: Delegating agent (multi-agent orchestration)
            customer_org_id: End-customer org (multi-tenant platforms)
            agent_version: Agent code version (e.g., "2.1.0")
            context: Arbitrary segmentation context dict
            env: Environment (prod, staging, dev)
            event_properties: Additional properties for all events
            user_properties: User properties to set
            groups: Group analytics properties

        When created inside an active :class:`~amplitude_ai.client.Session`,
        unset fields are auto-inherited from :func:`get_active_context`.
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError("LangChain not installed. Install with: pip install langchain")

        super().__init__()

        # Auto-inherit from SessionContext for any unset fields
        ctx = get_active_context()

        self.amplitude = amplitude
        self.privacy_config = privacy_config or PrivacyConfig()

        self.user_id = user_id or (ctx.user_id if ctx else None) or user_id
        self.conversation_id = conversation_id or (ctx.session_id if ctx else None) or str(uuid.uuid4())
        self.trace_id = trace_id or (ctx.trace_id if ctx else None) or str(uuid.uuid4())
        self.agent_id = agent_id or (ctx.agent_id if ctx else None)
        self.parent_agent_id = parent_agent_id or (ctx.parent_agent_id if ctx else None)
        self.customer_org_id = customer_org_id or (ctx.customer_org_id if ctx else None)
        self.agent_version = agent_version or (ctx.agent_version if ctx else None)
        self.context = context or (ctx.context if ctx else None)
        self.env = env or (ctx.env if ctx else None)
        self.event_properties = event_properties or {}
        self.user_properties = user_properties
        self.groups = groups or (ctx.groups if ctx else None)

        # Track state across callback calls
        self._run_data: dict[str, dict[str, Any]] = {}
        self._turn_counter = 1

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        **kwargs: Any,
    ) -> None:
        """Called when LLM starts running."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))

        # Store run metadata
        self._run_data[str(run_id)] = {
            "start_time": time.time(),
            "model_name": self._extract_model_name(serialized, kwargs),
            "provider": self._extract_provider_name(serialized, kwargs),
            "prompts": prompts,
            "turn_id": self._turn_counter,
        }

        # Track user messages from prompts
        for i, prompt in enumerate(prompts):
            track_user_message(
                amplitude=self.amplitude,
                user_id=self.user_id,
                message_content=prompt,
                conversation_id=self.conversation_id,
                trace_id=self.trace_id,
                turn_id=self._turn_counter + i * 2,  # Even numbers for user messages
                agent_id=self.agent_id,
                parent_agent_id=self.parent_agent_id,
                customer_org_id=self.customer_org_id,
                agent_version=self.agent_version,
                context=self.context,
                env=self.env,
                event_properties={
                    **self.event_properties,
                    "langchain_run_id": str(run_id),
                    "prompt_index": i,
                    "total_prompts": len(prompts),
                },
                user_properties=self.user_properties,
                groups=self.groups,
                privacy_config=self.privacy_config,
            )

    def on_llm_end(
        self,
        response: LLMResult,
        **kwargs: Any,
    ) -> None:
        """Called when LLM ends running."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))
        run_data = self._run_data.get(str(run_id), {})

        if not run_data:
            get_logger(self.amplitude).warning(f"No run data found for LLM end callback: {run_id}")
            return

        # Calculate latency
        start_time = run_data.get("start_time", time.time())
        latency_ms = (time.time() - start_time) * 1000

        # Extract response data
        model_name = run_data.get("model_name", "unknown")
        provider = run_data.get("provider", "unknown")
        turn_id = run_data.get("turn_id", self._turn_counter)

        # Track AI messages for each generation
        for gen_idx, generation in enumerate(response.generations):
            for choice_idx, gen in enumerate(generation):
                response_content = gen.text if hasattr(gen, "text") else str(gen)

                # Extract token usage if available
                token_usage = self._extract_token_usage(response)

                # Calculate cost if we have token usage.
                # Use explicit None checks -- 0 is a valid value (e.g. fully-cached requests).
                total_cost_usd = None
                if token_usage.get("input_tokens") is not None and token_usage.get("output_tokens") is not None:
                    total_cost_usd = calculate_cost(
                        model_name=model_name,
                        input_tokens=token_usage["input_tokens"],
                        output_tokens=token_usage["output_tokens"],
                        reasoning_tokens=token_usage.get("reasoning_tokens", 0),
                        cache_read_input_tokens=token_usage.get("cache_read_input_tokens") or 0,
                        cache_creation_input_tokens=token_usage.get("cache_creation_input_tokens") or 0,
                    )

                track_ai_message(
                    amplitude=self.amplitude,
                    user_id=self.user_id,
                    model_name=model_name,
                    provider=provider,
                    response_content=response_content,
                    latency_ms=latency_ms,
                    conversation_id=self.conversation_id,
                    trace_id=self.trace_id,
                    turn_id=turn_id + gen_idx * 2 + choice_idx * 2 + 1,  # Odd numbers for AI messages
                    total_cost_usd=total_cost_usd,
                    agent_id=self.agent_id,
                    parent_agent_id=self.parent_agent_id,
                    customer_org_id=self.customer_org_id,
                    agent_version=self.agent_version,
                    context=self.context,
                    env=self.env,
                    event_properties={
                        **self.event_properties,
                        "langchain_run_id": str(run_id),
                        "generation_index": gen_idx,
                        "choice_index": choice_idx,
                    },
                    user_properties=self.user_properties,
                    groups=self.groups,
                    privacy_config=self.privacy_config,
                    **token_usage,
                )

        # Update turn counter
        self._turn_counter += len(response.generations) * 2

        # Clean up run data
        self._run_data.pop(str(run_id), None)

    def on_llm_error(
        self,
        error: Exception | KeyboardInterrupt,
        **kwargs: Any,
    ) -> None:
        """Called when LLM errors."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))
        run_data = self._run_data.get(str(run_id), {})

        if not run_data:
            get_logger(self.amplitude).warning(f"No run data found for LLM error callback: {run_id}")
            return

        # Calculate latency
        start_time = run_data.get("start_time", time.time())
        latency_ms = (time.time() - start_time) * 1000

        # Track error as AI message
        track_ai_message(
            amplitude=self.amplitude,
            user_id=self.user_id,
            model_name=run_data.get("model_name", "unknown"),
            provider=run_data.get("provider", "unknown"),
            response_content="",
            latency_ms=latency_ms,
            conversation_id=self.conversation_id,
            trace_id=self.trace_id,
            turn_id=run_data.get("turn_id", self._turn_counter) + 1,
            is_error=True,
            error_message=str(error),
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            event_properties={
                **self.event_properties,
                "langchain_run_id": str(run_id),
                "error_type": type(error).__name__,
            },
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

        # Clean up run data
        if str(run_id) in self._run_data:
            del self._run_data[str(run_id)]

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Called when tool starts running."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))
        tool_name = serialized.get("name", serialized.get("id", "unknown_tool"))

        # Store tool run data
        self._run_data[str(run_id)] = {
            "start_time": time.time(),
            "tool_name": tool_name,
            "input": input_str,
            "type": "tool",
        }

    def on_tool_end(
        self,
        output: str,
        **kwargs: Any,
    ) -> None:
        """Called when tool ends running."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))
        run_data = self._run_data.get(str(run_id), {})

        if not run_data or run_data.get("type") != "tool":
            return

        # Calculate latency
        start_time = run_data.get("start_time", time.time())
        latency_ms = (time.time() - start_time) * 1000

        # Track tool call
        track_tool_call(
            amplitude=self.amplitude,
            user_id=self.user_id,
            tool_name=run_data.get("tool_name", "unknown"),
            success=True,
            latency_ms=latency_ms,
            conversation_id=self.conversation_id,
            trace_id=self.trace_id,
            tool_input=run_data.get("input"),
            tool_output=output,
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            event_properties={**self.event_properties, "langchain_run_id": str(run_id)},
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

        # Clean up run data
        self._run_data.pop(str(run_id), None)

    def on_tool_error(
        self,
        error: Exception | KeyboardInterrupt,
        **kwargs: Any,
    ) -> None:
        """Called when tool errors."""
        run_id = kwargs.get("run_id", str(uuid.uuid4()))
        run_data = self._run_data.get(str(run_id), {})

        if not run_data or run_data.get("type") != "tool":
            return

        # Calculate latency
        start_time = run_data.get("start_time", time.time())
        latency_ms = (time.time() - start_time) * 1000

        # Track failed tool call
        track_tool_call(
            amplitude=self.amplitude,
            user_id=self.user_id,
            tool_name=run_data.get("tool_name", "unknown"),
            success=False,
            latency_ms=latency_ms,
            conversation_id=self.conversation_id,
            trace_id=self.trace_id,
            tool_input=run_data.get("input"),
            error_message=str(error),
            agent_id=self.agent_id,
            parent_agent_id=self.parent_agent_id,
            customer_org_id=self.customer_org_id,
            agent_version=self.agent_version,
            context=self.context,
            env=self.env,
            event_properties={
                **self.event_properties,
                "langchain_run_id": str(run_id),
                "error_type": type(error).__name__,
            },
            user_properties=self.user_properties,
            groups=self.groups,
            privacy_config=self.privacy_config,
        )

        # Clean up run data
        if str(run_id) in self._run_data:
            del self._run_data[str(run_id)]

    def on_chain_start(
        self,
        serialized: dict[str, Any],
        inputs: dict[str, Any],
        **kwargs: Any,
    ) -> None:
        """Called when chain starts running."""
        # For chains, we mainly want to track the overall execution
        # Individual LLM and tool calls will be tracked by their respective callbacks

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        **kwargs: Any,
    ) -> None:
        """Called when chain ends running."""
        # Chain completion tracking could be added here if needed

    def on_chain_error(
        self,
        error: Exception | KeyboardInterrupt,
        **kwargs: Any,
    ) -> None:
        """Called when chain errors."""
        # Chain error tracking could be added here if needed

    def _extract_model_name(self, serialized: dict[str, Any], kwargs: dict[str, Any]) -> str:
        """Extract model name from serialized data or kwargs."""
        # Try model locations in order of preference
        for source in [kwargs, kwargs.get("invocation_params", {}), serialized]:
            for key in ["model", "model_name"]:
                if source.get(key):
                    return source[key]

        # Fallback to class name
        return serialized.get("id", serialized.get("name", "unknown"))

    def _extract_provider_name(self, serialized: dict[str, Any], kwargs: dict[str, Any]) -> str:
        """Extract provider name from serialized data."""
        class_name = serialized.get("id", serialized.get("name", "")).lower()

        # Provider mapping using dict lookup (more efficient)
        provider_map = {
            "openai": "openai",
            "anthropic": "anthropic",
            "claude": "anthropic",
            "gemini": "gemini",
            "google": "gemini",
            "cohere": "cohere",
            "huggingface": "huggingface",
            "hf": "huggingface",
        }

        for keyword, provider in provider_map.items():
            if keyword in class_name:
                return provider

        return "langchain"

    def _extract_token_usage(self, response: LLMResult) -> dict[str, int | None]:
        """Extract token usage from LangChain LLMResult.

        Returns input_tokens as TOTAL (including cached tokens) to match
        calculate_cost() expectations. Cache tokens are extracted as subsets.
        """
        usage_dict: dict[str, int | None] = {
            "input_tokens": None,
            "output_tokens": None,
            "total_tokens": None,
            "reasoning_tokens": None,
            "cache_read_input_tokens": None,
            "cache_creation_input_tokens": None,
        }

        if not (hasattr(response, "llm_output") and response.llm_output):
            return usage_dict

        llm_output = response.llm_output

        # Handle two main formats: OpenAI (token_usage) and Anthropic (usage)
        if "token_usage" in llm_output:
            token_usage = llm_output["token_usage"]
            usage_dict.update(
                {
                    # OpenAI prompt_tokens already includes cached tokens (total)
                    "input_tokens": token_usage.get("prompt_tokens"),
                    "output_tokens": token_usage.get("completion_tokens"),
                    "total_tokens": token_usage.get("total_tokens"),
                    "reasoning_tokens": token_usage.get("completion_tokens_details", {}).get(
                        "reasoning_tokens"
                    ),
                }
            )
            # Extract OpenAI cache details from prompt_tokens_details
            prompt_details = token_usage.get("prompt_tokens_details") or {}
            if isinstance(prompt_details, dict):
                cached = prompt_details.get("cached_tokens")
                if cached is not None:
                    usage_dict["cache_read_input_tokens"] = cached
        elif "usage" in llm_output:
            usage = llm_output["usage"]
            # Anthropic raw input_tokens is non-cached only; normalize to total.
            # Use explicit None checks -- 0 is a valid value (fully-cached requests).
            raw_input = usage.get("input_tokens")
            output_tokens = usage.get("output_tokens")
            cache_read = usage.get("cache_read_input_tokens") or 0
            cache_creation = usage.get("cache_creation_input_tokens") or 0
            total_input = (raw_input or 0) + cache_read + cache_creation if raw_input is not None else None
            usage_dict.update(
                {
                    "input_tokens": total_input,
                    "output_tokens": output_tokens,
                    "total_tokens": (total_input + output_tokens)
                    if total_input is not None and output_tokens is not None
                    else None,
                    "cache_read_input_tokens": cache_read if cache_read is not None else None,
                    "cache_creation_input_tokens": cache_creation if cache_creation is not None else None,
                }
            )

        return usage_dict


def create_amplitude_callback(
    amplitude_api_key: str,
    user_id: str,
    privacy_mode: bool = False,
    conversation_id: str | None = None,
    trace_id: str | None = None,
    agent_id: str | None = None,
    parent_agent_id: str | None = None,
    customer_org_id: str | None = None,
    agent_version: str | None = None,
    context: dict[str, Any] | None = None,
    env: str | None = None,
    event_properties: dict[str, Any] | None = None,
    user_properties: dict[str, Any] | None = None,
    groups: dict[str, Any] | None = None,
    **amplitude_kwargs: Any,
) -> AmplitudeCallbackHandler:
    """
    Create an Amplitude callback handler for LangChain.

    Args:
        amplitude_api_key: Amplitude project API key
        user_id: User ID for all tracked events
        privacy_mode: Enable privacy mode (content hashing)
        conversation_id: Optional conversation ID
        trace_id: Optional trace ID
        agent_id: Agent identifier
        parent_agent_id: Delegating agent (multi-agent orchestration)
        customer_org_id: End-customer org (multi-tenant platforms)
        agent_version: Agent code version (e.g., "2.1.0")
        context: Arbitrary segmentation context dict
        env: Environment (prod, staging, dev)
        event_properties: Additional properties for all events
        user_properties: User properties to set
        groups: Group analytics properties
        **amplitude_kwargs: Additional kwargs for Amplitude client

    Returns:
        Configured AmplitudeCallbackHandler
    """
    from amplitude import Amplitude

    amplitude = Amplitude(amplitude_api_key, **amplitude_kwargs)
    privacy_config = PrivacyConfig(privacy_mode=privacy_mode)

    return AmplitudeCallbackHandler(
        amplitude=amplitude,
        user_id=user_id,
        privacy_config=privacy_config,
        conversation_id=conversation_id,
        trace_id=trace_id,
        agent_id=agent_id,
        parent_agent_id=parent_agent_id,
        customer_org_id=customer_org_id,
        agent_version=agent_version,
        context=context,
        env=env,
        event_properties=event_properties,
        user_properties=user_properties,
        groups=groups,
    )
