"""
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

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,
        action_index: int | None = None,
    ) -> TestBuilder:
        """Record an agent tool call action with expected vs actual comparison.

        Args:
            expected: Expected tool call (tool_name, arguments).
            actual: Actual tool call made by the agent.
            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)
        self._action_results.append({
            "actor": "agent",
            "action_type": "tool_call",
            "action_index": idx,
            "expected": expected,
            "actual": actual,
            "match_status": match_status,
            "divergence_notes": divergence_notes,
        })
        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 completed.

        Args:
            status: Final status - "completed" or "failed".
        """
        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_...", base_url="https://...")
        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, tenant_id=1, 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 completed.

        Args:
            status: Final status - "completed" or "failed".
        """
        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,
        tenant_id: int,
        dataset_id: int,
        runner_id: int | None = None,
    ) -> dict[str, Any]:
        """Build the result and submit it as a new run via the API.

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

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


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."""
    total = len(tests)
    passed = 0
    failed = 0
    total_tool_call_divergence = 0
    total_response_divergence = 0
    similarity_scores: list[float] = []

    for test in tests:
        if test.get("status") == "completed":
            passed += 1
        elif test.get("status") == "failed":
            failed += 1

        for action in test.get("action_results", []):
            match_status = action.get("match_status")
            action_type = action.get("action_type")

            if action_type == "tool_call" and match_status and match_status != "exact":
                total_tool_call_divergence += 1
            elif action_type == "text" and action.get("actor") == "agent":
                if match_status and match_status != "exact":
                    total_response_divergence += 1
                if "semantic_similarity" in action:
                    similarity_scores.append(action["semantic_similarity"])

    avg_similarity = (
        round(sum(similarity_scores) / len(similarity_scores), 2) if similarity_scores else None
    )

    return {
        "total_tests": total,
        "tests_passed": passed,
        "tests_failed": failed,
        "average_similarity_score": avg_similarity,
        "total_tool_call_divergence": total_tool_call_divergence,
        "total_response_divergence": total_response_divergence,
    }
