"""Tests for MCPTrace and MCPToolCall data models."""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from ai_qa.mcp.trace import MCPToolCall, MCPTrace


class TestMCPToolCall:
    def test_create_minimal(self) -> None:
        call = MCPToolCall(tool="search_diagnosis")
        assert call.tool == "search_diagnosis"
        assert call.arguments == {}
        assert call.result is None
        assert call.error is None
        assert call.succeeded is True
        assert call.has_result is False

    def test_create_full(self) -> None:
        call = MCPToolCall(
            tool="search_diagnosis",
            arguments={"text": "diabetes"},
            result={"success": True, "count": 5},
            latency_ms=150.0,
            server="di-agentic-service",
        )
        assert call.tool == "search_diagnosis"
        assert call.arguments == {"text": "diabetes"}
        assert call.succeeded is True
        assert call.has_result is True
        assert call.latency_ms == 150.0

    def test_failed_call(self) -> None:
        call = MCPToolCall(
            tool="create_audience",
            arguments={"filter": "{}"},
            error="Permission denied",
        )
        assert call.succeeded is False
        assert call.has_result is False

    def test_empty_tool_name_raises(self) -> None:
        with pytest.raises(ValueError, match="non-empty"):
            MCPToolCall(tool="")

    def test_to_context_dict(self) -> None:
        call = MCPToolCall(
            tool="search_drug",
            arguments={"text": "keytruda"},
            result={"count": 1},
        )
        d = call.to_context_dict()
        assert d["name"] == "search_drug"
        assert d["arguments"] == {"text": "keytruda"}
        assert d["result"] == {"count": 1}
        assert "error" not in d

    def test_to_context_dict_with_error(self) -> None:
        call = MCPToolCall(
            tool="submit_query",
            arguments={"query_text": "SELECT 1"},
            error="Timeout",
        )
        d = call.to_context_dict()
        assert d["error"] == "Timeout"
        assert "result" not in d


class TestMCPTrace:
    def test_empty_trace(self) -> None:
        trace = MCPTrace()
        assert trace.calls == []
        assert trace.tool_names == []
        assert trace.unique_tools == set()
        assert trace.total_latency_ms == 0.0

    def test_trace_with_calls(self) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(tool="search_diagnosis", latency_ms=100),
                MCPToolCall(tool="get_hcp_stats", latency_ms=200),
                MCPToolCall(tool="create_audience", latency_ms=300),
            ],
            user_input="Find oncologists",
            agent_output="Created audience with 500 HCPs",
        )
        assert trace.tool_names == [
            "search_diagnosis",
            "get_hcp_stats",
            "create_audience",
        ]
        assert trace.unique_tools == {
            "search_diagnosis",
            "get_hcp_stats",
            "create_audience",
        }
        assert trace.total_latency_ms == 600.0

    def test_to_tool_calls(self) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(
                    tool="search_diagnosis",
                    arguments={"text": "diabetes"},
                    result={"count": 3},
                ),
            ]
        )
        tool_calls = trace.to_tool_calls()
        assert len(tool_calls) == 1
        assert tool_calls[0]["name"] == "search_diagnosis"
        assert tool_calls[0]["arguments"] == {"text": "diabetes"}

    def test_failed_calls_property(self) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(tool="a"),
                MCPToolCall(tool="b", error="fail"),
                MCPToolCall(tool="c"),
            ]
        )
        assert len(trace.failed_calls) == 1
        assert trace.failed_calls[0].tool == "b"
        assert len(trace.succeeded_calls) == 2

    def test_get_calls_by_tool(self) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(tool="search_drug"),
                MCPToolCall(tool="search_diagnosis"),
                MCPToolCall(tool="search_drug"),
            ]
        )
        drug_calls = trace.get_calls_by_tool("search_drug")
        assert len(drug_calls) == 2

    def test_get_call_by_index(self) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(tool="a"),
                MCPToolCall(tool="b"),
            ]
        )
        assert trace.get_call_by_index(0).tool == "a"
        assert trace.get_call_by_index(1).tool == "b"
        assert trace.get_call_by_index(5) is None

    def test_from_list(self) -> None:
        trace = MCPTrace.from_list([
            {"tool": "search_diagnosis", "arguments": {"text": "asthma"}},
            {"name": "create_audience", "arguments": {"filter": "{}"}},
        ])
        assert len(trace.calls) == 2
        assert trace.calls[0].tool == "search_diagnosis"
        assert trace.calls[1].tool == "create_audience"

    def test_from_jsonl(self, tmp_path: Path) -> None:
        jsonl_file = tmp_path / "trace.jsonl"
        lines = [
            json.dumps({"session_id": "sess-001"}),
            json.dumps({"user_input": "Find oncologists in NY"}),
            json.dumps({
                "tool": "search_hcp_specialities",
                "arguments": {"text": "oncology"},
                "result": {"success": True, "count": 2},
            }),
            json.dumps({
                "tool": "create_audience",
                "arguments": {"audience_name": "Onco NY"},
                "result": {"success": True, "audienceId": 123},
            }),
            json.dumps({"agent_output": "Created audience with ID 123"}),
        ]
        jsonl_file.write_text("\n".join(lines))

        trace = MCPTrace.from_jsonl(str(jsonl_file))
        assert trace.session_id == "sess-001"
        assert trace.user_input == "Find oncologists in NY"
        assert trace.agent_output == "Created audience with ID 123"
        assert len(trace.calls) == 2
        assert trace.calls[0].tool == "search_hcp_specialities"
        assert trace.calls[1].tool == "create_audience"

    def test_from_jsonl_file_not_found(self) -> None:
        with pytest.raises(FileNotFoundError):
            MCPTrace.from_jsonl("/nonexistent/trace.jsonl")

    def test_to_jsonl_roundtrip(self, tmp_path: Path) -> None:
        trace = MCPTrace(
            calls=[
                MCPToolCall(
                    tool="search_diagnosis",
                    arguments={"text": "diabetes"},
                    result={"count": 5},
                    latency_ms=100,
                ),
            ],
            user_input="Find diabetic patients",
            agent_output="Found 5 diagnosis codes",
            session_id="test-session",
        )
        out_path = tmp_path / "output.jsonl"
        trace.to_jsonl(out_path)

        loaded = MCPTrace.from_jsonl(out_path)
        assert loaded.session_id == "test-session"
        assert loaded.user_input == "Find diabetic patients"
        assert loaded.agent_output == "Found 5 diagnosis codes"
        assert len(loaded.calls) == 1
        assert loaded.calls[0].tool == "search_diagnosis"
