"""Lightweight experiment tracking: one directory per run under `runs/`.

    runs/<timestamp>-<name>/
        config.yaml      resolved config
        env.json         git commit, hardware, library versions, runtime profile
        metrics.csv      step-level training/validation metrics
        tb/              TensorBoard events (if tensorboard is installed)
        summary.json     final metrics, timings, cost
        checkpoints/     last.pt (resumable), best.pt, final.pt (calibrated)
"""
from __future__ import annotations

import csv
import json
import platform
import subprocess
import time
from pathlib import Path

import yaml


def git_commit() -> str | None:
    try:
        out = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=5,
                             cwd=Path(__file__).resolve().parent)
        dirty = subprocess.run(["git", "status", "--porcelain"], capture_output=True, text=True, timeout=5,
                               cwd=Path(__file__).resolve().parent).stdout.strip()
        return (out.stdout.strip() + ("-dirty" if dirty else "")) if out.returncode == 0 else None
    except Exception:
        return None


def environment_info() -> dict:
    import numpy
    info = {"python": platform.python_version(), "platform": platform.platform(),
            "processor": platform.processor(), "numpy": numpy.__version__, "git_commit": git_commit()}
    try:
        import torch
        info["torch"] = torch.__version__
        info["cuda"] = torch.version.cuda
        if torch.cuda.is_available():
            info["gpu"] = torch.cuda.get_device_name(0)
    except ImportError:
        pass
    return info


class RunLogger:
    def __init__(self, run_dir: str | Path, cfg: dict | None = None, resume: bool = False):
        self.dir = Path(run_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        (self.dir / "checkpoints").mkdir(exist_ok=True)
        if cfg is not None and not (resume and (self.dir / "config.yaml").exists()):
            (self.dir / "config.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False))
        self._csv_path = self.dir / "metrics.csv"
        self._fields: list[str] | None = None
        if resume and self._csv_path.exists():
            with open(self._csv_path) as f:
                self._fields = next(csv.reader(f), None)
        self.tb = None
        try:
            from torch.utils.tensorboard import SummaryWriter
            self.tb = SummaryWriter(str(self.dir / "tb"))
        except Exception:
            pass

    def write_json(self, name: str, obj) -> None:
        (self.dir / name).write_text(json.dumps(obj, indent=2, default=str))

    def log(self, step: int, **metrics) -> None:
        row = {"step": step, "time": round(time.time(), 1), **metrics}
        if self._fields is None:
            self._fields = list(row)
            with open(self._csv_path, "w", newline="") as f:
                csv.writer(f).writerow(self._fields)
        for k in row:
            if k not in self._fields:  # new metric: rewrite header
                self._fields.append(k)
                rows = list(csv.DictReader(open(self._csv_path)))
                with open(self._csv_path, "w", newline="") as f:
                    w = csv.DictWriter(f, self._fields)
                    w.writeheader()
                    w.writerows(rows)
        with open(self._csv_path, "a", newline="") as f:
            csv.DictWriter(f, self._fields).writerow(row)
        if self.tb is not None:
            for k, v in metrics.items():
                if isinstance(v, (int, float)):
                    self.tb.add_scalar(k, v, step)

    def close(self) -> None:
        if self.tb is not None:
            self.tb.close()


def new_run_dir(name: str, root: str | Path = "runs") -> Path:
    return Path(root) / f"{time.strftime('%Y%m%d-%H%M%S')}-{name}"