# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""Core benchmark definitions and runner logic."""

import asyncio
import dataclasses
import datetime
import enum
import logging
import time
import typing
from collections.abc import Callable, Sequence
from typing import Any

from skill_eval import exceptions, scenario, scorer


class Timer:
    """A context manager for timing blocks of code."""

    def __init__(self) -> None:
        self.duration = datetime.timedelta()
        self._start_time = None

    def __enter__(self) -> typing.Any:
        self._start_time = time.perf_counter()
        return self

    def __exit__(
        self, exc_type: typing.Any, exc_val: typing.Any, exc_tb: typing.Any
    ) -> typing.Any:
        if self._start_time is not None:
            elapsed = time.perf_counter() - self._start_time
            self.duration = datetime.timedelta(seconds=elapsed)
            self._start_time = None

    @property
    def elapsed(self) -> datetime.timedelta:
        if self._start_time is None:
            return self.duration
        return datetime.timedelta(
            seconds=time.perf_counter() - self._start_time
        )


class ExecutionStatus(enum.Enum):
    """Status of a benchmark execution."""

    PENDING = "PENDING"
    INITIALIZING = "INITIALIZING"
    CONVERSING = "CONVERSING"
    GRADING = "GRADING"
    RUNNING = "RUNNING"
    FINISHED = "FINISHED"
    FAILED = "FAILED"


@dataclasses.dataclass(frozen=True)
class ToolInteraction:
    """Represents a single tool interaction within a turn."""

    name: str
    args: dict[str, Any]
    output: str | None = None
    thought: str | None = None


@dataclasses.dataclass(frozen=True)
class TrajectoryEvent:
    """Base class for any tracked event during a conversation turn."""


@dataclasses.dataclass(frozen=True)
class ToolCallEvent(TrajectoryEvent):
    """Represents the intent to call a tool."""

    call_id: str
    name: str
    args: dict[str, Any]
    source: str = "main"


@dataclasses.dataclass(frozen=True)
class ToolResultEvent(TrajectoryEvent):
    """Represents the outcome of a tool execution."""

    call_id: str
    output: str
    is_error: bool = False
    source: str = "main"


@dataclasses.dataclass(frozen=True)
class AgentMessageEvent(TrajectoryEvent):
    """Represents text generated by the agent (thought or spoken)."""

    text: str
    is_thought: bool = False
    source: str = "main"


@dataclasses.dataclass(frozen=True)
class SubagentEvent(TrajectoryEvent):
    """Represents subagent lifecycle events."""

    subagent_id: str
    action: str  # e.g., "spawned", "finished"
    prompt: str | None = None
    source: str = "main"


@dataclasses.dataclass(frozen=True)
class TurnMetrics:
    """Metrics for a single conversation turn."""

    latency_sec: float
    simulator_latency_sec: float
    tool_calls: int
    tool_interactions: list[ToolInteraction] = dataclasses.field(
        default_factory=list
    )
    events: list[TrajectoryEvent] = dataclasses.field(default_factory=list)


@dataclasses.dataclass(frozen=True)
class RubricCriterion:
    """A single scoring criterion from the rubric."""

    criteria: str
    score: int
    reasoning: str


@dataclasses.dataclass(frozen=True)
class RubricResult:
    """The full result of a rubric-based evaluation."""

    scores: list[RubricCriterion]
    summary: str
    total_score: int
    max_score: int


@dataclasses.dataclass(frozen=True)
class ConversationResult:
    """The final result of an agent benchmarking run."""

    scenario_name: str
    scenario_text: str
    head_name: str
    turns: int
    total_time_sec: float
    messages: list[dict[str, Any]]
    turn_metrics: list[TurnMetrics]
    status: ExecutionStatus = ExecutionStatus.FINISHED
    success: bool | None = None
    failure_reason: str | None = None
    rubric_results: RubricResult | None = None
    scoring_latency_sec: float = 0.0
    start_time: float = 0.0
    end_time: float = 0.0
    log_files: list[str] = dataclasses.field(default_factory=list)
    trajectory_id: str | None = None

    # New Orchestration Metrics:
    init_queued_latency: datetime.timedelta = datetime.timedelta()
    init_active_latency: datetime.timedelta = datetime.timedelta()
    conversing_latency: datetime.timedelta = datetime.timedelta()
    scoring_latency: datetime.timedelta = datetime.timedelta()
    cleanup_latency: datetime.timedelta = datetime.timedelta()


class BaseAgentHead:
    """Base class for an Agent Head (MetaAgent, Antigravity, or Fake)."""

    async def send_message(self, message: str) -> str:
        """Sends a message to the agent and returns its response."""
        raise NotImplementedError()

    def get_tool_calls_count_last_turn(self) -> int:
        """Returns the number of tool calls made in the last turn."""
        raise NotImplementedError()

    def get_tool_interactions_last_turn(self) -> Sequence[ToolInteraction]:
        """Returns detailed tool interactions from the last turn."""
        return []

    def get_events_last_turn(self) -> Sequence[TrajectoryEvent]:
        """Returns detailed trajectory events from the last turn."""
        return []

    def get_log_files(self) -> Sequence[str]:
        """Returns a list of diagnostic log files associated with this head."""
        return []

    def get_trajectory_id(self) -> str | None:
        """Returns the trajectory ID associated with this head."""
        return None

    async def initialize(self) -> None:
        """Initializes the agent's internal state/history."""
        raise NotImplementedError()

    @property
    def name(self) -> str:
        """Returns the display name of this head."""
        raise NotImplementedError()

    async def close(self) -> None:
        """Performs final cleanup (closing connections, stopping processes)."""
        pass


class BenchmarkRunner:
    """Runs a benchmark scenario against a specific agent head."""

    def __init__(
        self,
        scen: scenario.Scenario,
        scorer_instance: Any | None = None,
    ) -> None:
        self.scenario = scen
        self.scenario_name = scen.name
        self.scenario_text = scen.text
        self.scorer = scorer_instance

    async def _notify_progress(
        self,
        callback: Callable[[ConversationResult], Any],
        result: ConversationResult,
    ) -> None:
        """Safely invokes the progress callback."""
        try:
            await callback(result)
        except Exception:
            logging.exception("Failed to update progress report")

    def _create_result(
        self,
        head: BaseAgentHead,
        turns: int,
        duration: float,
        messages: list[dict[str, Any]],
        metrics: list[TurnMetrics],
        success: bool | None = None,
        status: ExecutionStatus = ExecutionStatus.FINISHED,
        failure_reason: str | None = None,
        rubric_results: RubricResult | None = None,
        scoring_latency_sec: float = 0.0,
        start_time: float = 0.0,
        end_time: float = 0.0,
        init_queued_latency: datetime.timedelta = datetime.timedelta(),
        init_active_latency: datetime.timedelta = datetime.timedelta(),
        conversing_latency: datetime.timedelta = datetime.timedelta(),
        scoring_latency: datetime.timedelta = datetime.timedelta(),
        cleanup_latency: datetime.timedelta = datetime.timedelta(),
    ) -> ConversationResult:
        """Helper to create a ConversationResult with current head state."""
        return ConversationResult(
            scenario_name=self.scenario_name,
            scenario_text=self.scenario_text,
            head_name=head.name,
            turns=turns,
            total_time_sec=duration,
            messages=messages,
            turn_metrics=metrics,
            status=status,
            success=success,
            failure_reason=failure_reason,
            rubric_results=rubric_results,
            scoring_latency_sec=scoring_latency_sec,
            start_time=start_time,
            end_time=end_time,
            log_files=list(head.get_log_files()),
            trajectory_id=head.get_trajectory_id(),
            init_queued_latency=init_queued_latency,
            init_active_latency=init_active_latency,
            conversing_latency=conversing_latency,
            scoring_latency=scoring_latency,
            cleanup_latency=cleanup_latency,
        )

    async def run(
        self,
        agent_head: BaseAgentHead,
        on_turn_completed: Callable[[ConversationResult], Any] | None = None,
        init_semaphore: asyncio.Semaphore | None = None,
    ) -> ConversationResult:
        """Runs the benchmark and returns the final results."""

        init_queued_timer = Timer()
        init_active_timer = Timer()
        conversing_timer = Timer()
        scoring_timer = Timer()

        async def _notify_state(status: ExecutionStatus) -> None:
            if on_turn_completed:
                await self._notify_progress(
                    on_turn_completed,
                    self._create_result(
                        agent_head,
                        0,
                        0.0,
                        [],
                        [],
                        status=status,
                    ),
                )

        if init_semaphore:
            await _notify_state(ExecutionStatus.INITIALIZING)
            with init_queued_timer:
                await init_semaphore.acquire()
            try:
                logging.info(
                    "[%s] [%s] Initializing head...",
                    self.scenario_name,
                    agent_head.name,
                )
                with init_active_timer:
                    await agent_head.initialize()
            finally:
                init_semaphore.release()
        else:
            await _notify_state(ExecutionStatus.INITIALIZING)
            with init_active_timer:
                await agent_head.initialize()

        await _notify_state(ExecutionStatus.CONVERSING)

        messages: list[dict[str, Any]] = []
        turn_metrics_list: list[TurnMetrics] = []
        total_start_time = time.time()

        user_message = self.scenario.prompt
        messages.append({"role": "user", "content": user_message})

        try:
            with conversing_timer:
                current_agent_response = await agent_head.send_message(
                    user_message
                )

            turn_metrics_list.append(
                TurnMetrics(
                    latency_sec=conversing_timer.duration.total_seconds(),
                    simulator_latency_sec=0.0,
                    tool_calls=agent_head.get_tool_calls_count_last_turn(),
                    tool_interactions=list(
                        agent_head.get_tool_interactions_last_turn()
                    ),
                    events=list(agent_head.get_events_last_turn()),
                )
            )

            messages.append(
                {"role": "agent", "content": current_agent_response}
            )

            if on_turn_completed:
                await self._notify_progress(
                    on_turn_completed,
                    self._create_result(
                        agent_head,
                        1,
                        time.time() - total_start_time,
                        messages,
                        turn_metrics_list,
                        status=ExecutionStatus.CONVERSING,
                        start_time=total_start_time,
                        init_queued_latency=init_queued_timer.duration,
                        init_active_latency=init_active_timer.duration,
                        conversing_latency=conversing_timer.duration,
                    ),
                )

            # FINAL STEP: The Scorer (Auditor) decides Success/Failure.
            if self.scorer is None:
                total_end_time = time.time()
                return self._create_result(
                    agent_head,
                    1,
                    total_end_time - total_start_time,
                    messages,
                    turn_metrics_list,
                    success=True,  # Default success if no scorer
                    status=ExecutionStatus.FINISHED,
                    failure_reason=None,
                    rubric_results=None,
                    scoring_latency_sec=0.0,
                    start_time=total_start_time,
                    end_time=total_end_time,
                    init_queued_latency=init_queued_timer.duration,
                    init_active_latency=init_active_timer.duration,
                    conversing_latency=conversing_timer.duration,
                )

            # Notify grading state but keep history
            if on_turn_completed:
                await self._notify_progress(
                    on_turn_completed,
                    self._create_result(
                        agent_head,
                        1,
                        time.time() - total_start_time,
                        messages,
                        turn_metrics_list,
                        status=ExecutionStatus.GRADING,
                        start_time=total_start_time,
                        init_queued_latency=init_queued_timer.duration,
                        init_active_latency=init_active_timer.duration,
                        conversing_latency=conversing_timer.duration,
                    ),
                )

            logging.info("[%s] Grading with rubric...", agent_head.name)
            structured_turns = [
                scorer.ScorerTurn(
                    user_message=user_message,
                    agent_response=current_agent_response,
                    tool_interactions=turn_metrics_list[0].tool_interactions,
                    events=turn_metrics_list[0].events,
                )
            ]
            with scoring_timer:
                rubric_results = await self.scorer.grade_conversation(
                    self.scenario, structured_turns
                )
            # Success logic: Must be 100%
            if rubric_results.max_score > 0:
                success = rubric_results.total_score >= rubric_results.max_score
            else:
                success = False

            total_end_time = time.time()
            return self._create_result(
                agent_head,
                1,
                total_end_time - total_start_time,
                messages,
                turn_metrics_list,
                success=success,
                status=ExecutionStatus.FINISHED,
                failure_reason=None,
                rubric_results=rubric_results,
                scoring_latency_sec=scoring_timer.duration.total_seconds(),
                start_time=total_start_time,
                end_time=total_end_time,
                init_queued_latency=init_queued_timer.duration,
                init_active_latency=init_active_timer.duration,
                conversing_latency=conversing_timer.duration,
                scoring_latency=scoring_timer.duration,
            )

        except exceptions.ScorerError as e:
            logging.exception("[%s] Scorer failure", agent_head.name)
            return self._create_result(
                agent_head,
                1,
                time.time() - total_start_time,
                messages,
                turn_metrics_list,
                success=False,
                status=ExecutionStatus.FAILED,
                failure_reason=f"[Scorer] {e!r}",
                start_time=total_start_time,
                end_time=time.time(),
                init_queued_latency=init_queued_timer.duration,
                init_active_latency=init_active_timer.duration,
                conversing_latency=conversing_timer.duration,
                scoring_latency=scoring_timer.duration,
            )
        except Exception as e:
            logging.exception("[%s] Error during benchmark", agent_head.name)
            return self._create_result(
                agent_head,
                1,
                time.time() - total_start_time,
                messages,
                turn_metrics_list,
                success=False,
                status=ExecutionStatus.FAILED,
                failure_reason=repr(e),
                start_time=total_start_time,
                end_time=time.time(),
                init_queued_latency=init_queued_timer.duration,
                init_active_latency=init_active_timer.duration,
                conversing_latency=conversing_timer.duration,
                scoring_latency=scoring_timer.duration,
            )