#!/usr/bin/env python3
"""Verify downloaded GLM-5.2 ``UD-IQ1_S`` shards against the pinned manifest.

Every shard must match the exact byte size and LFS SHA-256 that DGR-017 locked
in ``target-manifest.json``. Size is checked first because it is free and
catches truncated resumes; the digest is then streamed so a 49 GB file never
loads into memory. A shard that fails is reported with what was measured — the
file is never deleted or "healed" by this script.

Fail-closed: exit 0 only when every requested shard is present and matches.

Usage::

    python scripts/verify_glm_shards.py --model-dir /mnt/models/glm-5.2
    python scripts/verify_glm_shards.py --model-dir DIR --shard 6   # one shard
    python scripts/verify_glm_shards.py --model-dir DIR --json report.json
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
import time
from pathlib import Path

_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_ROOT / "packages/node"))

from meshnet_node.glm_alpha.manifest import (  # noqa: E402
    Shard,
    TargetManifest,
    load_target_manifest,
)

_CHUNK_BYTES = 8 * 1024 * 1024


def resolve_shard_file(model_dir: Path, shard: Shard) -> Path:
    """Accept the repository layout (``UD-IQ1_S/<name>``) or a flat directory."""
    nested = model_dir / shard.path
    if nested.exists():
        return nested
    return model_dir / Path(shard.path).name


def streaming_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while chunk := handle.read(_CHUNK_BYTES):
            digest.update(chunk)
    return digest.hexdigest()


def verify_shard(model_dir: Path, shard: Shard) -> dict:
    result: dict = {
        "shard_index": shard.index,
        "path": shard.path,
        "expected_size_bytes": shard.size_bytes,
        "expected_sha256": shard.sha256,
    }
    file = resolve_shard_file(model_dir, shard)
    result["file"] = str(file)
    if not file.exists():
        result.update(status="missing")
        return result

    size = file.stat().st_size
    result["measured_size_bytes"] = size
    if size != shard.size_bytes:
        result.update(status="size_mismatch")
        return result

    started = time.monotonic()
    measured = streaming_sha256(file)
    result["measured_sha256"] = measured
    result["hash_seconds"] = round(time.monotonic() - started, 1)
    result.update(status="ok" if measured == shard.sha256 else "sha256_mismatch")
    return result


def verify_shards(
    model_dir: Path, manifest: TargetManifest, indices: list[int] | None = None
) -> dict:
    shards = manifest.shards if indices is None else [manifest.shard(i) for i in indices]
    results = [verify_shard(model_dir, shard) for shard in shards]
    return {
        "generated_by": "scripts/verify_glm_shards.py",
        "model_dir": str(model_dir),
        "gguf_repo_id": manifest.gguf_repo_id,
        "gguf_revision": manifest.gguf_revision,
        "quantization": manifest.quantization,
        "manifest_sha256": manifest.digest,
        "shards": results,
        "verdict": "pass" if all(r["status"] == "ok" for r in results) else "fail",
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", required=True, help="directory holding the shards")
    parser.add_argument(
        "--shard",
        type=int,
        action="append",
        dest="shards",
        help="verify only this shard index (repeatable); default: all six",
    )
    parser.add_argument("--json", dest="json_path", help="also write the full report here")
    args = parser.parse_args()

    report = verify_shards(Path(args.model_dir), load_target_manifest(), args.shards)

    if args.json_path:
        Path(args.json_path).write_text(
            json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
        )

    for result in report["shards"]:
        line = f"shard {result['shard_index']}: {result['status']}"
        if result["status"] == "size_mismatch":
            line += (
                f" (expected {result['expected_size_bytes']}, "
                f"measured {result['measured_size_bytes']})"
            )
        elif result["status"] == "sha256_mismatch":
            line += f" (measured {result['measured_sha256']})"
        print(line)
    print(f"verdict: {report['verdict']}")
    return 0 if report["verdict"] == "pass" else 1


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