"""Tests for EvalRunner and Agent protocol."""

import json
import pytest
from unittest.mock import MagicMock, patch

from ashr_labs.eval import Agent, EvalRunner
from ashr_labs.run_builder import RunBuilder


# ── Test fixtures ────────────────────────────────────────────────────────


class MockAgent:
    """A mock agent that returns scripted responses."""

    def __init__(self, responses: list[dict] | None = None):
        self._responses = responses or []
        self._call_index = 0
        self.reset_count = 0

    def respond(self, message: str) -> dict:
        if self._call_index < len(self._responses):
            resp = self._responses[self._call_index]
            self._call_index += 1
            return resp
        return {"text": "default response", "tool_calls": []}

    def reset(self) -> None:
        self._call_index = 0
        self.reset_count += 1


def _make_dataset_source(scenarios: dict) -> dict:
    """Wrap scenarios into dataset_source format."""
    return {"runs": scenarios}


def _simple_scenario():
    """A single-scenario dataset with one user action and one agent expectation."""
    return _make_dataset_source({
        "test_1": {
            "title": "Order Lookup",
            "actions": [
                {
                    "actor": "user",
                    "content": "What's the status of ORD-123?",
                    "name": "ask_order_status",
                },
                {
                    "actor": "agent",
                    "content": "Your order ORD-123 has shipped and is on the way.",
                    "expected_response": {
                        "tool_calls": [
                            {
                                "name": "lookup_order",
                                "arguments_json": '{"order_id": "ORD-123"}',
                            }
                        ],
                    },
                },
            ],
        }
    })


def _multi_scenario():
    """A dataset with two scenarios."""
    return _make_dataset_source({
        "scenario_a": {
            "title": "Lookup",
            "actions": [
                {"actor": "user", "content": "Check ORD-001", "name": "ask"},
                {
                    "actor": "agent",
                    "content": "Order shipped.",
                    "expected_response": {
                        "tool_calls": [
                            {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-001"}'}
                        ],
                    },
                },
            ],
        },
        "scenario_b": {
            "title": "Inventory",
            "actions": [
                {"actor": "user", "content": "Is WavePad in stock?", "name": "ask"},
                {
                    "actor": "agent",
                    "content": "WavePad is available.",
                    "expected_response": {
                        "tool_calls": [
                            {"name": "check_inventory", "arguments_json": '{"product_name": "WavePad"}'}
                        ],
                    },
                },
            ],
        },
    })


def _get_tool_calls(built, test_index=0):
    """Extract the flattened list of tool call entries from built results."""
    tcs = []
    for ar in built["tests"][test_index]["action_results"]:
        if ar["action_type"] == "tool_call":
            tcs.extend(ar.get("tool_calls", []))
    return tcs


# ── Agent protocol tests ─────────────────────────────────────────────


class TestAgentProtocol:
    def test_mock_agent_satisfies_protocol(self):
        agent = MockAgent()
        assert isinstance(agent, Agent)

    def test_object_without_respond_fails(self):
        class Bad:
            def reset(self) -> None: ...

        assert not isinstance(Bad(), Agent)

    def test_object_without_reset_fails(self):
        class Bad:
            def respond(self, message: str) -> dict:
                return {}

        assert not isinstance(Bad(), Agent)


# ── EvalRunner.run() tests ───────────────────────────────────────────


class TestEvalRunnerRun:
    def test_basic_run_returns_run_builder(self):
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                ],
            }
        ])
        runner = EvalRunner(_simple_scenario())
        result = runner.run(agent)
        assert isinstance(result, RunBuilder)

    def test_basic_run_records_tests(self):
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                ],
            }
        ])
        runner = EvalRunner(_simple_scenario())
        result = runner.run(agent)
        built = result.build()
        assert built["status"] == "completed"
        assert len(built["tests"]) == 1
        assert built["tests"][0]["test_id"] == "test_1"

    def test_tool_match_exact(self):
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                ],
            }
        ])
        runner = EvalRunner(_simple_scenario())
        built = runner.run(agent).build()
        tcs = _get_tool_calls(built)
        assert len(tcs) == 1
        assert tcs[0]["match_status"] == "exact"

    def test_tool_match_with_arguments_dict_format(self):
        """Agents can return tool_calls with 'arguments' (dict) instead of 'arguments_json' (str)."""
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments": {"order_id": "ORD-123"}}
                ],
            }
        ])
        runner = EvalRunner(_simple_scenario())
        built = runner.run(agent).build()
        tcs = _get_tool_calls(built)
        assert len(tcs) == 1
        assert tcs[0]["match_status"] == "exact"
        # Verify arguments were recorded (not empty)
        actual_args = json.loads(tcs[0]["actual"]["arguments_json"])
        assert actual_args.get("order_id") == "ORD-123"

    def test_tool_not_called(self):
        agent = MockAgent([
            {"text": "I can help with that.", "tool_calls": []}
        ])
        runner = EvalRunner(_simple_scenario())
        built = runner.run(agent).build()
        tcs = _get_tool_calls(built)
        assert len(tcs) == 1
        assert tcs[0]["match_status"] == "mismatch"
        assert "NOT_CALLED" in tcs[0]["actual"]["name"]

    def test_text_similarity_recorded(self):
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped and is on the way.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                ],
            }
        ])
        runner = EvalRunner(_simple_scenario())
        built = runner.run(agent).build()
        text_actions = [
            a for a in built["tests"][0]["action_results"]
            if a["action_type"] == "text" and a["actor"] == "agent"
        ]
        assert len(text_actions) == 1
        assert "semantic_similarity" in text_actions[0]
        assert text_actions[0]["semantic_similarity"] > 0.5

    def test_multi_scenario_resets_agent(self):
        agent = MockAgent([
            {"text": "Shipped.", "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}]},
            {"text": "In stock.", "tool_calls": [{"name": "check_inventory", "arguments_json": '{}'}]},
        ])
        runner = EvalRunner(_multi_scenario())
        runner.run(agent)
        # Agent should have been reset once per scenario (2 scenarios)
        assert agent.reset_count == 2

    def test_agent_error_handled(self):
        class ErrorAgent:
            def respond(self, message: str) -> dict:
                raise RuntimeError("API timeout")

            def reset(self) -> None:
                pass

        runner = EvalRunner(_simple_scenario())
        built = runner.run(ErrorAgent()).build()
        # Should still complete without raising
        assert built["status"] == "completed"
        assert len(built["tests"]) == 1

    def test_empty_scenario_skipped(self):
        source = _make_dataset_source({
            "empty": {"title": "Empty", "actions": []},
            "valid": {
                "title": "Valid",
                "actions": [
                    {"actor": "user", "content": "Hi", "name": "greet"},
                ],
            },
        })
        agent = MockAgent([{"text": "Hello!", "tool_calls": []}])
        built = EvalRunner(source).run(agent).build()
        # Only the valid scenario should produce a test
        assert len(built["tests"]) == 1
        assert built["tests"][0]["test_id"] == "valid"

    def test_on_scenario_callback(self):
        agent = MockAgent([
            {"text": "ok", "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}]},
        ])
        scenarios_seen = []
        runner = EvalRunner(_simple_scenario())
        runner.run(agent, on_scenario=lambda sid, s: scenarios_seen.append(sid))
        assert scenarios_seen == ["test_1"]

    def test_on_action_callback(self):
        agent = MockAgent([
            {"text": "ok", "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}]},
        ])
        actions_seen = []
        runner = EvalRunner(_simple_scenario())
        runner.run(agent, on_action=lambda idx, a: actions_seen.append(idx))
        assert actions_seen == [0, 1]  # user action + agent action

    def test_aggregate_metrics(self):
        agent = MockAgent([
            {
                "text": "Your order ORD-123 has shipped.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                ],
            }
        ])
        built = EvalRunner(_simple_scenario()).run(agent).build()
        metrics = built["aggregate_metrics"]
        assert metrics["total_tests"] == 1
        assert metrics["tests_passed"] == 1
        assert metrics["average_similarity_score"] is not None

    def test_tool_pool_persists_across_agent_actions(self):
        """Tool calls should persist and be matchable across multiple agent actions."""
        source = _make_dataset_source({
            "multi_agent": {
                "title": "Multi-agent actions",
                "actions": [
                    {"actor": "user", "content": "Refund ORD-123", "name": "ask"},
                    {
                        "actor": "agent",
                        "content": "Looking up your order.",
                        "expected_response": {
                            "tool_calls": [
                                {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'}
                            ],
                        },
                    },
                    {
                        "actor": "agent",
                        "content": "Refund processed.",
                        "expected_response": {
                            "tool_calls": [
                                {"name": "process_refund", "arguments_json": '{"order_id": "ORD-123", "reason": "damaged"}'}
                            ],
                        },
                    },
                ],
            }
        })
        agent = MockAgent([
            {
                "text": "Let me look that up and process the refund.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments_json": '{"order_id": "ORD-123"}'},
                    {"name": "process_refund", "arguments_json": '{"order_id": "ORD-123", "reason": "damaged"}'},
                ],
            }
        ])
        built = EvalRunner(source).run(agent).build()
        tcs = _get_tool_calls(built)
        # Both tools should be exact matches
        statuses = [tc["match_status"] for tc in tcs]
        assert statuses.count("exact") == 2

    def test_tool_calls_batched_in_action_result(self):
        """Multiple tool calls for the same action should be in one tool_calls array."""
        source = _make_dataset_source({
            "batch": {
                "title": "Batch test",
                "actions": [
                    {"actor": "user", "content": "Help me", "name": "ask"},
                    {
                        "actor": "agent",
                        "content": "Done.",
                        "expected_response": {
                            "tool_calls": [
                                {"name": "lookup_order", "arguments_json": '{"order_id": "A"}'},
                                {"name": "check_inventory", "arguments_json": '{"product_name": "B"}'},
                            ],
                        },
                    },
                ],
            }
        })
        agent = MockAgent([
            {
                "text": "Done.",
                "tool_calls": [
                    {"name": "lookup_order", "arguments": {"order_id": "A"}},
                    {"name": "check_inventory", "arguments": {"product_name": "B"}},
                ],
            }
        ])
        built = EvalRunner(source).run(agent).build()
        tool_actions = [
            a for a in built["tests"][0]["action_results"]
            if a["action_type"] == "tool_call"
        ]
        # Should be ONE action_result with TWO entries in tool_calls
        assert len(tool_actions) == 1
        assert len(tool_actions[0]["tool_calls"]) == 2
        assert tool_actions[0]["tool_calls"][0]["match_status"] == "exact"
        assert tool_actions[0]["tool_calls"][1]["match_status"] == "exact"


class TestEvalRunnerCustomComparators:
    def test_custom_tool_comparator(self):
        def always_exact(exp, act):
            return "exact", None

        agent = MockAgent([
            {"text": "ok", "tool_calls": [{"name": "lookup_order", "arguments_json": '{"order_id": "WRONG"}'}]}
        ])
        runner = EvalRunner(_simple_scenario(), tool_comparator=always_exact)
        built = runner.run(agent).build()
        tcs = _get_tool_calls(built)
        assert tcs[0]["match_status"] == "exact"

    def test_custom_text_comparator(self):
        def always_one(a, b):
            return 1.0

        agent = MockAgent([
            {"text": "completely different", "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}]}
        ])
        runner = EvalRunner(_simple_scenario(), text_comparator=always_one)
        built = runner.run(agent).build()
        text_actions = [
            a for a in built["tests"][0]["action_results"]
            if a["action_type"] == "text" and a["actor"] == "agent"
        ]
        assert text_actions[0]["match_status"] == "exact"
        assert text_actions[0]["semantic_similarity"] == 1.0

    def test_custom_thresholds(self):
        agent = MockAgent([
            {
                "text": "I looked it up for you.",
                "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}],
            }
        ])
        # Set very high thresholds so everything is divergent
        runner = EvalRunner(
            _simple_scenario(),
            similarity_thresholds={"exact": 0.99, "similar": 0.98},
        )
        built = runner.run(agent).build()
        text_actions = [
            a for a in built["tests"][0]["action_results"]
            if a["action_type"] == "text" and a["actor"] == "agent"
        ]
        assert text_actions[0]["match_status"] == "divergent"


class TestEvalRunnerFromDataset:
    def test_from_dataset(self):
        mock_client = MagicMock()
        mock_client.get_dataset.return_value = {
            "id": 322,
            "dataset_source": _simple_scenario(),
        }
        runner = EvalRunner.from_dataset(mock_client, dataset_id=322)
        assert isinstance(runner, EvalRunner)
        mock_client.get_dataset.assert_called_once_with(
            dataset_id=322, include_signed_urls=False
        )


class TestEvalRunnerRunAndDeploy:
    def test_run_and_deploy(self):
        agent = MockAgent([
            {
                "text": "Shipped.",
                "tool_calls": [{"name": "lookup_order", "arguments_json": '{}'}],
            }
        ])
        mock_client = MagicMock()
        mock_client.create_run.return_value = {"id": 99, "status": "ok"}

        runner = EvalRunner(_simple_scenario())
        result = runner.run_and_deploy(agent, mock_client, dataset_id=42)

        assert result["id"] == 99
        mock_client.create_run.assert_called_once()
        call_kwargs = mock_client.create_run.call_args
        assert call_kwargs.kwargs.get("dataset_id") == 42 or call_kwargs[1].get("dataset_id") == 42
