#!/usr/bin/env python3
"""DGR-018 whole-model preflight: storage, memory, and the ordered download plan.

Answers, before a single artifact byte moves, the three questions DGR-018's
finish contract asks:

1. Is there one filesystem, outside every forbidden prefix (``/home``), with at
   least 250 GB free for the 216.715 GB ``UD-IQ1_S`` artifact plus headroom?
2. Does this host have the 224 GiB runtime-accessible memory the whole-model
   oracle load needs (the experimental hard-fit floor from DGR-017)?
3. In what order do the six shards download, and against which exact revision
   URL, size, and LFS SHA-256 is each one verified?

Everything is resolved from the pinned target manifest
(``meshnet_node.glm_alpha``) — this script never contacts the network and never
invents a size or digest. It fails closed: a missing requirement is a non-zero
exit and an explicit reason, never a downgrade to a smaller target.

Usage::

    python scripts/glm_whole_model_preflight.py                 # scan all mounts
    python scripts/glm_whole_model_preflight.py --dest DIR      # judge one destination
    python scripts/glm_whole_model_preflight.py --storage-only  # download-host mode
    python scripts/glm_whole_model_preflight.py --json out.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import dataclass
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
    GIB,
    TargetManifest,
    load_target_manifest,
)

GB = 1000**3

# 250 GB: the 216.715 GB artifact plus resume/temp headroom. DGR-018's issue
# text names this number; it is a storage requirement, not a tunable.
REQUIRED_FREE_BYTES = 250 * GB

# DGR-017's experimental hard-fit floor for whole-model runtime-accessible
# memory. 224 GiB fits weights + Q8_0 KV at 16k context with nothing to spare.
REQUIRED_MEMORY_GIB = 224.0

# Download order, by shard index. Rationale:
#   1 first  — 9.4 MB split-metadata shard: proves revision URLs, auth-free
#              access, and tooling end-to-end for the cost of a rounding error.
#   6 second — 19.2 GB, the smallest weight shard: proves resume + hash
#              verification at real scale before any ~49 GB transfer starts.
#   2..5     — the four ~49 GB shards, sequentially, verify-after-each, so at
#              most one unverified partial ever exists on disk.
DOWNLOAD_ORDER = (1, 6, 2, 3, 4, 5)

# Filesystems that can hold a 49 GB file and survive a resume. FAT variants are
# excluded by omission (4 GiB file limit); network/pseudo filesystems likewise.
_USABLE_FSTYPES = {"ext4", "ext3", "xfs", "btrfs", "ntfs", "ntfs3", "fuseblk", "exfat", "f2fs"}


@dataclass(frozen=True)
class MountCandidate:
    mountpoint: str
    fstype: str
    total_bytes: int
    free_bytes: int
    forbidden: bool

    @property
    def eligible(self) -> bool:
        return not self.forbidden and self.free_bytes >= REQUIRED_FREE_BYTES

    def to_dict(self) -> dict:
        return {
            "mountpoint": self.mountpoint,
            "fstype": self.fstype,
            "total_gb": round(self.total_bytes / GB, 1),
            "free_gb": round(self.free_bytes / GB, 1),
            "free_bytes": self.free_bytes,
            "forbidden": self.forbidden,
            "eligible": self.eligible,
        }


def forbidden_prefixes(manifest: TargetManifest) -> tuple[str, ...]:
    """The storage policy pinned in the manifest itself — not a script opinion."""
    storage = manifest.raw.get("storage", {})
    return tuple(storage.get("forbidden_path_prefixes", ("/home",)))


def is_forbidden(path: str, prefixes: tuple[str, ...]) -> bool:
    resolved = os.path.realpath(path)
    return any(resolved == p or resolved.startswith(p.rstrip("/") + "/") for p in prefixes)


def _statvfs_free(path: str) -> tuple[int, int]:
    st = os.statvfs(path)
    return st.f_frsize * st.f_blocks, st.f_frsize * st.f_bavail


def scan_mounts(
    prefixes: tuple[str, ...], mounts_file: str = "/proc/mounts"
) -> list[MountCandidate]:
    """Every real, writable-class filesystem on the host, deduplicated by device."""
    candidates: dict[str, MountCandidate] = {}
    with open(mounts_file, encoding="utf-8") as handle:
        for line in handle:
            fields = line.split()
            if len(fields) < 3:
                continue
            device, mountpoint, fstype = fields[0], fields[1], fields[2]
            mountpoint = mountpoint.replace("\\040", " ")
            if fstype not in _USABLE_FSTYPES or device in candidates:
                continue
            try:
                total, free = _statvfs_free(mountpoint)
            except OSError:
                continue
            candidates[device] = MountCandidate(
                mountpoint=mountpoint,
                fstype=fstype,
                total_bytes=total,
                free_bytes=free,
                forbidden=is_forbidden(mountpoint, prefixes),
            )
    return sorted(candidates.values(), key=lambda c: c.free_bytes, reverse=True)


def judge_destination(dest: str, prefixes: tuple[str, ...]) -> MountCandidate:
    total, free = _statvfs_free(dest)
    return MountCandidate(
        mountpoint=dest,
        fstype="(as-given)",
        total_bytes=total,
        free_bytes=free,
        forbidden=is_forbidden(dest, prefixes),
    )


def memory_total_gib(meminfo_file: str = "/proc/meminfo") -> float:
    with open(meminfo_file, encoding="utf-8") as handle:
        for line in handle:
            if line.startswith("MemTotal:"):
                return int(line.split()[1]) * 1024 / GIB
    raise RuntimeError(f"MemTotal not found in {meminfo_file}")


def download_plan(manifest: TargetManifest, dest_dir: str | None) -> list[dict]:
    """The six shards in download order, each bound to its pinned identity."""
    steps = []
    for position, index in enumerate(DOWNLOAD_ORDER, start=1):
        shard = manifest.shard(index)
        target = f"{dest_dir or '$GLM_DEST'}/{shard.path}"
        steps.append(
            {
                "step": position,
                "shard_index": shard.index,
                "path": shard.path,
                "size_bytes": shard.size_bytes,
                "size_gb": round(shard.size_bytes / GB, 3),
                "sha256": shard.sha256,
                "url": shard.url,
                "download_command": f'curl -L -C - --fail -o "{target}" "{shard.url}"',
                "verify_command": (
                    f"python scripts/verify_glm_shards.py "
                    f'--model-dir "{dest_dir or "$GLM_DEST"}" --shard {shard.index}'
                ),
            }
        )
    return steps


def build_report(*, dest: str | None, storage_only: bool) -> dict:
    manifest = load_target_manifest()
    prefixes = forbidden_prefixes(manifest)

    if dest is not None:
        mounts = [judge_destination(dest, prefixes)]
    else:
        mounts = scan_mounts(prefixes)
    eligible = [m for m in mounts if m.eligible]
    chosen = eligible[0] if eligible else None

    mem_gib = memory_total_gib()
    storage_pass = chosen is not None
    memory_pass = mem_gib >= REQUIRED_MEMORY_GIB

    checks = [
        {
            "check": "storage",
            "requirement": f">= {REQUIRED_FREE_BYTES / GB:.0f} GB free on one filesystem "
            f"outside {list(prefixes)}",
            "observed": (
                f"{chosen.free_bytes / GB:.1f} GB free at {chosen.mountpoint}"
                if chosen
                else "no eligible filesystem"
            ),
            "passes": storage_pass,
        },
        {
            "check": "memory",
            "requirement": f">= {REQUIRED_MEMORY_GIB:.0f} GiB runtime-accessible memory "
            "(DGR-017 experimental hard-fit floor)",
            "observed": f"{mem_gib:.1f} GiB MemTotal",
            "passes": memory_pass,
            "waived": storage_only,
        },
    ]
    overall = storage_pass and (memory_pass or storage_only)

    return {
        "generated_by": "scripts/glm_whole_model_preflight.py",
        "target": {
            "gguf_repo_id": manifest.gguf_repo_id,
            "gguf_revision": manifest.gguf_revision,
            "quantization": manifest.quantization,
            "shard_count": len(manifest.shards),
            "total_bytes": manifest.total_bytes,
            "total_gb": round(manifest.total_gb, 3),
            "manifest_sha256": manifest.digest,
        },
        "forbidden_path_prefixes": list(prefixes),
        "mounts": [m.to_dict() for m in mounts],
        "chosen_destination": chosen.to_dict() if chosen else None,
        "checks": checks,
        "download_authorized": overall,
        "storage_only": storage_only,
        "download_plan": download_plan(manifest, chosen.mountpoint if chosen else None),
        "verdict": "pass" if overall else "fail",
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dest", help="judge this destination directory instead of scanning")
    parser.add_argument(
        "--storage-only",
        action="store_true",
        help="download-host mode: require storage, report memory without gating on it",
    )
    parser.add_argument("--json", dest="json_path", help="also write the full report here")
    args = parser.parse_args()

    report = build_report(dest=args.dest, storage_only=args.storage_only)

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

    target = report["target"]
    print(f"target: {target['quantization']} {target['total_gb']} GB, "
          f"{target['shard_count']} shards @ {target['gguf_revision'][:12]}")
    for check in report["checks"]:
        status = "PASS" if check["passes"] else ("WAIVED" if check.get("waived") else "FAIL")
        print(f"[{status}] {check['check']}: need {check['requirement']}; "
              f"observed {check['observed']}")
    if report["chosen_destination"]:
        print(f"destination: {report['chosen_destination']['mountpoint']}")
    else:
        print("destination: NONE — no filesystem outside "
              f"{report['forbidden_path_prefixes']} has "
              f"{REQUIRED_FREE_BYTES / GB:.0f} GB free")
        for mount in report["mounts"]:
            why = "forbidden prefix" if mount["forbidden"] else f"{mount['free_gb']} GB free"
            print(f"  - {mount['mountpoint']} ({mount['fstype']}): {why}")
    print(f"verdict: {report['verdict']}")
    return 0 if report["download_authorized"] else 1


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