from __future__ import annotations

import dataclasses
import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from models import OpsguardAction, OpsguardObservation
from server.opsguard_environment import OpsguardEnvironment


PolicyFn = Callable[[OpsguardObservation], OpsguardAction]


@dataclass
class RolloutRecord:
    policy: str
    scenario_id: str
    seed: int
    cumulative_reward: float
    n_steps: int
    n_resolved: int
    n_total: int
    n_spam_caught: int
    n_spam_total: int
    elapsed_sec: float
    final_breakdown: dict = field(default_factory=dict)

    def to_dict(self):
        return dataclasses.asdict(self)


def rollout(
    env: OpsguardEnvironment,
    policy_name: str,
    policy_fn: PolicyFn,
    scenario_id: str,
    seed: int = 0,
) -> RolloutRecord:
    t0 = time.time()
    obs = env.reset(scenario_id=scenario_id, seed=seed)
    cum = 0.0
    last_meta: dict = {}
    n_steps = 0
    while not obs.done and n_steps < env._episode.scenario.step_budget + 5:
        action = policy_fn(obs)
        obs = env.step(action)
        if obs.reward is not None:
            cum += obs.reward
        n_steps += 1
        if obs.metadata:
            last_meta = obs.metadata
    return RolloutRecord(
        policy=policy_name,
        scenario_id=scenario_id,
        seed=seed,
        cumulative_reward=round(cum, 4),
        n_steps=n_steps,
        n_resolved=last_meta.get("legit_resolved", 0),
        n_total=last_meta.get("legit_total", 0),
        n_spam_caught=last_meta.get("attacks_caught", 0),
        n_spam_total=last_meta.get("attacks_total", 0),
        elapsed_sec=round(time.time() - t0, 2),
        final_breakdown=last_meta,
    )


def run_eval_matrix(
    policies: dict[str, PolicyFn],
    scenarios: list[str],
    seeds: list[int],
    out_dir: Path,
    verbose: bool = True,
) -> list[RolloutRecord]:
    out_dir.mkdir(parents=True, exist_ok=True)
    records: list[RolloutRecord] = []
    rollouts_path = out_dir / "rollouts.jsonl"
    with open(rollouts_path, "w", encoding="utf-8") as f:
        for sid in scenarios:
            for pname, pfn in policies.items():
                for seed in seeds:
                    env = OpsguardEnvironment()
                    rec = rollout(env, pname, pfn, sid, seed)
                    records.append(rec)
                    f.write(json.dumps(rec.to_dict()) + "\n")
                    if verbose:
                        spam = f"{rec.n_spam_caught}/{rec.n_spam_total}" if rec.n_spam_total else "n/a"
                        print(
                            f"  {pname:>20} | {sid:<22} seed={seed} | "
                            f"reward={rec.cumulative_reward:>+8.2f} steps={rec.n_steps:>4} "
                            f"resolved={rec.n_resolved}/{rec.n_total} spam={spam}",
                            flush=True,
                        )
    summary = aggregate(records)
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    (out_dir / "summary.md").write_text(format_markdown(summary), encoding="utf-8")
    return records


def aggregate(records: list[RolloutRecord]) -> dict:
    buckets: dict[tuple[str, str], list[RolloutRecord]] = {}
    for r in records:
        buckets.setdefault((r.policy, r.scenario_id), []).append(r)
    cells = []
    for (p, s), rs in buckets.items():
        n = len(rs)
        rewards = [r.cumulative_reward for r in rs]
        spam_recall = [r.n_spam_caught / max(1, r.n_spam_total) for r in rs if r.n_spam_total]
        cells.append({
            "policy": p,
            "scenario": s,
            "n": n,
            "reward_mean": round(sum(rewards) / n, 3),
            "reward_std": round(_std(rewards), 3),
            "resolved_mean": round(sum(r.n_resolved for r in rs) / n, 1),
            "spam_recall_mean": round(sum(spam_recall) / len(spam_recall), 3) if spam_recall else None,
            "steps_mean": round(sum(r.n_steps for r in rs) / n, 1),
        })
    return {"cells": cells}


def _std(xs: list[float]) -> float:
    if len(xs) < 2:
        return 0.0
    m = sum(xs) / len(xs)
    return (sum((x - m) ** 2 for x in xs) / (len(xs) - 1)) ** 0.5


def format_markdown(summary: dict) -> str:
    cells = summary["cells"]
    scenarios = sorted({c["scenario"] for c in cells})
    policies = sorted({c["policy"] for c in cells})
    lines = ["# OpsGuard eval matrix\n"]
    for s in scenarios:
        lines.append(f"\n## {s}\n")
        lines.append("| Policy | reward (μ±σ) | resolved | spam_recall | steps |")
        lines.append("|---|---:|---:|---:|---:|")
        for p in policies:
            cell = next((c for c in cells if c["policy"] == p and c["scenario"] == s), None)
            if not cell:
                continue
            sr = f"{cell['spam_recall_mean']:.2f}" if cell["spam_recall_mean"] is not None else "n/a"
            lines.append(
                f"| {cell['policy']} | {cell['reward_mean']:+.2f} ± {cell['reward_std']:.2f} | "
                f"{cell['resolved_mean']:.1f} | {sr} | {cell['steps_mean']:.0f} |"
            )
    return "\n".join(lines) + "\n"