"""Checkpoint save/load for resumable OmniAL training."""

import json
import os
from pathlib import Path
from typing import Any, Dict, Optional

import torch
from torch.amp import GradScaler


def ensure_run_dir(run_dir: str) -> Path:
    p = Path(run_dir)
    (p / "checkpoints").mkdir(parents=True, exist_ok=True)
    (p / "metrics").mkdir(parents=True, exist_ok=True)
    return p


def save_checkpoint(
    path: str,
    epoch: int,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scaler: Optional[GradScaler],
    history: Dict[str, Any],
    best_metric: float,
    cfg: dict,
    tag: str = "",
    scheduler=None,
) -> str:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    payload = {
        "epoch": epoch,
        "model_state": model.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "scaler_state": scaler.state_dict() if scaler is not None else None,
        "scheduler_state": scheduler.state_dict() if scheduler is not None else None,
        "history": history,
        "best_metric": best_metric,
        "cfg_train": cfg.get("train", {}),
        "tag": tag,
    }
    torch.save(payload, path)
    return path


def load_checkpoint(
    path: str,
    model: torch.nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scaler: Optional[GradScaler] = None,
    device: str = "cpu",
) -> Dict[str, Any]:
    # Full training checkpoints store history/optimizer (numpy scalars etc.).
    # PyTorch 2.6+ defaults weights_only=True, which rejects those — use False
    # only for checkpoints you trust (your own training runs).
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model_state"])
    if optimizer is not None and ckpt.get("optimizer_state"):
        optimizer.load_state_dict(ckpt["optimizer_state"])
    if scaler is not None and ckpt.get("scaler_state"):
        try:
            scaler.load_state_dict(ckpt["scaler_state"])
        except Exception as e:
            print(f"[WARN] could not restore GradScaler state: {e}")
    return ckpt


def restore_scheduler(scheduler, ckpt: Dict[str, Any]) -> None:
    if scheduler is None or not ckpt.get("scheduler_state"):
        return
    try:
        scheduler.load_state_dict(ckpt["scheduler_state"])
    except Exception as e:
        print(f"[WARN] could not restore scheduler state: {e}")


def save_metrics_json(run_dir: str, history: Dict[str, Any], filename: str = "metrics_history.json") -> str:
    path = Path(run_dir) / "metrics" / filename
    with open(path, "w", encoding="utf-8") as f:
        json.dump(history, f, indent=2, ensure_ascii=False)
    return str(path)


def append_epoch_metrics(run_dir: str, epoch_record: Dict[str, Any]) -> str:
    metrics_dir = Path(run_dir) / "metrics"
    metrics_dir.mkdir(parents=True, exist_ok=True)
    path = metrics_dir / "epoch_metrics.jsonl"
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(epoch_record, ensure_ascii=False) + "\n")
    return str(path)


def find_latest_checkpoint(run_dir: str) -> Optional[str]:
    ckpt_dir = Path(run_dir) / "checkpoints"
    if not ckpt_dir.is_dir():
        return None
    last = ckpt_dir / "last.pt"
    if last.is_file():
        return str(last)
    numbered = sorted(ckpt_dir.glob("epoch_*.pt"), key=lambda p: p.stat().st_mtime)
    return str(numbered[-1]) if numbered else None