"""LangChain integration for AgentCost.

Usage:
    from agentcost import AgentCost
    from agentcost.integrations.langchain import CostTracker

    ac = AgentCost(api_key="ac_xxxxx")
    tracker = CostTracker(ac, agent_name="my-agent")

    # Pass as a callback to any LangChain invocation
    agent.invoke(input, config={"callbacks": [tracker]})
"""

from __future__ import annotations

import time
import uuid
from typing import Any, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from agentcost import AgentCost

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult

    HAS_LANGCHAIN = True
except ImportError:
    HAS_LANGCHAIN = False


def _require_langchain():
    if not HAS_LANGCHAIN:
        raise ImportError(
            "langchain-core is required for the LangChain integration. "
            "Install it with: pip install agentcost[langchain]"
        )


if HAS_LANGCHAIN:

    class CostTracker(BaseCallbackHandler):
        """LangChain callback handler that tracks costs via AgentCost."""

        def __init__(
            self,
            ac: "AgentCost",
            agent_name: str = "langchain-agent",
            task: Optional[str] = None,
            **metadata: Any,
        ):
            super().__init__()
            self._ac = ac
            self._agent_name = agent_name
            self._run_id = str(uuid.uuid4())
            self._metadata = {"task": task, **metadata} if task else metadata
            self._tool_starts: dict[str, float] = {}
            self._llm_starts: dict[str, float] = {}

        # --- LLM callbacks ---

        def on_llm_start(
            self,
            serialized: dict[str, Any],
            prompts: list[str],
            *,
            run_id: Any = None,
            **kwargs: Any,
        ) -> None:
            key = str(run_id) if run_id else str(uuid.uuid4())
            self._llm_starts[key] = time.time()

        def on_llm_end(
            self,
            response: "LLMResult",
            *,
            run_id: Any = None,
            **kwargs: Any,
        ) -> None:
            key = str(run_id) if run_id else ""
            start_time = self._llm_starts.pop(key, None)
            duration_ms = (
                int((time.time() - start_time) * 1000) if start_time else None
            )

            # Extract token usage from response
            tokens_in = 0
            tokens_out = 0
            model = "unknown"

            if response.llm_output:
                usage = response.llm_output.get("token_usage", {})
                tokens_in = usage.get("prompt_tokens", 0)
                tokens_out = usage.get("completion_tokens", 0)
                model = response.llm_output.get(
                    "model_name", response.llm_output.get("model", "unknown")
                )

            # Calculate cost
            cost = self._ac.pricing.llm_cost(model, tokens_in, tokens_out)

            from ..event import CostEvent

            event = CostEvent(
                agent_id=self._agent_name,
                run_id=self._run_id,
                component_name=model,
                event_type="llm_call",
                cost_usd=cost,
                duration_ms=duration_ms,
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                metadata=self._metadata,
            )
            self._ac._client.enqueue(event)

        # --- Tool callbacks ---

        def on_tool_start(
            self,
            serialized: dict[str, Any],
            input_str: str,
            *,
            run_id: Any = None,
            **kwargs: Any,
        ) -> None:
            key = str(run_id) if run_id else str(uuid.uuid4())
            self._tool_starts[key] = time.time()

        def on_tool_end(
            self,
            output: str,
            *,
            run_id: Any = None,
            **kwargs: Any,
        ) -> None:
            key = str(run_id) if run_id else ""
            start_time = self._tool_starts.pop(key, None)
            duration_ms = (
                int((time.time() - start_time) * 1000) if start_time else None
            )

            # Get tool name from kwargs or serialized data
            tool_name = kwargs.get("name", "unknown_tool")

            cost = self._ac.pricing.tool_cost(tool_name)

            from ..event import CostEvent

            event = CostEvent(
                agent_id=self._agent_name,
                run_id=self._run_id,
                component_name=tool_name,
                event_type="tool_call",
                cost_usd=cost,
                duration_ms=duration_ms,
                metadata=self._metadata,
            )
            self._ac._client.enqueue(event)

else:
    # Stub class when langchain is not installed
    class CostTracker:  # type: ignore[no-redef]
        def __init__(self, *args: Any, **kwargs: Any):
            _require_langchain()
