"""Composite evaluation judge."""

from __future__ import annotations

import asyncio
from collections.abc import Sequence

from agentbyte.cancellation_token import CancellationToken

from agentbyte.eval.base import BaseEvalJudge
from agentbyte.eval.types import EvalScore, EvalTrajectory


class CompositeJudge(BaseEvalJudge):
    """Combine multiple judges into one weighted score."""

    def __init__(
        self,
        judges: Sequence[tuple[BaseEvalJudge, float]],
        *,
        name: str = "Composite",
        normalize_weights: bool = True,
    ) -> None:
        if not judges:
            raise ValueError("CompositeJudge requires at least one sub-judge")
        if normalize_weights:
            total_weight = sum(weight for _, weight in judges)
            if total_weight <= 0:
                raise ValueError("CompositeJudge weights must sum to > 0")
            normalized = [(judge, weight / total_weight) for judge, weight in judges]
        else:
            total_weight = sum(weight for _, weight in judges)
            if abs(total_weight - 1.0) > 0.01:
                raise ValueError(
                    "CompositeJudge weights must sum to 1.0 when normalize_weights=False"
                )
            normalized = list(judges)

        super().__init__(name=name, answer_strategy="last_non_empty")
        self.judges = normalized

    async def score(
        self,
        trajectory: EvalTrajectory,
        criteria: list[str] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> EvalScore:
        results = await asyncio.gather(
            *[
                judge.score(trajectory, criteria, cancellation_token)
                for judge, _ in self.judges
            ]
        )

        dimensions: dict[str, float] = {}
        dimension_weight_totals: dict[str, float] = {}
        reasoning: dict[str, str] = {}
        metadata_sub_judges: list[dict[str, float | str]] = []
        weighted_overall = 0.0

        for (judge, weight), result in zip(self.judges, results):
            metadata_sub_judges.append(
                {"name": judge.name, "weight": weight, "score": result.overall}
            )
            weighted_overall += result.overall * weight
            for dimension, score in result.dimensions.items():
                dimensions[dimension] = dimensions.get(dimension, 0.0) + (score * weight)
                dimension_weight_totals[dimension] = (
                    dimension_weight_totals.get(dimension, 0.0) + weight
                )
                detail = result.reasoning.get(dimension, "No reasoning provided")
                reasoning[f"{dimension} ({judge.name})"] = (
                    f"[weight: {weight:.2f}] {detail}"
                )

        normalized_dimensions = {
            dimension: dimensions[dimension] / max(dimension_weight_totals[dimension], 1e-9)
            for dimension in dimensions
        }
        return EvalScore(
            overall=weighted_overall,
            dimensions=normalized_dimensions,
            reasoning=reasoning,
            trajectory=trajectory,
            metadata={"judge": self.name, "sub_judges": metadata_sub_judges},
        )


__all__ = ["CompositeJudge"]
