from __future__ import annotations

import csv
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


@dataclass
class StepLogRow:
    episode: int
    step: int
    reward_total: float
    done: bool
    my_crowns: int
    opp_crowns: int
    opp_tilt: float
    my_elixir: float
    tower_damage: float
    crown_differential: float
    tilt_efficiency: float
    elixir_discipline: float
    step_reward: float
    invalid_action: int


def ensure_outputs_dirs(root: Path) -> tuple[Path, Path]:
    logs = root / "outputs" / "logs"
    plots = root / "outputs" / "plots"
    logs.mkdir(parents=True, exist_ok=True)
    plots.mkdir(parents=True, exist_ok=True)
    return logs, plots


def write_step_logs_csv(path: Path, rows: Iterable[StepLogRow]) -> None:
    rows = list(rows)
    with path.open("w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(
            [
                "episode",
                "step",
                "reward_total",
                "done",
                "my_crowns",
                "opp_crowns",
                "opp_tilt",
                "my_elixir",
                "tower_damage",
                "crown_differential",
                "tilt_efficiency",
                "elixir_discipline",
                "step_reward",
                "invalid_action",
            ]
        )
        for r in rows:
            w.writerow(
                [
                    r.episode,
                    r.step,
                    r.reward_total,
                    int(r.done),
                    r.my_crowns,
                    r.opp_crowns,
                    f"{r.opp_tilt:.4f}",
                    f"{r.my_elixir:.3f}",
                    f"{r.tower_damage:.6f}",
                    f"{r.crown_differential:.6f}",
                    f"{r.tilt_efficiency:.6f}",
                    f"{r.elixir_discipline:.6f}",
                    f"{r.step_reward:.6f}",
                    r.invalid_action,
                ]
            )