"""Shared training utilities: run ids + manifests, seeding, device selection, token data, checkpoints, logging."""
from __future__ import annotations

import hashlib, json, os, platform, random, subprocess, sys, time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator, Optional

import numpy as np
import torch
import yaml


# ----------------------------------------------------------------------------- environment
def git_sha() -> str:
    try:
        return subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=5).stdout.strip() or "unknown"
    except Exception:  # noqa: BLE001
        return "unknown"


def pick_device(pref: str = "auto") -> torch.device:
    if pref != "auto":
        return torch.device(pref)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def seed_all(seed: int) -> None:
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def sync(device: torch.device) -> None:
    if device.type == "mps":
        torch.mps.synchronize()
    elif device.type == "cuda":
        torch.cuda.synchronize()


def peak_memory_bytes(device: torch.device) -> Optional[int]:
    try:
        if device.type == "mps":
            return int(torch.mps.driver_allocated_memory())
        if device.type == "cuda":
            return int(torch.cuda.max_memory_allocated())
    except Exception:  # noqa: BLE001
        return None
    try:
        import resource
        return int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
    except Exception:  # noqa: BLE001
        return None


def sha256_file(p: str | Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def load_yaml(path: str | Path) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


# ----------------------------------------------------------------------------- run bookkeeping
class Run:
    """Immutable run id + manifest + jsonl metrics log under experiments/runs/<run_id>/."""

    def __init__(self, name: str, config: dict, root: str = "experiments/runs", seed: int = 0, resume_id: str | None = None):
        ts = time.strftime("%Y%m%d%H%M%S")
        self.run_id = resume_id or f"{name}_seed{seed:02d}_{ts}"
        self.dir = Path(root) / self.run_id
        self.dir.mkdir(parents=True, exist_ok=True)
        self.config = config
        self.t0 = time.time()
        self.manifest = {
            "run_id": self.run_id, "name": name, "seed": seed, "code_commit": git_sha(),
            "config_sha256": hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest(),
            "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "host": platform.node(), "python": sys.version.split()[0],
            "torch": torch.__version__, "command": " ".join(sys.argv), "state": "running",
        }
        (self.dir / "config.yaml").write_text(yaml.safe_dump(config, sort_keys=False))
        self._log = open(self.dir / "metrics.jsonl", "a")
        self.write_manifest()

    def write_manifest(self, **extra: Any) -> None:
        self.manifest.update(extra)
        self.manifest["wall_seconds"] = round(time.time() - self.t0, 2)
        (self.dir / "manifest.json").write_text(json.dumps(self.manifest, indent=2, default=str))

    def log(self, **kv: Any) -> None:
        kv["t"] = round(time.time() - self.t0, 2)
        self._log.write(json.dumps(kv, default=float) + "\n"); self._log.flush()

    def finish(self, state: str = "completed", **extra: Any) -> None:
        self.write_manifest(state=state, finished_at=time.strftime("%Y-%m-%dT%H:%M:%S%z"), **extra)
        self._log.close()


# ----------------------------------------------------------------------------- token data
class TokenBin:
    """Memory-mapped uint16 token stream (documents separated by EOT). Random contiguous windows of length T+1."""

    def __init__(self, path: str | Path, seq_len: int):
        self.path = Path(path)
        self.data = np.memmap(self.path, dtype=np.uint16, mode="r")
        self.seq_len = seq_len

    def __len__(self) -> int:
        return len(self.data)

    def batches(self, batch_size: int, seed: int, device: torch.device) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:
        rng = np.random.default_rng(seed)
        n = len(self.data) - self.seq_len - 1
        while True:
            idx = rng.integers(0, n, size=batch_size)
            x = np.stack([self.data[i:i + self.seq_len + 1].astype(np.int64) for i in idx])
            xt = torch.from_numpy(x)
            yield xt[:, :-1].to(device), xt[:, 1:].to(device)

    def sequential_batches(self, batch_size: int, device: torch.device, max_batches: int | None = None):
        """Deterministic non-overlapping windows for evaluation."""
        n = (len(self.data) - 1) // self.seq_len
        count = 0
        for start in range(0, n, batch_size):
            idx = [i * self.seq_len for i in range(start, min(start + batch_size, n))]
            x = np.stack([self.data[i:i + self.seq_len + 1].astype(np.int64) for i in idx])
            xt = torch.from_numpy(x)
            yield xt[:, :-1].to(device), xt[:, 1:].to(device)
            count += 1
            if max_batches is not None and count >= max_batches:
                break


# ----------------------------------------------------------------------------- optimisation helpers
def cosine_lr(step: int, warmup: int, total: int, lr: float, min_ratio: float = 0.1) -> float:
    if step < warmup:
        return lr * (step + 1) / warmup
    p = min(1.0, (step - warmup) / max(1, total - warmup))
    import math
    return lr * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * p)))


def make_adamw(params, lr: float, weight_decay: float, betas=(0.9, 0.95)):
    decay, no_decay = [], []
    for p in params:
        (decay if p.ndim >= 2 else no_decay).append(p)
    return torch.optim.AdamW([{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}],
                             lr=lr, betas=betas)


def save_checkpoint(path: str | Path, model: torch.nn.Module, cfg_dict: dict, extra: dict | None = None) -> str:
    from safetensors.torch import save_file
    path = Path(path); path.mkdir(parents=True, exist_ok=True)
    sd = {k: v.detach().to("cpu").contiguous() for k, v in model.state_dict().items()}
    save_file(sd, str(path / "model.safetensors"))
    meta = {"config": cfg_dict, "saved_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), **(extra or {})}
    meta["weights_sha256"] = sha256_file(path / "model.safetensors")
    (path / "checkpoint.json").write_text(json.dumps(meta, indent=2, default=str))
    return meta["weights_sha256"]


def load_checkpoint(path: str | Path, model: torch.nn.Module, strict: bool = True) -> dict:
    from safetensors.torch import load_file
    path = Path(path)
    sd = load_file(str(path / "model.safetensors"))
    own = model.state_dict()
    mismatched = [k for k, v in sd.items() if k in own and own[k].shape != v.shape]
    if mismatched and strict:
        raise RuntimeError(f"shape mismatch for {mismatched}")
    for k in mismatched:          # e.g. a controller re-shaped by a feature option: left at its fresh init
        sd.pop(k)
    missing, unexpected = model.load_state_dict(sd, strict=strict)
    meta = json.loads((path / "checkpoint.json").read_text())
    meta["_missing"] = list(missing) + mismatched; meta["_unexpected"] = list(unexpected); meta["_shape_mismatch"] = mismatched
    return meta