from __future__ import annotations

from pathlib import Path

from ..models import RolloutTrace
from ..storage import atomic_write_jsonl
from .agentrm import AgentRM
from .pre_score import PreScorer, adjust_agent_score
from .identity import composite_id


class TraceScorer:
    """Resolve pre-scores and AgentRM scores into one effective score per trace."""

    def __init__(
        self,
        pre_scorer: PreScorer,
        agentrm: AgentRM | None = None,
    ):
        self.pre_scorer = pre_scorer
        self.agentrm = agentrm if agentrm is not None else AgentRM()

    def score_all(
        self,
        traces: list[RolloutTrace],
        final_score_dir: Path,
        agentrm_dir: Path | None = None,
    ) -> dict[str, float]:
        final_score_dir.mkdir(parents=True, exist_ok=True)
        agentrm_dir = agentrm_dir or final_score_dir
        agentrm_dir.mkdir(parents=True, exist_ok=True)
        results = self.pre_scorer.score_all(traces)
        pre_scores = {result.trace_id: result for result in results}
        if set(pre_scores) != {trace.trace_id for trace in traces}:
            raise ValueError("pre-scorer returned incomplete or duplicate results")
        atomic_write_jsonl(
            final_score_dir / "pre_scores.jsonl",
            [pre_scores[trace.trace_id].to_dict() for trace in traces],
        )
        agentrm_traces = [
            trace for trace in traces if pre_scores[trace.trace_id].route == "agentrm"
        ]
        requests = [trace.agentrm_request() for trace in agentrm_traces]
        local_scores = agentrm_dir / "agentrm_scores.jsonl"
        found = self.agentrm.score_requests(requests)
        atomic_write_jsonl(local_scores, found)

        by_key = {composite_id(row): float(row["score"]) for row in found}
        resolved = []
        scores: dict[str, float] = {}
        for trace in traces:
            pre_score = pre_scores[trace.trace_id]
            if pre_score.route == "fixed_score":
                if pre_score.score is None:
                    raise ValueError(f"fixed pre-score is missing for {trace.trace_id}")
                score = pre_score.score
                source = "pre_score"
            else:
                score = adjust_agent_score(by_key[trace.key.as_tuple()], pre_score)
                source = "agentrm"
            scores[trace.trace_id] = score
            resolved.append({
                "task_name": trace.task_name,
                "compile_type": trace.compile_type,
                "test_name": trace.test_name,
                "trace_id": trace.trace_id,
                "effective_score": score,
                "score_source": source,
                "pre_score_reason": pre_score.reason,
                "agentrm_score": by_key.get(trace.key.as_tuple()),
                "efficiency_cost": pre_score.efficiency_cost,
                "efficiency_penalty": pre_score.efficiency_penalty,
            })
        atomic_write_jsonl(final_score_dir / "effective_scores.jsonl", resolved)
        return scores