#!/usr/bin/env python
"""Sharded checkpoint tools: inspect, verify, export to HF safetensors, reshard.

    python scripts/ckpt_convert.py inspect RUN/checkpoints/micro_N
    python scripts/ckpt_convert.py verify  RUN/checkpoints/micro_N
    python scripts/ckpt_convert.py export  RUN/checkpoints/micro_N OUT_DIR
    python scripts/ckpt_convert.py reshard RUN/checkpoints/micro_N NEW_DIR --world-size 2

CPU only, no GPU, no process group, no heavy-job lock: `export` is a streaming
byte copy (peak RSS ~ a few MiB plus the Python runtime, independent of model
size) and `reshard` re-cuts the weights the same way. The ONE part that holds
real memory is the optimizer re-cut inside `reshard` (one NEW rank's optimizer
state, e.g. 12 B/param/new_ws for fp32-master AdamW); it prints that estimate
and refuses to start when the box cannot afford it (CLAUDE.md: 45 GB
MemAvailable floor for anything heavy, no swap).

WHY THIS EXISTS. `train.checkpoint_format = "sharded"` writes per-rank files
that eval, the rollout weight sync and external tools cannot read. `export`
turns one into exactly the directory the gathered writer would have produced
(byte-identical; tested by file hash) without ever holding the model in RAM --
the gathered path's rank-0 gather measured +28 GB RSS on OLMoE (docs/
MEASUREMENTS.md, Track M). `reshard` is the ONLY way to resume at a different
world size: resume refuses it, because sharded weights are rank-local row
slices and silently loading them under another partitioning trains plausibly
on scrambled weights.

Exit status: 0 ok, 1 refused/corrupt (message on stderr), 2 usage.
"""

from __future__ import annotations

import argparse
import json
import os
import sys

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # never touch a GPU from a converter

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from align.train import sharded_ckpt as sc  # noqa: E402


def _mem_available_gib() -> float:
    with open("/proc/meminfo") as f:
        for line in f:
            if line.startswith("MemAvailable:"):
                return int(line.split()[1]) / 2**20
    return 0.0


def _optimizer_bytes(src: str, ws: int) -> int:
    return sum(os.path.getsize(os.path.join(src, f"optimizer_rank{r}.pt")) for r in range(ws) if os.path.exists(os.path.join(src, f"optimizer_rank{r}.pt")))


def cmd_inspect(a) -> int:
    m = sc.read_manifest(a.src)
    total = sum(f["bytes"] for f in m["files"].values())
    print(
        json.dumps(
            {
                "world_size": m["world_size"],
                "tensors": len(m["tensors"]),
                "weight_bytes": total,
                "export_family": (m.get("export") or {}).get("family"),
                "resharded_from": m.get("resharded_from"),
                "problem": sc.sharded_problem(a.src, lambda p: None),
            },
            indent=1,
        )
    )
    return 0


def cmd_verify(a) -> int:
    m = sc.read_manifest(a.src)
    bad = []
    for r in range(m["world_size"]):
        p = sc.verify_rank_file(a.src, r)
        print(f"rank {r}: {'OK' if p is None else p}")
        if p:
            bad.append(p)
    return 1 if bad else 0


def cmd_export(a) -> int:
    from align.utils.guard import assert_paths_off_root

    assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))])
    res = sc.export_hf(a.src, a.dst, verify=not a.no_verify)
    print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()}))
    return 0


def cmd_reshard(a) -> int:
    from align.utils.guard import assert_paths_off_root

    assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))])
    m = sc.read_manifest(a.src)
    if not a.drop_optimizer:
        est = _optimizer_bytes(a.src, m["world_size"]) / max(1, a.world_size) / 2**30
        avail = _mem_available_gib()
        print(f"optimizer re-cut peak ~ {est:.1f} GiB; MemAvailable {avail:.0f} GiB", file=sys.stderr)
        if not a.allow_large and avail < est + 45:
            print(
                f"refusing: need MemAvailable >= {est + 45:.0f} GiB (estimate + the 45 GiB box floor) and have "
                f"{avail:.0f}. Wait for the box, or pass --allow-large if you accept the risk, or --drop-optimizer "
                "(fresh moments; NOT lossless).",
                file=sys.stderr,
            )
            return 1
    res = sc.reshard(a.src, a.dst, a.world_size, drop_optimizer=a.drop_optimizer)
    print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()}))
    return 0


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name in ("inspect", "verify"):
        p = sub.add_parser(name)
        p.add_argument("src")
    p = sub.add_parser("export")
    p.add_argument("src")
    p.add_argument("dst")
    p.add_argument("--no-verify", action="store_true", help="skip the per-file sha256 pass (not recommended)")
    p = sub.add_parser("reshard")
    p.add_argument("src")
    p.add_argument("dst")
    p.add_argument("--world-size", type=int, required=True)
    p.add_argument("--drop-optimizer", action="store_true", help="omit optimizer files: resume starts FRESH moments (NOT lossless)")
    p.add_argument("--allow-large", action="store_true")
    a = ap.parse_args(argv)
    try:
        return {"inspect": cmd_inspect, "verify": cmd_verify, "export": cmd_export, "reshard": cmd_reshard}[a.cmd](a)
    except (sc.ShardedCheckpointError, FileExistsError, FileNotFoundError, KeyError) as e:
        print(f"{type(e).__name__}: {e}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())