"""Pairwise evaluation runner and comparison utilities."""

from __future__ import annotations

import asyncio

from agentbyte.cancellation_token import CancellationToken

from agentbyte.eval.base import BaseEvalTarget
from agentbyte.eval.judges.pairwise import PairwiseJudge
from agentbyte.eval.types import EvalTask, PairwiseResult


class PairwiseRunner:
    """Run two targets against the same task suite and compare them pairwise.

    For each task both targets are executed concurrently, then ``judge.compare()``
    is called on the resulting trajectories. A target failure produces a failed
    trajectory; the judge still runs and will typically prefer the non-failed side.
    """

    def __init__(self, judge: PairwiseJudge, parallel: bool = True) -> None:
        self.judge = judge
        self.parallel = parallel

    async def compare(
        self,
        target_a: BaseEvalTarget,
        target_b: BaseEvalTarget,
        tasks: list[EvalTask],
        cancellation_token: CancellationToken | None = None,
    ) -> list[PairwiseResult]:
        """Run both targets on every task and return one PairwiseResult per task."""
        if self.parallel:
            return list(
                await asyncio.gather(
                    *[
                        self._compare_single(
                            target_a,
                            target_b,
                            task,
                            cancellation_token=cancellation_token,
                        )
                        for task in tasks
                    ]
                )
            )

        results: list[PairwiseResult] = []
        for task in tasks:
            if cancellation_token and cancellation_token.is_cancelled():
                break
            results.append(
                await self._compare_single(
                    target_a,
                    target_b,
                    task,
                    cancellation_token=cancellation_token,
                )
            )
        return results

    async def _compare_single(
        self,
        target_a: BaseEvalTarget,
        target_b: BaseEvalTarget,
        task: EvalTask,
        *,
        cancellation_token: CancellationToken | None,
    ) -> PairwiseResult:
        traj_a, traj_b = await asyncio.gather(
            target_a.run(task, cancellation_token),
            target_b.run(task, cancellation_token),
        )
        return await self.judge.compare(traj_a, traj_b, cancellation_token)


def compare_pairwise(
    results: list[PairwiseResult],
    label_a: str = "baseline",
    label_b: str = "candidate",
) -> dict[str, float | str | int]:
    """Aggregate pairwise results into a win/loss summary.

    Returns a dict with keys:
    - ``label_a``, ``label_b``
    - ``wins_a``, ``wins_b``, ``ties``, ``total``
    - ``win_rate_a``, ``win_rate_b``
    """
    wins_a = sum(1 for r in results if r.winner == "a")
    wins_b = sum(1 for r in results if r.winner == "b")
    ties = sum(1 for r in results if r.winner == "tie")
    total = len(results)
    return {
        "label_a": label_a,
        "label_b": label_b,
        "wins_a": wins_a,
        "wins_b": wins_b,
        "ties": ties,
        "total": total,
        "win_rate_a": wins_a / total if total else 0.0,
        "win_rate_b": wins_b / total if total else 0.0,
    }


__all__ = ["PairwiseRunner", "compare_pairwise"]
