#!/usr/bin/env python3
"""Small helpers for prepare_model_dir.sh (standard library only).

Run as:  python3 -I prepare_helpers.py <command> ...
Commands: vae-config, rewrite-source, st-check, shards-check, verify
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import struct
import sys
import tempfile
from pathlib import Path

ENCODER_SOURCE = "hf://FastVideo/FastH3-Omni-NVFP4-QAD/text_encoder"
MANIFEST = "mlx_h3_nvfp4_encoder.json"


def _atomic_write(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, tmp = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    with os.fdopen(fd, "w") as f:
        f.write(text)
    os.replace(tmp, path)


def safetensors_header(path: Path) -> dict:
    """Parse the header and check the file size matches the declared tensor data."""
    size = path.stat().st_size
    with path.open("rb") as f:
        raw = f.read(8)
        if len(raw) != 8:
            raise ValueError("file shorter than 8 bytes")
        (n,) = struct.unpack("<Q", raw)
        if n <= 0 or 8 + n > size:
            raise ValueError(f"bad header length {n} for file size {size}")
        header = json.loads(f.read(n))
    end = max((v["data_offsets"][1] for k, v in header.items() if k != "__metadata__"), default=0)
    if 8 + n + end != size:
        raise ValueError(f"truncated or padded: expected {8 + n + end} bytes, found {size}")
    return header


def cmd_vae_config(src: str, dst: str) -> None:
    meta = safetensors_header(Path(src)).get("__metadata__") or {}
    config = json.loads(meta["architecture"])["teacher_config"]
    _atomic_write(Path(dst), json.dumps(config, indent=2))  # no trailing newline, matches our build


def cmd_rewrite_source(manifest: str, source: str) -> None:
    path = Path(manifest)
    data = json.loads(path.read_text())
    if data.get("source_dir") != source:
        data["source_dir"] = source
        _atomic_write(path, json.dumps(data, indent=2) + "\n")


def cmd_st_check(files: list[str]) -> int:
    bad = 0
    for name in files:
        try:
            safetensors_header(Path(name))
        except Exception as exc:  # noqa: BLE001 - report every bad file
            print(f"BAD safetensors {name}: {exc}", file=sys.stderr)
            bad += 1
    return 1 if bad else 0


def cmd_shards_check(directory: str, prefix: str) -> int:
    """All `<prefix>-0000k-of-0000N.safetensors` shards present and well-formed."""
    d = Path(directory)
    pat = re.compile(rf"^{re.escape(prefix)}-(\d+)-of-(\d+)\.safetensors$")
    found = {}
    total = 0
    for p in d.iterdir() if d.is_dir() else []:
        m = pat.match(p.name)
        if m:
            found[int(m.group(1))] = p
            total = int(m.group(2))
    if not found or sorted(found) != list(range(1, total + 1)):
        print(f"shards incomplete in {d}: have {sorted(found)} of {total}", file=sys.stderr)
        return 1
    return cmd_st_check([str(p) for p in found.values()])


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while chunk := f.read(16 << 20):
            h.update(chunk)
    return h.hexdigest()


def _is_converter_output(rel: str) -> bool:
    return rel.startswith("dit/") or rel in (f"root/text_encoder/{MANIFEST}", "root/text_encoder/model.safetensors")


def cmd_verify(model_dir: str, expected_file: str, formats: list[str]) -> int:
    root = Path(model_dir)
    fails, warns, ok = [], [], 0
    for line in Path(expected_file).read_text().splitlines():
        if not line.strip():
            continue
        digest, rel = line.split(None, 1)
        rel = rel.strip()
        if rel.startswith("dit/") and rel.split("/")[1] not in formats:
            continue  # format not requested
        path = root / rel
        soft = _is_converter_output(rel)
        if not path.is_file():
            fails.append(f"missing: {rel}")
            print(f"  MISSING {rel}", flush=True)
            continue
        got = sha256_file(path)
        if got == digest:
            ok += 1
            print(f"  ok      {rel}", flush=True)
        elif soft:
            warns.append(f"converter output differs (not bit-identical): {rel}")
            print(f"  WARN    {rel} (converter output differs from reference)", flush=True)
        else:
            fails.append(f"hash mismatch: {rel}")
            print(f"  FAIL    {rel}\n          expected {digest}\n          got      {got}", flush=True)
    # structural checks for outputs that have no reference hash (int8 / bf16 and everything safetensors)
    for fmt in formats:
        d = root / "dit" / fmt
        for name in ("mlx_h3_dit.json", "mlx_h3_dit.safetensors"):
            if not (d / name).is_file():
                fails.append(f"missing: dit/{fmt}/{name}")
        if (d / "mlx_h3_dit.safetensors").is_file() and cmd_st_check([str(d / "mlx_h3_dit.safetensors")]):
            fails.append(f"dit/{fmt}/mlx_h3_dit.safetensors is not a complete safetensors file")
    enc = root / "root/text_encoder/model.safetensors"
    if enc.is_file() and cmd_st_check([str(enc)]):
        fails.append("root/text_encoder/model.safetensors is not a complete safetensors file")
    for rel in ("root/transformer_ref/config.json", "root/vae/config.json", "root/fastvideo_inference.json"):
        try:
            json.loads((root / rel).read_text())
        except Exception as exc:  # noqa: BLE001
            fails.append(f"unreadable JSON {rel}: {exc}")
    print(f"\nverify: {ok} exact matches, {len(warns)} warnings, {len(fails)} failures")
    for w in warns:
        print(f"  warning: {w}")
    for f in fails:
        print(f"  FAILURE: {f}")
    return 1 if fails else 0


def main(argv: list[str]) -> int:
    if not argv:
        print(__doc__)
        return 2
    cmd, args = argv[0], argv[1:]
    if cmd == "vae-config":
        cmd_vae_config(*args)
    elif cmd == "rewrite-source":
        cmd_rewrite_source(*args)
    elif cmd == "st-check":
        return cmd_st_check(args)
    elif cmd == "shards-check":
        return cmd_shards_check(*args)
    elif cmd == "verify":
        return cmd_verify(args[0], args[1], args[2:])
    else:
        print(f"unknown command {cmd}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))