"""
LangChain integration for AITracer.

This module provides a callback handler that automatically logs
LangChain LLM calls to AITracer.

Usage:
    from aitracer.integrations.langchain import AITracerCallbackHandler
    from langchain_openai import ChatOpenAI

    handler = AITracerCallbackHandler(api_key="your-api-key", project="my-project")

    llm = ChatOpenAI(callbacks=[handler])
    response = llm.invoke("Hello!")
"""

from __future__ import annotations

import time
import uuid
from decimal import Decimal
from typing import Any, Dict, List, Optional, Union

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
    from langchain_core.messages import BaseMessage
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    # Create a dummy base class for type hints
    class BaseCallbackHandler:
        pass


class AITracerCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler that logs LLM calls to AITracer.

    This handler captures:
    - LLM calls (start, end, error)
    - Chain execution
    - Tool usage
    - Token usage and costs

    Example:
        >>> from aitracer.integrations.langchain import AITracerCallbackHandler
        >>> from langchain_openai import ChatOpenAI
        >>>
        >>> handler = AITracerCallbackHandler(
        ...     api_key="at-xxx",
        ...     project="my-langchain-app"
        ... )
        >>> llm = ChatOpenAI(callbacks=[handler])
        >>> response = llm.invoke("What is the capital of France?")
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        project: Optional[str] = None,
        base_url: str = "https://api.aitracer.co",
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        flush_on_chain_end: bool = True,
    ):
        """
        Initialize the AITracer callback handler.

        Args:
            api_key: AITracer API key. If not provided, uses AITRACER_API_KEY env var.
            project: Project name for logging.
            base_url: AITracer API base URL.
            session_id: Optional session ID for grouping related calls.
            user_id: Optional user ID for tracking user interactions.
            metadata: Additional metadata to include with all logs.
            flush_on_chain_end: Whether to flush logs when a chain completes.
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain is not installed. Install it with: pip install langchain-core"
            )

        super().__init__()

        # Import AITracer client
        from aitracer import AITracer

        self.client = AITracer(api_key=api_key, project=project, base_url=base_url)
        self.project = project
        self.session_id = session_id or str(uuid.uuid4())
        self.user_id = user_id
        self.default_metadata = metadata or {}
        self.flush_on_chain_end = flush_on_chain_end

        # Track active runs
        self._runs: Dict[str, Dict[str, Any]] = {}
        self._chain_stack: List[str] = []

    def _get_trace_id(self) -> str:
        """Get the current trace ID (from chain or generate new)."""
        if self._chain_stack:
            return self._chain_stack[0]
        return str(uuid.uuid4())

    def _get_parent_run_id(self) -> Optional[str]:
        """Get the parent run ID if in a chain."""
        if len(self._chain_stack) > 1:
            return self._chain_stack[-1]
        return None

    def _extract_model_info(self, serialized: Dict[str, Any]) -> Dict[str, str]:
        """Extract provider and model from serialized LLM info."""
        # Default values
        provider = "unknown"
        model = "unknown"

        # Try to extract from kwargs
        kwargs = serialized.get("kwargs", {})

        # Check for model name in various places
        if "model" in kwargs:
            model = kwargs["model"]
        elif "model_name" in kwargs:
            model = kwargs["model_name"]

        # Determine provider from class name or model name
        class_name = serialized.get("name", "").lower()
        if "openai" in class_name:
            provider = "openai"
        elif "anthropic" in class_name or "claude" in class_name:
            provider = "anthropic"
        elif "google" in class_name or "gemini" in class_name:
            provider = "google"
        elif "mistral" in class_name:
            provider = "mistral"
        elif "cohere" in class_name:
            provider = "cohere"

        # Also check model name for provider hints
        model_lower = model.lower()
        if "gpt" in model_lower or model_lower.startswith("o1"):
            provider = "openai"
        elif "claude" in model_lower:
            provider = "anthropic"
        elif "gemini" in model_lower:
            provider = "google"

        return {"provider": provider, "model": model}

    def _format_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
        """Format messages for logging."""
        formatted = []
        for msg in messages:
            if isinstance(msg, dict):
                formatted.append(msg)
            elif hasattr(msg, "type") and hasattr(msg, "content"):
                formatted.append({
                    "role": msg.type,
                    "content": msg.content,
                })
            elif hasattr(msg, "role") and hasattr(msg, "content"):
                formatted.append({
                    "role": msg.role,
                    "content": msg.content,
                })
            else:
                formatted.append({"role": "unknown", "content": str(msg)})
        return formatted

    # ========== LLM Callbacks ==========

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when an LLM starts processing."""
        run_id_str = str(run_id)
        model_info = self._extract_model_info(serialized)

        self._runs[run_id_str] = {
            "type": "llm",
            "start_time": time.time(),
            "provider": model_info["provider"],
            "model": model_info["model"],
            "input": {"prompts": prompts},
            "trace_id": self._get_trace_id(),
            "parent_run_id": str(parent_run_id) if parent_run_id else self._get_parent_run_id(),
            "tags": tags or [],
            "metadata": {**self.default_metadata, **(metadata or {})},
        }

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a chat model starts processing."""
        run_id_str = str(run_id)
        model_info = self._extract_model_info(serialized)

        # Format messages for logging
        formatted_messages = []
        for msg_list in messages:
            formatted_messages.extend(self._format_messages(msg_list))

        self._runs[run_id_str] = {
            "type": "chat",
            "start_time": time.time(),
            "provider": model_info["provider"],
            "model": model_info["model"],
            "input": {"messages": formatted_messages},
            "trace_id": self._get_trace_id(),
            "parent_run_id": str(parent_run_id) if parent_run_id else self._get_parent_run_id(),
            "tags": tags or [],
            "metadata": {**self.default_metadata, **(metadata or {})},
        }

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when an LLM finishes processing."""
        run_id_str = str(run_id)
        run_data = self._runs.pop(run_id_str, None)

        if not run_data:
            return

        latency_ms = int((time.time() - run_data["start_time"]) * 1000)

        # Extract token usage
        token_usage = {}
        if response.llm_output:
            token_usage = response.llm_output.get("token_usage", {})

        # Extract output
        output_text = ""
        if response.generations:
            for gen_list in response.generations:
                for gen in gen_list:
                    output_text += gen.text

        # Log to AITracer
        self.client.log(
            trace_id=run_data["trace_id"],
            span_id=run_id_str,
            parent_span_id=run_data.get("parent_run_id"),
            provider=run_data["provider"],
            model=run_data["model"],
            input=run_data["input"],
            output={"content": output_text},
            input_tokens=token_usage.get("prompt_tokens", 0),
            output_tokens=token_usage.get("completion_tokens", 0),
            latency_ms=latency_ms,
            status="success",
            metadata={
                **run_data["metadata"],
                "tags": run_data["tags"],
                "session_id": self.session_id,
                "user_id": self.user_id,
                "source": "langchain",
            },
        )

    def on_llm_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when an LLM errors."""
        run_id_str = str(run_id)
        run_data = self._runs.pop(run_id_str, None)

        if not run_data:
            return

        latency_ms = int((time.time() - run_data["start_time"]) * 1000)

        # Log error to AITracer
        self.client.log(
            trace_id=run_data["trace_id"],
            span_id=run_id_str,
            parent_span_id=run_data.get("parent_run_id"),
            provider=run_data["provider"],
            model=run_data["model"],
            input=run_data["input"],
            output=None,
            latency_ms=latency_ms,
            status="error",
            error_message=str(error),
            metadata={
                **run_data["metadata"],
                "tags": run_data["tags"],
                "session_id": self.session_id,
                "user_id": self.user_id,
                "source": "langchain",
            },
        )

    # ========== Chain Callbacks ==========

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a chain starts."""
        run_id_str = str(run_id)
        self._chain_stack.append(run_id_str)

        self._runs[run_id_str] = {
            "type": "chain",
            "start_time": time.time(),
            "name": serialized.get("name", "unknown"),
            "inputs": inputs,
            "tags": tags or [],
            "metadata": {**self.default_metadata, **(metadata or {})},
        }

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a chain ends."""
        run_id_str = str(run_id)

        if run_id_str in self._chain_stack:
            self._chain_stack.remove(run_id_str)

        self._runs.pop(run_id_str, None)

        # Flush logs if configured
        if self.flush_on_chain_end and not self._chain_stack:
            self.client.flush()

    def on_chain_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a chain errors."""
        run_id_str = str(run_id)

        if run_id_str in self._chain_stack:
            self._chain_stack.remove(run_id_str)

        self._runs.pop(run_id_str, None)

    # ========== Tool Callbacks ==========

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a tool starts."""
        run_id_str = str(run_id)

        self._runs[run_id_str] = {
            "type": "tool",
            "start_time": time.time(),
            "name": serialized.get("name", "unknown"),
            "input": input_str,
            "tags": tags or [],
            "metadata": {**self.default_metadata, **(metadata or {})},
        }

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a tool ends."""
        run_id_str = str(run_id)
        run_data = self._runs.pop(run_id_str, None)

        # Tools are tracked but not logged as separate entries
        # They appear in the chain context

    def on_tool_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> Any:
        """Called when a tool errors."""
        run_id_str = str(run_id)
        self._runs.pop(run_id_str, None)

    # ========== Utility Methods ==========

    def flush(self) -> None:
        """Flush any pending logs to AITracer."""
        self.client.flush()

    def set_session_id(self, session_id: str) -> None:
        """Set the session ID for subsequent calls."""
        self.session_id = session_id

    def set_user_id(self, user_id: str) -> None:
        """Set the user ID for subsequent calls."""
        self.user_id = user_id

    def add_metadata(self, key: str, value: Any) -> None:
        """Add metadata that will be included with all subsequent logs."""
        self.default_metadata[key] = value
