"""
Eval runner for testing agents against Ashr Labs datasets.

Provides the Agent protocol and EvalRunner class that encapsulate the full
eval loop: iterate scenarios, feed user actions to the agent, compare expected
vs actual tool calls and text responses, and produce a RunBuilder result.
"""

from __future__ import annotations

import json
from typing import Any, Callable, Protocol, TYPE_CHECKING, runtime_checkable

from .comparators import compare_tool_args, extract_tool_args, text_similarity
from .run_builder import RunBuilder

if TYPE_CHECKING:
    from .client import AshrLabsClient


@runtime_checkable
class Agent(Protocol):
    """Protocol for agents that can be evaluated.

    Any object with ``respond()`` and ``reset()`` methods works. The
    ``respond()`` method receives a user message and must return a dict
    with ``"text"`` (str) and ``"tool_calls"`` (list of dicts with
    ``"name"`` (str) and ``"arguments"`` (dict) keys).

    Example::

        class MyAgent:
            def respond(self, message: str) -> dict[str, Any]:
                # call your LLM, return {"text": "...", "tool_calls": [...]}
                ...

            def reset(self) -> None:
                # clear conversation history for a new scenario
                ...
    """

    def respond(self, message: str) -> dict[str, Any]: ...
    def reset(self) -> None: ...


# Type aliases for callback signatures
OnScenarioCallback = Callable[[str, dict], None]
OnActionCallback = Callable[[int, dict], None]

# Type alias for custom comparator functions
ToolComparator = Callable[[dict, dict], tuple[str, str | None]]
TextComparator = Callable[[str, str], float]


class EvalRunner:
    """Runs an agent against every scenario in a dataset and records results.

    This is the missing middle layer between fetching a dataset and submitting
    results. It handles:

    1. Iterating all scenarios in the dataset
    2. Feeding user actions to the agent
    3. Matching expected tool calls against actual (by name, popping matched)
    4. Comparing expected text responses using similarity scoring
    5. Producing a populated RunBuilder ready for ``.build()`` or ``.deploy()``

    Args:
        dataset_source: The ``dataset_source`` dict from a dataset, containing
            a ``"runs"`` key mapping scenario IDs to scenario dicts.
        tool_comparator: Optional custom function for comparing tool calls.
            Signature: ``(expected, actual) -> (match_status, divergence_notes)``.
            Defaults to :func:`comparators.compare_tool_args`.
        text_comparator: Optional custom function for comparing text.
            Signature: ``(expected_text, actual_text) -> float``.
            Defaults to :func:`comparators.text_similarity`.
        similarity_thresholds: Optional dict with ``"exact"`` and ``"similar"``
            float thresholds. Defaults to ``{"exact": 0.70, "similar": 0.40}``.

    Example::

        from ashr_labs import AshrLabsClient, EvalRunner

        client = AshrLabsClient(api_key="tp_...")
        runner = EvalRunner.from_dataset(client, dataset_id=322)
        run = runner.run(my_agent)
        run.deploy(client, dataset_id=322)
    """

    def __init__(
        self,
        dataset_source: dict[str, Any],
        *,
        tool_comparator: ToolComparator | None = None,
        text_comparator: TextComparator | None = None,
        similarity_thresholds: dict[str, float] | None = None,
    ) -> None:
        self._source = dataset_source
        self._tool_comparator = tool_comparator or compare_tool_args
        self._text_comparator = text_comparator or text_similarity
        thresholds = similarity_thresholds or {}
        self._exact_threshold = thresholds.get("exact", 0.70)
        self._similar_threshold = thresholds.get("similar", 0.40)

    @classmethod
    def from_dataset(
        cls,
        client: AshrLabsClient,
        dataset_id: int,
        **kwargs: Any,
    ) -> EvalRunner:
        """Create an EvalRunner by fetching a dataset from the API.

        Args:
            client: An authenticated AshrLabsClient.
            dataset_id: The dataset ID to fetch.
            **kwargs: Additional keyword arguments passed to EvalRunner.__init__.

        Returns:
            A configured EvalRunner ready to call ``.run()``.
        """
        ds = client.get_dataset(dataset_id=dataset_id, include_signed_urls=False)
        source = ds.get("dataset_source", {})
        return cls(source, **kwargs)

    def run(
        self,
        agent: Agent,
        *,
        on_scenario: OnScenarioCallback | None = None,
        on_action: OnActionCallback | None = None,
    ) -> RunBuilder:
        """Run the agent against every scenario and return a populated RunBuilder.

        Args:
            agent: An object implementing the Agent protocol.
            on_scenario: Optional callback called at the start of each scenario
                with ``(scenario_id, scenario_dict)``.
            on_action: Optional callback called for each action with
                ``(action_index, action_dict)``.

        Returns:
            A RunBuilder with all test results recorded, ready for
            ``.build()`` or ``.deploy()``.
        """
        runs_data = self._source.get("runs", {})
        run = RunBuilder()
        run.start()

        for run_id, scenario in runs_data.items():
            actions = scenario.get("actions", [])
            if not actions:
                continue

            if on_scenario:
                on_scenario(run_id, scenario)

            agent.reset()
            test = run.add_test(run_id)
            test.start()

            # Track agent response state across actions within this scenario
            agent_text = ""
            agent_tools: list[dict[str, Any]] = []

            for i, action in enumerate(actions):
                if on_action:
                    on_action(i, action)

                actor = action.get("actor")
                content = action.get("content", "")

                if actor == "user":
                    test.add_user_text(
                        text=content,
                        description=action.get("name", f"user_action_{i}"),
                        action_index=i,
                    )
                    try:
                        result = agent.respond(content)
                        agent_text = result.get("text", "")
                        agent_tools = list(result.get("tool_calls", []))
                    except Exception as e:
                        agent_text = f"[error: {e}]"
                        agent_tools = []

                elif actor == "agent":
                    expected = action.get("expected_response", {})
                    expected_tools = expected.get("tool_calls", [])
                    expected_text = content

                    # Compare tool calls — pop matched from pool
                    for exp_tc in expected_tools:
                        exp_name = exp_tc.get("name", "")
                        matched = None
                        matched_idx = None
                        for ti, t in enumerate(agent_tools):
                            if t.get("name") == exp_name:
                                matched = t
                                matched_idx = ti
                                break

                        if matched is not None:
                            status, notes = self._tool_comparator(exp_tc, matched)
                            actual_args = extract_tool_args(matched)
                            test.add_tool_call(
                                expected=exp_tc,
                                actual={
                                    "name": matched.get("name", ""),
                                    "arguments_json": json.dumps(actual_args),
                                },
                                match_status=status,
                                divergence_notes=notes,
                                action_index=i,
                            )
                            agent_tools.pop(matched_idx)
                        else:
                            test.add_tool_call(
                                expected=exp_tc,
                                actual={
                                    "name": "NOT_CALLED",
                                    "arguments_json": "{}",
                                },
                                match_status="mismatch",
                                divergence_notes=f"Agent never called {exp_name}",
                                action_index=i,
                            )

                    # Compare text response
                    if expected_text and agent_text:
                        similarity = self._text_comparator(
                            expected_text, agent_text
                        )
                        if similarity > self._exact_threshold:
                            match_status = "exact"
                        elif similarity > self._similar_threshold:
                            match_status = "similar"
                        else:
                            match_status = "divergent"

                        test.add_agent_response(
                            expected_response={"text": expected_text},
                            actual_response={"text": agent_text},
                            match_status=match_status,
                            semantic_similarity=similarity,
                            action_index=i,
                        )

                    # Clear text after comparing; keep tool pool for later actions
                    agent_text = ""

            test.complete()

        run.complete()
        return run

    def run_and_deploy(
        self,
        agent: Agent,
        client: AshrLabsClient,
        dataset_id: int | None = None,
        *,
        on_scenario: OnScenarioCallback | None = None,
        on_action: OnActionCallback | None = None,
        **deploy_kwargs: Any,
    ) -> dict[str, Any]:
        """Run the eval and submit results in one call.

        Args:
            agent: An object implementing the Agent protocol.
            client: An authenticated AshrLabsClient.
            dataset_id: The dataset to submit against. If None, you must have
                created this runner via ``from_dataset`` (but dataset_id is
                still required for deploy).
            on_scenario: Optional callback per scenario.
            on_action: Optional callback per action.
            **deploy_kwargs: Extra kwargs passed to ``RunBuilder.deploy()``.

        Returns:
            The created run object from the API.
        """
        run = self.run(agent, on_scenario=on_scenario, on_action=on_action)
        return run.deploy(client, dataset_id=dataset_id, **deploy_kwargs)
