"""LangChain tracker — captures tool calls, RAG context, and answer via callbacks."""

from __future__ import annotations

import json
from collections.abc import Sequence
from typing import Any
from uuid import UUID

from evalkit.contrib.tracker import BaseTracker
from evalkit.models import ToolCall

try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.documents import Document
except ImportError:
    raise ImportError(
        "LangChain is required for LangChainTracker. " "Install it with: pip install evalkit-llm[langchain]"
    )


class LangChainTracker(BaseTracker, BaseCallbackHandler):
    """Captures tool calls, RAG context, and the final answer from a LangChain run.

    Pass as ``callbacks=[tracker]`` in the config dict::

        from evalkit.contrib.langchain import LangChainTracker

        tracker = LangChainTracker()
        chain.invoke(question, config={"callbacks": [tracker]})
        tracker.tool_calls  # [ToolCall(name="search", ...)]
        tracker.context     # ["chunk1", "chunk2"]
        tracker.answer      # "The answer is..."
    """

    def __init__(self) -> None:
        super().__init__()
        self._question: str = ""
        self._tool_calls: list[ToolCall] = []
        self._context: list[str] = []
        self._answer: str = ""
        self._pending_tools: dict[UUID, str] = {}
        self._root_run_id: UUID | None = None

    @property
    def question(self) -> str:
        return self._question

    @property
    def tool_calls(self) -> list[ToolCall]:
        return list(self._tool_calls)

    @property
    def context(self) -> list[str]:
        return list(self._context)

    @property
    def answer(self) -> str:
        return self._answer

    def on_chain_start(
        self,
        serialized: dict[str, Any],
        inputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        if self._root_run_id is None:
            self._root_run_id = run_id
        if not self._question and parent_run_id is None:
            if isinstance(inputs, dict):
                for key in ("input", "question", "query"):
                    if key in inputs and isinstance(inputs[key], str):
                        self._question = inputs[key]
                        break
            if not self._question and isinstance(inputs, str):
                self._question = inputs

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        inputs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        tool_name: str = serialized.get("name", "")
        args: dict[str, Any] = {}
        if inputs is not None and isinstance(inputs, dict):
            args = dict(inputs)
        elif input_str:
            try:
                parsed = json.loads(input_str)
                if isinstance(parsed, dict):
                    args = parsed
            except (json.JSONDecodeError, TypeError):
                pass

        self._pending_tools[run_id] = tool_name
        self._tool_calls.append(ToolCall(name=tool_name, parameters=args))

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> Any:
        self._pending_tools.pop(run_id, None)
        if output is not None:
            output_str = str(output) if not isinstance(output, str) else output
            if output_str:
                self._context.append(output_str)

    def on_retriever_end(
        self,
        documents: Sequence[Document],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> Any:
        for doc in documents:
            content: str = doc.page_content
            if content:
                self._context.append(content)

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> Any:
        if run_id == self._root_run_id and isinstance(outputs, dict):
            for key in ("output", "answer", "text", "result"):
                if key in outputs:
                    value = outputs[key]
                    if isinstance(value, str):
                        self._answer = value
                        break
