"""Tests for bounded agent control flow."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import pytest
from langchain_core.messages import AIMessage, HumanMessage

from src.agent import (
    MAX_ITERATIONS,
    _build_model,
    build_agent_graph,
    finalize_output,
    llm_call,
    run_agent,
    should_continue,
    stream_agent_steps,
    tool_node,
)
from src.schemas import SystemOutput


@dataclass
class FakeBoundModel:
    """Return a fixed message when the bound model is invoked."""

    message: AIMessage

    def invoke(self, _: list[Any]) -> AIMessage:
        """Return the configured AI message."""
        return self.message


class FakeAgentLLM:
    """Mimic a model object that supports tool binding or structured output."""

    def __init__(
        self,
        *,
        message: AIMessage | None = None,
        output: SystemOutput | None = None,
    ) -> None:
        """Store fake model outputs for the requested surface."""
        self._message = message
        self._output = output

    def bind_tools(self, _: list[Any]) -> FakeBoundModel:
        """Return the fake bound-tool model."""
        assert self._message is not None
        return FakeBoundModel(self._message)

    def with_structured_output(self, _: Any) -> FakeBoundModel:
        """Return the fake structured-output model."""
        assert self._output is not None
        return FakeBoundModel(self._output)  # type: ignore[arg-type]


def test_should_continue_respects_tools_and_iteration_cap() -> None:
    """Choose the next graph step from tool calls and iteration limits."""
    with_tool_call = {
        "messages": [
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "search_kb",
                        "args": {"query": "refund"},
                        "id": "1",
                        "type": "tool_call",
                    }
                ],
            )
        ],
        "used_tools": [],
        "metadata": {"ticket_id": "T-1"},
        "iterations": 0,
    }
    capped = {
        "messages": [AIMessage(content="done")],
        "used_tools": [],
        "metadata": {"ticket_id": "T-1"},
        "iterations": MAX_ITERATIONS,
    }

    assert should_continue(with_tool_call) == "tool_node"
    assert should_continue(capped) == "finalize_output"


def test_tool_node_executes_requested_tool() -> None:
    """Execute requested tools and record normalized tool names."""
    state = {
        "messages": [
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "get_policy",
                        "args": {"policy_name": "refund_policy"},
                        "id": "call-1",
                        "type": "tool_call",
                    }
                ],
            )
        ],
        "used_tools": [],
        "metadata": {"ticket_id": "T-2"},
        "iterations": 0,
    }

    update = tool_node(state)

    assert update["used_tools"] == ["get_policy"]
    assert "refund_policy" in update["messages"][0].content


def test_tool_node_requires_ai_message() -> None:
    """Reject tool execution when the latest message is not from the model."""
    state = {
        "messages": [HumanMessage(content="hello")],
        "used_tools": [],
        "metadata": {"ticket_id": "T-3"},
        "iterations": 0,
    }

    with pytest.raises(ValueError):
        tool_node(state)


def test_llm_call_uses_bound_model(monkeypatch) -> None:
    """Increment iterations and append the model response."""
    fake_message = AIMessage(content="I can answer directly.")
    monkeypatch.setattr("src.agent._build_model", lambda: FakeBoundModel(fake_message))

    update = llm_call(
        {
            "messages": [HumanMessage(content="Hi")],
            "metadata": {"ticket_id": "T-4"},
            "used_tools": [],
            "iterations": 0,
        }
    )

    assert update["iterations"] == 1
    assert update["messages"][0].content == "I can answer directly."


def test_finalize_output_returns_structured_result(monkeypatch) -> None:
    """Convert agent history into the final structured output."""
    fake_output = SystemOutput(
        ticket_id="placeholder",
        intent="order_status",
        route="logistics",
        should_escalate=False,
        response="Your order is shipped.",
    )
    monkeypatch.setattr(
        "src.agent.build_langchain_llm",
        lambda temperature=0: FakeAgentLLM(output=fake_output),
    )

    result = finalize_output(
        {
            "messages": [HumanMessage(content="Where is my order?")],
            "metadata": {"ticket_id": "T-5"},
            "used_tools": ["search_kb"],
            "iterations": 1,
        }
    )

    assert result["final_output"].ticket_id == "T-5"
    assert result["final_output"].tool_call_count == 1


class FakeApp:
    """Mimic the compiled LangGraph application surface."""

    def invoke(self, _: dict[str, Any]) -> dict[str, SystemOutput]:
        """Return a stable final output for tests."""
        return {
            "final_output": SystemOutput(
                ticket_id="T-6",
                intent="refund_request",
                route="general",
                should_escalate=False,
                response="Please share your order ID.",
            )
        }

    def stream(self, _: dict[str, Any], stream_mode: str) -> list[dict[str, Any]]:
        """Return a simple list of streaming updates."""
        assert stream_mode == "updates"
        return [{"llm_call": {"messages": ["step"]}}]


class FakeGraph:
    """Capture graph-building calls before compile."""

    def __init__(self, state_type: Any) -> None:
        """Store the provided state type and graph mutations."""
        self.state_type = state_type
        self.nodes: list[str] = []
        self.edges: list[tuple[str, str]] = []
        self.conditional: list[tuple[str, Any, dict[str, str]]] = []

    def add_node(self, name: str, _: Any) -> None:
        """Record node additions."""
        self.nodes.append(name)

    def add_edge(self, start: str, end: str) -> None:
        """Record edge additions."""
        self.edges.append((start, end))

    def add_conditional_edges(self, source: str, condition: Any, mapping: dict[str, str]) -> None:
        """Record conditional edges."""
        self.conditional.append((source, condition, mapping))

    def compile(self) -> dict[str, Any]:
        """Return a snapshot of the configured graph."""
        return {
            "nodes": self.nodes,
            "edges": self.edges,
            "conditional": self.conditional,
        }


def test_run_agent_and_stream_agent_steps(monkeypatch) -> None:
    """Use the compiled app surface for invoke and stream helpers."""
    monkeypatch.setattr("src.agent.build_agent_graph", lambda: FakeApp())

    result = run_agent("Need a refund", {"ticket_id": "T-6"})
    updates = list(stream_agent_steps("Need a refund", {"ticket_id": "T-6"}))

    assert result.intent == "refund_request"
    assert updates[0]["llm_call"]["messages"] == ["step"]


def test_build_model_binds_all_langchain_tools(monkeypatch) -> None:
    """Bind the configured tool list to the underlying model."""
    captured: dict[str, Any] = {}

    class FakeModel:
        """Capture tool binding calls."""

        def bind_tools(self, tools: list[Any]) -> str:
            """Store the provided tools and return a sentinel."""
            captured["tools"] = tools
            return "bound-model"

    monkeypatch.setattr("src.agent.build_langchain_llm", lambda temperature=0: FakeModel())

    result = _build_model()

    assert result == "bound-model"
    assert len(captured["tools"]) == 5


def test_build_agent_graph_wires_expected_nodes(monkeypatch) -> None:
    """Compile a graph with the expected nodes and edges."""
    monkeypatch.setattr("src.agent.StateGraph", FakeGraph)

    compiled = build_agent_graph()

    assert compiled["nodes"] == ["llm_call", "tool_node", "finalize_output"]
    assert ("tool_node", "llm_call") in compiled["edges"]
