"""Tests for evalkit.contrib.langchain.LangChainTracker."""

from __future__ import annotations

from typing import Any
from uuid import uuid4

import pytest

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


@pytest.fixture()
def _import_guard() -> None:
    """Skip if langchain-core is not installed."""
    pytest.importorskip("langchain_core")


def _make_doc(page_content: str, metadata: dict[str, Any] | None = None) -> Any:
    """Create a mock LangChain Document."""
    from langchain_core.documents import Document

    return Document(page_content=page_content, metadata=metadata or {})


@pytest.mark.usefixtures("_import_guard")
class TestLangChainTracker:
    def _make_tracker(self) -> Any:
        from evalkit.contrib.langchain import LangChainTracker

        return LangChainTracker()

    def test_is_base_tracker(self) -> None:
        tracker = self._make_tracker()
        assert isinstance(tracker, BaseTracker)

    def test_initial_state(self) -> None:
        tracker = self._make_tracker()
        assert tracker.question == ""
        assert tracker.tool_calls == []
        assert tracker.context == []
        assert tracker.answer == ""

    def test_on_chain_start_captures_question_from_input_key(self) -> None:
        tracker = self._make_tracker()
        run_id = uuid4()
        tracker.on_chain_start(
            serialized={},
            inputs={"input": "What is Python?"},
            run_id=run_id,
        )
        assert tracker.question == "What is Python?"

    def test_on_chain_start_captures_question_from_question_key(self) -> None:
        tracker = self._make_tracker()
        run_id = uuid4()
        tracker.on_chain_start(
            serialized={},
            inputs={"question": "What is 2+2?"},
            run_id=run_id,
        )
        assert tracker.question == "What is 2+2?"

    def test_on_chain_start_captures_question_from_query_key(self) -> None:
        tracker = self._make_tracker()
        run_id = uuid4()
        tracker.on_chain_start(
            serialized={},
            inputs={"query": "Tell me about AI"},
            run_id=run_id,
        )
        assert tracker.question == "Tell me about AI"

    def test_on_chain_start_ignores_child_chains(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        child_id = uuid4()
        tracker.on_chain_start(
            serialized={},
            inputs={"input": "Root question"},
            run_id=root_id,
        )
        tracker.on_chain_start(
            serialized={},
            inputs={"input": "Child input"},
            run_id=child_id,
            parent_run_id=root_id,
        )
        assert tracker.question == "Root question"

    def test_on_tool_start_captures_tool_with_inputs_dict(self) -> None:
        tracker = self._make_tracker()
        tracker.on_tool_start(
            serialized={"name": "get_weather"},
            input_str="",
            run_id=uuid4(),
            inputs={"city": "Paris"},
        )
        assert tracker.tool_calls == [
            ToolCall(name="get_weather", parameters={"city": "Paris"}),
        ]

    def test_on_tool_start_parses_json_input_str(self) -> None:
        tracker = self._make_tracker()
        tracker.on_tool_start(
            serialized={"name": "search"},
            input_str='{"query": "python"}',
            run_id=uuid4(),
        )
        assert tracker.tool_calls[0].parameters == {"query": "python"}

    def test_on_tool_start_handles_non_json_input_str(self) -> None:
        tracker = self._make_tracker()
        tracker.on_tool_start(
            serialized={"name": "search"},
            input_str="raw text input",
            run_id=uuid4(),
        )
        assert tracker.tool_calls[0].parameters == {}

    def test_on_tool_end_captures_context(self) -> None:
        tracker = self._make_tracker()
        run_id = uuid4()
        tracker.on_tool_start(
            serialized={"name": "search"},
            input_str="",
            run_id=run_id,
            inputs={"q": "python"},
        )
        tracker.on_tool_end(output="Python is a language.", run_id=run_id)
        assert tracker.context == ["Python is a language."]

    def test_on_tool_end_converts_non_string_output(self) -> None:
        tracker = self._make_tracker()
        run_id = uuid4()
        tracker.on_tool_end(output=42, run_id=run_id)
        assert tracker.context == ["42"]

    def test_on_retriever_end_captures_documents(self) -> None:
        tracker = self._make_tracker()
        docs = [
            _make_doc("Chunk 1 text"),
            _make_doc("Chunk 2 text"),
        ]
        tracker.on_retriever_end(documents=docs, run_id=uuid4())
        assert tracker.context == ["Chunk 1 text", "Chunk 2 text"]

    def test_on_retriever_end_skips_empty_content(self) -> None:
        tracker = self._make_tracker()
        docs = [_make_doc(""), _make_doc("Real content")]
        tracker.on_retriever_end(documents=docs, run_id=uuid4())
        assert tracker.context == ["Real content"]

    def test_on_chain_end_captures_answer_from_output_key(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        # Set up root chain
        tracker.on_chain_start(serialized={}, inputs={"input": "Q"}, run_id=root_id)
        tracker.on_chain_end(outputs={"output": "The answer"}, run_id=root_id)
        assert tracker.answer == "The answer"

    def test_on_chain_end_captures_answer_from_answer_key(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        tracker.on_chain_start(serialized={}, inputs={"input": "Q"}, run_id=root_id)
        tracker.on_chain_end(outputs={"answer": "42"}, run_id=root_id)
        assert tracker.answer == "42"

    def test_on_chain_end_captures_answer_from_text_key(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        tracker.on_chain_start(serialized={}, inputs={"input": "Q"}, run_id=root_id)
        tracker.on_chain_end(outputs={"text": "Result"}, run_id=root_id)
        assert tracker.answer == "Result"

    def test_on_chain_end_ignores_child_chains(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        child_id = uuid4()
        tracker.on_chain_start(serialized={}, inputs={"input": "Q"}, run_id=root_id)
        tracker.on_chain_end(outputs={"output": "child answer"}, run_id=child_id)
        assert tracker.answer == ""

    def test_multiple_tool_calls(self) -> None:
        tracker = self._make_tracker()
        tracker.on_tool_start(
            serialized={"name": "weather"},
            input_str="",
            run_id=uuid4(),
            inputs={"city": "Paris"},
        )
        tracker.on_tool_start(
            serialized={"name": "search"},
            input_str="",
            run_id=uuid4(),
            inputs={"q": "python"},
        )
        assert len(tracker.tool_calls) == 2
        assert tracker.tool_calls[0].name == "weather"
        assert tracker.tool_calls[1].name == "search"

    def test_full_flow(self) -> None:
        tracker = self._make_tracker()
        root_id = uuid4()
        tool_run_id = uuid4()

        # Chain start
        tracker.on_chain_start(
            serialized={},
            inputs={"input": "What is Python?"},
            run_id=root_id,
        )

        # Retriever
        docs = [_make_doc("Python is a programming language.")]
        tracker.on_retriever_end(documents=docs, run_id=uuid4())

        # Tool call
        tracker.on_tool_start(
            serialized={"name": "search_kb"},
            input_str="",
            run_id=tool_run_id,
            inputs={"query": "python"},
        )
        tracker.on_tool_end(
            output="Python was created by Guido van Rossum.",
            run_id=tool_run_id,
        )

        # Chain end
        tracker.on_chain_end(
            outputs={"output": "Python is a high-level language."},
            run_id=root_id,
        )

        assert tracker.question == "What is Python?"
        assert tracker.tool_calls == [
            ToolCall(name="search_kb", parameters={"query": "python"}),
        ]
        assert "Python is a programming language." in tracker.context
        assert "Python was created by Guido van Rossum." in tracker.context
        assert tracker.answer == "Python is a high-level language."
