"""
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 copy
import inspect
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, Protocol, TYPE_CHECKING, runtime_checkable

from .comparators import extract_tool_args, compare_args_structural
from .run_builder import RunBuilder, TestBuilder

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]
OnEnvironmentCallback = Callable[[str, dict], dict[str, Any] | None]


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

    The SDK does **not** grade results locally. It handles:

    1. Iterating all scenarios in the dataset
    2. Feeding user actions to the agent
    3. Pairing expected tool calls against actual (by name)
    4. Pairing expected text responses against actual
    5. Producing a populated RunBuilder ready for ``.deploy()``

    All scoring and grading is performed server-side after submission.
    The backend uses LLM-based semantic matching for tool arguments and
    response comparison, which is more accurate than local heuristics.

    Args:
        dataset_source: The ``dataset_source`` dict from a dataset, containing
            a ``"runs"`` key mapping scenario IDs to scenario dicts.

    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],
    ) -> None:
        self._source = dataset_source

    @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)

    @staticmethod
    def _accepts_scenario_id(method: Any) -> bool:
        """Check if a method accepts a scenario_id parameter."""
        try:
            sig = inspect.signature(method)
            params = list(sig.parameters.values())
            # Check for scenario_id param or **kwargs
            for p in params:
                if p.name == "scenario_id":
                    return True
                if p.kind == inspect.Parameter.VAR_KEYWORD:
                    return True
            return False
        except (ValueError, TypeError):
            return False

    def _run_scenario(
        self,
        agent: Agent,
        run_id: str,
        scenario: dict[str, Any],
        *,
        on_scenario: OnScenarioCallback | None = None,
        on_action: OnActionCallback | None = None,
        on_environment: OnEnvironmentCallback | None = None,
    ) -> TestBuilder:
        """Run a single scenario against an agent and return a populated TestBuilder.

        The SDK does NOT grade results locally. It pairs expected vs actual
        tool calls and text responses, then submits them for backend grading.
        """
        if on_scenario:
            on_scenario(run_id, scenario)

        # Pass scenario_id if the agent supports it (backward-compatible)
        if self._accepts_scenario_id(agent.reset):
            agent.reset(scenario_id=run_id)
        else:
            agent.reset()
        test = TestBuilder(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(scenario.get("actions", [])):
            if on_action:
                on_action(i, action)

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

            if actor not in ("agent", "environment"):
                # Treat any non-agent, non-environment actor as user input.
                # Datasets may use custom actor names like "planner",
                # "traveler", "companion", etc. — all feed into the agent.
                test.add_user_text(
                    text=content,
                    description=action.get("name", f"user_action_{i}"),
                    action_index=i,
                )
                try:
                    if self._accepts_scenario_id(agent.respond):
                        result = agent.respond(content, scenario_id=run_id)
                    else:
                        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

                # Pair expected tool calls with actual by name — no scoring
                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:
                        actual_args = extract_tool_args(matched)
                        actual_tc = {
                            "name": matched.get("name", ""),
                            "arguments_json": json.dumps(actual_args),
                        }
                        status, arg_comp = compare_args_structural(exp_tc, actual_tc)
                        test.add_tool_call(
                            expected=exp_tc,
                            actual=actual_tc,
                            match_status=status,
                            argument_comparison=arg_comp,
                            action_index=i,
                        )
                        agent_tools.pop(matched_idx)
                    else:
                        not_called = {
                            "name": "NOT_CALLED",
                            "arguments_json": "{}",
                        }
                        status, arg_comp = compare_args_structural(exp_tc, not_called)
                        test.add_tool_call(
                            expected=exp_tc,
                            actual=not_called,
                            match_status=status,
                            argument_comparison=arg_comp,
                            action_index=i,
                        )

                # Check if the NEXT action is also an agent action — if so,
                # carry leftover tool calls forward instead of marking them as extra.
                # This handles cases like retries where a single respond() produces
                # tool calls that span multiple expected agent actions.
                next_action = None
                if i + 1 < len(scenario.get("actions", [])):
                    next_action = scenario["actions"][i + 1]

                if agent_tools and next_action and next_action.get("actor") == "agent":
                    # Carry remaining tool calls to the next agent action
                    pass
                else:
                    # No more consecutive agent actions — record extras as unexpected
                    for extra_tc in agent_tools:
                        extra_args = extract_tool_args(extra_tc)
                        none_expected = {
                            "name": "NONE_EXPECTED",
                            "arguments_json": "{}",
                        }
                        actual_extra = {
                            "name": extra_tc.get("name", ""),
                            "arguments_json": json.dumps(extra_args),
                        }
                        status, arg_comp = compare_args_structural(none_expected, actual_extra)
                        test.add_tool_call(
                            expected=none_expected,
                            actual=actual_extra,
                            match_status=status,
                            argument_comparison=arg_comp,
                            action_index=i,
                        )
                    agent_tools = []

                # Pair expected vs actual text — no similarity scoring
                if expected_text or agent_text:
                    test.add_agent_response(
                        expected_response={"text": expected_text},
                        actual_response={"text": agent_text},
                        match_status="pending",
                        action_index=i,
                    )

                # Only clear text if the next action is NOT another agent action.
                # When consecutive agent actions exist (e.g., tools+text then text-only),
                # carry the agent's text forward so the next action can pair with it.
                next_is_agent = (
                    next_action is not None
                    and next_action.get("actor") == "agent"
                )
                if not next_is_agent:
                    agent_text = ""

            elif actor == "environment":
                # Environment actions inject external context (e.g., tool results,
                # system events). By default they are skipped. If the user provides
                # an on_environment callback, it receives the content and the full
                # action dict. Returning a dict with "text"/"tool_calls" keys will
                # update the agent state as if the agent had responded.
                if on_environment:
                    env_result = on_environment(content, action)
                    if env_result and isinstance(env_result, dict):
                        agent_text = env_result.get("text", agent_text)
                        new_tools = env_result.get("tool_calls")
                        if new_tools is not None:
                            agent_tools = list(new_tools)

        test.complete()
        return test

    def run(
        self,
        agent: Agent,
        *,
        on_scenario: OnScenarioCallback | None = None,
        on_action: OnActionCallback | None = None,
        on_environment: OnEnvironmentCallback | None = None,
        max_workers: int = 1,
    ) -> 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)``.
            on_environment: Optional callback for environment actions. Receives
                ``(content, action_dict)``. By default environment actions are
                skipped. Return a dict with ``"text"`` and/or ``"tool_calls"``
                to feed the environment context to the agent.
            max_workers: Number of scenarios to run in parallel. Defaults to 1
                (sequential). When >1, each scenario gets a ``copy.deepcopy``
                of the agent so they can run independently. Callbacks may be
                invoked from multiple threads.

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

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

        if max_workers <= 1:
            # Sequential — use the agent directly (no copy)
            for run_id, scenario in scenarios:
                test = self._run_scenario(
                    agent, run_id, scenario,
                    on_scenario=on_scenario, on_action=on_action,
                    on_environment=on_environment,
                )
                run._tests.append(test)
        else:
            # Parallel — deep-copy the agent per scenario
            def _worker(run_id: str, scenario: dict) -> TestBuilder:
                try:
                    agent_copy = copy.deepcopy(agent)
                except Exception as e:
                    raise RuntimeError(
                        f"Failed to deep-copy agent for parallel execution "
                        f"(max_workers={max_workers}). LLM clients with "
                        f"connection pools (Anthropic, OpenAI) cannot be "
                        f"deep-copied. Use max_workers=1 (sequential), or "
                        f"implement __deepcopy__ on your agent to create "
                        f"fresh clients. Original error: {e}"
                    ) from e
                return self._run_scenario(
                    agent_copy, run_id, scenario,
                    on_scenario=on_scenario, on_action=on_action,
                    on_environment=on_environment,
                )

            # Preserve original scenario order in results
            results: list[TestBuilder | None] = [None] * len(scenarios)
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_idx = {
                    executor.submit(_worker, rid, sc): idx
                    for idx, (rid, sc) in enumerate(scenarios)
                }
                for future in as_completed(future_to_idx):
                    idx = future_to_idx[future]
                    try:
                        results[idx] = future.result()
                    except Exception:
                        # Scenario raised — record as a failed test
                        rid, _ = scenarios[idx]
                        failed = TestBuilder(rid)
                        failed.start()
                        failed.complete(status="failed")
                        results[idx] = failed

            for test in results:
                if test is not None:
                    run._tests.append(test)

        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,
        on_environment: OnEnvironmentCallback | None = None,
        max_workers: int = 1,
        **deploy_kwargs: Any,
    ) -> dict[str, Any]:
        """Run the eval and submit results in one call.

        After submission, the backend grader asynchronously scores results
        (typically 1-3 minutes). The returned run will have
        ``grading_status: "pending"`` — poll with ``client.get_run(run_id)``
        to check when final pass/fail verdicts are available.

        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.
            on_environment: Optional callback for environment actions.
            max_workers: Number of scenarios to run in parallel (default 1).
            **deploy_kwargs: Extra kwargs passed to ``RunBuilder.deploy()``.

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