#!/usr/bin/env python3
"""Write down exactly what your model estate is, so it can be rebuilt.

Purpose: produce one JSON file listing every model file with its size and hash, every
    engine with its version, every container image with its digest, and every
    configuration file with its hash, so that a new machine can be brought to the same
    state and an old one can be proved to have changed.
Platform: all (spark, strix, mac, nvidia). Pure Python, no dependencies.
Minimum memory: 8 GB, which is the estate being described; this script needs almost none.
Assumes: a model library laid out as Part 4 built it, the gateway directory from Part 9,
    and whichever engines are installed on PATH. Hashing a large model library reads every
    byte of it, so hashes are taken from the .sha256 sidecar files Part 4's downloader
    wrote where those exist, and computed only for files with --hash. Nothing is modified.

Usage: python3 estate-manifest.py --models ~/models --gateway ~/gateway
       python3 estate-manifest.py --models ~/models --gateway ~/gateway --output estate.json
       python3 estate-manifest.py --models ~/models --hash --output estate.json
"""
import argparse
import hashlib
import json
import platform
import re
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

WEIGHT_SUFFIXES = {".gguf", ".safetensors", ".bin", ".pt", ".pth", ".npz"}
CONFIG_SUFFIXES = {".yaml", ".yml", ".json", ".conf", ".toml", ".md", ".txt"}

# Engines the course installs, and the argument that makes each print its version. A tool
# that is not on PATH is recorded as absent rather than omitted, because "not installed" is
# part of the description of a machine.
ENGINE_VERSION_COMMANDS = [
    ("llama-server", ["llama-server", "--version"]),
    ("llama-swap", ["llama-swap", "--version"]),
    ("vllm", ["vllm", "--version"]),
    ("litellm", ["litellm", "--version"]),
    ("mlx_lm.server", ["mlx_lm.server", "--help"]),
    ("ollama", ["ollama", "--version"]),
    ("docker", ["docker", "--version"]),
]


def run(cmd, timeout=20):
    if shutil.which(cmd[0]) is None:
        return None
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, check=False)
    except (OSError, subprocess.TimeoutExpired):
        return None
    text = (out.stdout or "") + (out.stderr or "")
    return text.strip().splitlines()[0] if text.strip() else ""


def sha256_of(path, chunk=1024 * 1024):
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(chunk), b""):
            digest.update(block)
    return digest.hexdigest()


def sidecar_hash(path):
    """Part 4's downloader writes <file>.sha256 beside each model file it verified."""
    sidecar = path.with_name(path.name + ".sha256")
    if not sidecar.is_file():
        return None
    text = sidecar.read_text(encoding="utf-8", errors="replace").strip()
    match = re.search(r"\b[0-9a-f]{64}\b", text)
    return match.group(0) if match else None


def describe_models(models_dir, compute_hashes, verbose):
    rows = []
    if models_dir is None:
        return rows
    root = Path(models_dir).expanduser()
    if not root.is_dir():
        print(f"estate-manifest: no such model directory: {root}", file=sys.stderr)
        return rows
    for path in sorted(root.rglob("*")):
        if not path.is_file() or path.suffix.lower() not in WEIGHT_SUFFIXES:
            continue
        stat = path.stat()
        digest = sidecar_hash(path)
        source = "sidecar" if digest else None
        if digest is None and compute_hashes:
            if verbose:
                print(f"    hashing {path.name} ...", file=sys.stderr)
            digest = sha256_of(path)
            source = "computed"
        rows.append({
            "path": str(path.relative_to(root)),
            "bytes": stat.st_size,
            "modified": datetime.fromtimestamp(stat.st_mtime, timezone.utc).isoformat(),
            "sha256": digest,
            "sha256_source": source or "not recorded",
        })
    return rows


def describe_engines():
    rows = []
    for name, cmd in ENGINE_VERSION_COMMANDS:
        version = run(cmd)
        rows.append({
            "tool": name,
            "present": version is not None,
            "version_line": version if version else "not installed",
        })
    return rows


def image_references(gateway_dir):
    """Every image name mentioned in the gateway's .env or compose files."""
    names = set()
    if gateway_dir is None:
        return names
    root = Path(gateway_dir).expanduser()
    for candidate in list(root.glob("*.yaml")) + list(root.glob("*.yml")) + [root / ".env"]:
        if not candidate.is_file():
            continue
        for line in candidate.read_text(encoding="utf-8", errors="replace").splitlines():
            match = re.search(r"(?:image:\s*|_IMAGE=)[\"']?([A-Za-z0-9._/-]+(?::[A-Za-z0-9._-]+)?)",
                              line)
            if not match:
                continue
            reference = match.group(1).strip("\"'")
            # A bare word is a Compose variable that was not expanded, not an image.
            if "/" in reference or ":" in reference:
                names.add(reference)
    return names


def describe_images(gateway_dir):
    rows = []
    if shutil.which("docker") is None:
        return rows
    for name in sorted(image_references(gateway_dir)):
        digest = run(["docker", "image", "inspect", "--format",
                      "{{if .RepoDigests}}{{index .RepoDigests 0}}{{end}}", name])
        rows.append({
            "reference": name,
            "digest": digest if digest else "not pulled on this machine",
        })
    return rows


def describe_configs(gateway_dir):
    rows = []
    if gateway_dir is None:
        return rows
    root = Path(gateway_dir).expanduser()
    if not root.is_dir():
        print(f"estate-manifest: no such gateway directory: {root}", file=sys.stderr)
        return rows
    for path in sorted(root.rglob("*")):
        if not path.is_file() or path.suffix.lower() not in CONFIG_SUFFIXES:
            continue
        # A secrets file is named here so that its absence from the list is deliberate
        # rather than accidental, and its contents are never read.
        if path.name.startswith(".env"):
            rows.append({"path": str(path.relative_to(root)), "sha256": "not recorded",
                         "note": "excluded on purpose: this file holds keys"})
            continue
        rows.append({"path": str(path.relative_to(root)), "sha256": sha256_of(path)})
    return rows


def describe_machine():
    return {
        "system": platform.system(),
        "release": platform.release(),
        "machine": platform.machine(),
        "python": platform.python_version(),
        "processor": platform.processor(),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--models", default=None, help="model library directory")
    parser.add_argument("--gateway", default=None, help="the Part 9 gateway directory")
    parser.add_argument("--output", default=None, help="write here instead of standard output")
    parser.add_argument("--hash", dest="compute_hashes", action="store_true",
                        help="compute SHA-256 for model files with no .sha256 beside them")
    parser.add_argument("--quiet", action="store_true", help="do not report progress")
    args = parser.parse_args()

    if args.models is None and args.gateway is None:
        parser.error("give --models, --gateway, or both; there is nothing to describe otherwise")

    manifest = {
        "written": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "machine": describe_machine(),
        "engines": describe_engines(),
        "images": describe_images(args.gateway),
        "configs": describe_configs(args.gateway),
        "models": describe_models(args.models, args.compute_hashes, not args.quiet),
    }
    manifest["summary"] = {
        "model_files": len(manifest["models"]),
        "model_bytes": sum(m["bytes"] for m in manifest["models"]),
        "model_files_without_hash": sum(1 for m in manifest["models"] if not m["sha256"]),
        "config_files": len(manifest["configs"]),
        "engines_present": sum(1 for e in manifest["engines"] if e["present"]),
    }

    text = json.dumps(manifest, indent=2, sort_keys=True)
    if args.output:
        Path(args.output).expanduser().write_text(text + "\n", encoding="utf-8")
        if not args.quiet:
            print(f"wrote {args.output}: {manifest['summary']['model_files']} model file(s), "
                  f"{manifest['summary']['config_files']} configuration file(s)")
            missing = manifest["summary"]["model_files_without_hash"]
            if missing:
                print(f"{missing} model file(s) have no recorded hash. Run again with --hash "
                      "when you have time to read every byte.")
    else:
        print(text)


if __name__ == "__main__":
    main()