"""
Run builder for incrementally constructing test run results.

Provides a builder pattern for constructing run result objects as an agent
executes tests. Once complete, the result can be deployed via the client.
"""

from __future__ import annotations

import warnings
from datetime import datetime, timezone
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from .client import AshrLabsClient


class TestBuilder:
    """Builds a single test result incrementally as the agent runs through actions."""

    def __init__(self, test_id: str) -> None:
        self._test_id = test_id
        self._status = "pending"
        self._started_at: str | None = None
        self._completed_at: str | None = None
        self._action_results: list[dict[str, Any]] = []
        self._next_action_index = 0

    def start(self) -> TestBuilder:
        """Mark this test as started."""
        self._status = "running"
        self._started_at = _now()
        return self

    def add_user_file(
        self,
        file_path: str,
        description: str,
        action_index: int | None = None,
    ) -> TestBuilder:
        """Record a user file input action."""
        idx = self._resolve_index(action_index)
        self._action_results.append({
            "actor": "user",
            "action_type": "file",
            "action_index": idx,
            "description": description,
            "file_path": file_path,
            "input_provided": True,
        })
        return self

    def add_user_text(
        self,
        text: str,
        description: str,
        action_index: int | None = None,
    ) -> TestBuilder:
        """Record a user text input action."""
        idx = self._resolve_index(action_index)
        self._action_results.append({
            "actor": "user",
            "action_type": "text",
            "action_index": idx,
            "description": description,
            "text": text,
            "input_provided": True,
        })
        return self

    def add_tool_call(
        self,
        expected: dict[str, Any],
        actual: dict[str, Any],
        match_status: str,
        divergence_notes: str | None = None,
        argument_comparison: dict[str, Any] | None = None,
        action_index: int | None = None,
    ) -> TestBuilder:
        """Record an agent tool call action with expected vs actual comparison.

        Multiple calls with the same ``action_index`` are batched into a single
        action result with a ``tool_calls`` array (the structure the frontend
        expects).

        Args:
            expected: Expected tool call. Accepts either ``{"name", "arguments_json"}``
                (dataset format) or ``{"tool_name", "arguments"}`` (display format).
                Both are normalized automatically.
            actual: Actual tool call made by the agent. Same format flexibility as expected.
            match_status: One of "exact", "partial", "mismatch".
            divergence_notes: Optional notes explaining the divergence.
            action_index: Explicit index, or auto-incremented if omitted.
        """
        idx = self._resolve_index(action_index)
        tc_entry: dict[str, Any] = {
            "expected": _normalize_tool_call(expected),
            "actual": _normalize_tool_call(actual),
            "match_status": match_status,
            "divergence_notes": divergence_notes,
        }
        if argument_comparison is not None:
            tc_entry["argument_comparison"] = argument_comparison
        else:
            warnings.warn(
                f"Tool call '{expected.get('name', '?')}' has no "
                f"argument_comparison — the backend grader may skip it. "
                f"Use compare_args_structural() to generate this field.",
                stacklevel=2,
            )

        # Find an existing tool_call action_result for this action_index
        existing = None
        for ar in self._action_results:
            if (
                ar.get("action_type") == "tool_call"
                and ar.get("action_index") == idx
            ):
                existing = ar
                break

        if existing is not None:
            existing["tool_calls"].append(tc_entry)
        else:
            self._action_results.append({
                "actor": "agent",
                "action_type": "tool_call",
                "action_index": idx,
                "tool_calls": [tc_entry],
            })
        return self

    def add_agent_response(
        self,
        expected_response: dict[str, Any],
        actual_response: dict[str, Any],
        match_status: str,
        semantic_similarity: float | None = None,
        divergence_notes: str | None = None,
        action_index: int | None = None,
    ) -> TestBuilder:
        """Record an agent text response with expected vs actual comparison.

        Args:
            expected_response: The expected response content.
            actual_response: The actual response from the agent.
            match_status: One of "exact", "similar", "divergent".
            semantic_similarity: Optional similarity score (0.0 to 1.0).
            divergence_notes: Optional notes explaining the divergence.
            action_index: Explicit index, or auto-incremented if omitted.
        """
        idx = self._resolve_index(action_index)
        result: dict[str, Any] = {
            "actor": "agent",
            "action_type": "text",
            "action_index": idx,
            "expected_response": expected_response,
            "actual_response": actual_response,
            "match_status": match_status,
            "divergence_notes": divergence_notes,
        }
        if semantic_similarity is not None:
            result["semantic_similarity"] = semantic_similarity
        self._action_results.append(result)
        return self

    def complete(self, status: str = "completed") -> TestBuilder:
        """Mark this test as finished executing.

        Args:
            status: Execution status. Defaults to ``"completed"``. The backend
                grader will update this to ``"failed"`` if tool calls or
                responses diverge. Use ``"failed"`` if the agent errored
                during execution.
        """
        self._status = status
        self._completed_at = _now()
        return self

    def build(self) -> dict[str, Any]:
        """Serialize this test to a dict matching the run result schema."""
        result: dict[str, Any] = {
            "test_id": self._test_id,
            "status": self._status,
            "action_results": list(self._action_results),
        }
        if self._started_at:
            result["started_at"] = self._started_at
        if self._completed_at:
            result["completed_at"] = self._completed_at
        return result

    def _resolve_index(self, explicit: int | None) -> int:
        if explicit is not None:
            self._next_action_index = explicit + 1
            return explicit
        idx = self._next_action_index
        self._next_action_index += 1
        return idx


class RunBuilder:
    """Incrementally builds a run result object over the lifetime of a test execution.

    Example:
        ```python
        from ashr_labs import AshrLabsClient, RunBuilder

        client = AshrLabsClient(api_key="tp_...")
        run = RunBuilder()
        run.start()

        test = run.add_test("bank_analysis")
        test.start()
        test.add_user_file(file_path="datasets/.../action_0.pdf", description="Upload PDF")
        test.add_user_text(text="Analyze this", description="User asks for analysis")
        test.add_tool_call(
            expected={"tool_name": "extract_pdf", "arguments": {"file": "a.pdf"}},
            actual={"tool_name": "extract_pdf", "arguments": {"file": "a.pdf", "pages": "all"}},
            match_status="partial",
            divergence_notes="Extra pages argument",
        )
        test.add_agent_response(
            expected_response={"summary": "Expected text..."},
            actual_response={"summary": "Actual text..."},
            match_status="similar",
            semantic_similarity=0.89,
        )
        test.complete()

        run.complete()
        run.deploy(client, dataset_id=42)
        ```
    """

    def __init__(self) -> None:
        self._status = "pending"
        self._started_at: str | None = None
        self._completed_at: str | None = None
        self._tests: list[TestBuilder] = []

    def start(self) -> RunBuilder:
        """Mark the run as started."""
        self._status = "running"
        self._started_at = _now()
        return self

    def add_test(self, test_id: str) -> TestBuilder:
        """Create and register a new test within this run.

        Args:
            test_id: Unique identifier for the test case.

        Returns:
            A TestBuilder for incrementally building the test result.
        """
        test = TestBuilder(test_id)
        self._tests.append(test)
        return test

    def complete(self, status: str = "completed") -> RunBuilder:
        """Mark the run as finished executing.

        Args:
            status: Execution status. Defaults to ``"completed"``.
                The backend grader re-scores results asynchronously after
                deploy (typically 1-3 minutes).
        """
        self._status = status
        self._completed_at = _now()
        return self

    def build(self) -> dict[str, Any]:
        """Serialize the full run result to a dict.

        Returns:
            A dict matching the run result schema, ready to be passed
            to ``client.create_run(result=...)``.
        """
        tests = [t.build() for t in self._tests]
        return {
            "tests": tests,
            "status": self._status,
            "started_at": self._started_at,
            "completed_at": self._completed_at,
            "aggregate_metrics": _compute_aggregate(tests),
        }

    def deploy(
        self,
        client: AshrLabsClient,
        dataset_id: int,
        tenant_id: int | None = None,
        runner_id: int | None = None,
    ) -> dict[str, Any]:
        """Build the result and submit it as a new run via the API.

        After submission, the backend grader asynchronously scores tool call
        arguments and text responses (typically 1-3 minutes). The returned
        run object will have ``grading_status: "pending"`` — use
        ``client.get_run(run_id)`` to poll for the final graded result.

        Args:
            client: An authenticated AshrLabsClient instance.
            dataset_id: The dataset this run is for.
            tenant_id: The tenant to create the run under (auto-resolved if omitted).
            runner_id: Optional ID of the user who ran the test.

        Returns:
            The created run object from the API (grading still pending).
        """
        result = self.build()
        return client.create_run(
            dataset_id=dataset_id,
            result=result,
            tenant_id=tenant_id,
            runner_id=runner_id,
        )


def _normalize_tool_call(tc: dict[str, Any]) -> dict[str, Any]:
    """Normalize a tool call dict so it contains the fields the platform expects.

    Accepts either the dataset format (``name``, ``arguments_json``) or the
    display format (``tool_name``, ``arguments``) and ensures both are present.
    """
    import json as _json

    out = dict(tc)

    # Ensure tool_name is set (frontend reads this field)
    if "tool_name" not in out and "name" in out:
        out["tool_name"] = out["name"]
    elif "name" not in out and "tool_name" in out:
        out["name"] = out["tool_name"]

    # Ensure arguments is a dict (frontend reads this field)
    if "arguments" not in out and "arguments_json" in out:
        try:
            val = out["arguments_json"]
            out["arguments"] = _json.loads(val) if isinstance(val, str) else val
        except (_json.JSONDecodeError, TypeError):
            out["arguments"] = {}
    elif "arguments_json" not in out and "arguments" in out:
        args = out["arguments"]
        out["arguments_json"] = _json.dumps(args) if isinstance(args, dict) else str(args)

    return out


def _now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _compute_aggregate(tests: list[dict[str, Any]]) -> dict[str, Any]:
    """Compute aggregate metrics from a list of built test dicts.

    The backend grader re-scores results asynchronously after deploy
    (typically 1-3 minutes). It updates ``tests_passed``, ``tests_failed``,
    and similarity/divergence metrics. Poll with ``client.get_run(run_id)``
    to check when final graded results are available.
    """
    total = len(tests)
    errored = sum(1 for t in tests if t.get("status") == "failed")

    return {
        "total_tests": total,
        "tests_completed": total - errored,
        "tests_errored": errored,
    }
