from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable

from rl_harness.types import RolloutStep


class JsonlTrajectoryWriter:
    def __init__(self, path: str | Path):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._fh = self.path.open("a", encoding="utf-8")

    def write(self, step: RolloutStep):
        self._fh.write(json.dumps(step.to_json(), ensure_ascii=False) + "\n")
        self._fh.flush()

    def close(self):
        self._fh.close()


def summarize_rollout(steps: Iterable[RolloutStep]) -> dict:
    rows = list(steps)
    total = len(rows)
    reward = sum(s.reward.value for s in rows)
    exact = sum(1 for s in rows if s.reward.details.get("exact"))
    action = sum(1 for s in rows if s.reward.details.get("action_ok"))
    by_group: dict[str, list[int]] = {}
    for s in rows:
        group = s.reward.details.get("group", "unknown")
        hit, count = by_group.get(group, [0, 0])
        by_group[group] = [hit + int(bool(s.reward.details.get("exact"))), count + 1]
    return {
        "total": total,
        "mean_reward": reward / max(total, 1),
        "exact": exact / max(total, 1),
        "action": action / max(total, 1),
        "by_group": {k: {"hit": v[0], "total": v[1], "acc": v[0] / max(v[1], 1)} for k, v in by_group.items()},
    }