#!/usr/bin/env python3
"""Stage an existing local VideoMAE-v2 Large snapshot; no network access."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import tempfile
import shutil


MODEL_ID = "VideoMAEv2-Large"
REVISION = "9981a9c8f77118c421e5228e1b219468a4b0238d"


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while block := handle.read(8 * 1024 * 1024):
            digest.update(block)
    return digest.hexdigest()


def canonical_sha256(value) -> str:
    payload = json.dumps(
        value, sort_keys=True, separators=(",", ":"), allow_nan=False
    ).encode("utf-8")
    return hashlib.sha256(payload).hexdigest()


def atomic_json(path: Path, value) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, name = tempfile.mkstemp(
        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            json.dump(value, handle, sort_keys=True, indent=2)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(name, path)
    finally:
        Path(name).unlink(missing_ok=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output-dir",
        default=(
            "assets/models/"
            "VideoMAEv2-Large"),
    )
    parser.add_argument("--source-dir", required=True, type=Path,
                        help="Existing local snapshot; not a hosted repository ID")
    args = parser.parse_args()
    output = Path(args.output_dir).expanduser().resolve()
    source = args.source_dir.expanduser().resolve()
    if not source.is_dir():
        parser.error("--source-dir must be an existing local directory")
    if source == output or source in output.parents or output in source.parents:
        parser.error("Source and output directories must be separate and non-nested")
    if output.exists() and (not output.is_dir() or any(output.iterdir())):
        parser.error("--output-dir must be absent or empty")
    if not (source / "config.json").is_file():
        parser.error("The local snapshot must include config.json")
    if not any(source.glob("*.safetensors")) and not any(source.glob("*.bin")):
        parser.error("The local snapshot must include model weights")
    shutil.copytree(source, output, dirs_exist_ok=True,
                    ignore=shutil.ignore_patterns(".git", ".cache", "__pycache__", ".DS_Store"))
    # Existing frozen banks depend on the original manifest digest. Copy that
    # manifest unchanged when supplied, rather than rewriting its identity.
    existing_manifest = output / "download_manifest.json"
    if existing_manifest.is_file():
        manifest = json.loads(existing_manifest.read_text())
        payload = dict(manifest)
        saved = payload.pop("content_hash", None)
        if manifest.get("revision") != REVISION or canonical_sha256(payload) != saved:
            raise ValueError("Invalid pinned local snapshot manifest")
        for item in manifest["files"]:
            path = (output / item["path"]).resolve()
            if (output not in path.parents or path.stat().st_size != item["size_bytes"]
                    or file_sha256(path) != item["sha256"]):
                raise ValueError(f"Local snapshot file identity mismatch: {path}")
        print(json.dumps({"output_dir": str(output), "file_count": len(manifest["files"]),
                          "content_hash": saved, "preserved_manifest": True}, sort_keys=True))
        return
    files = []
    for path in sorted(output.rglob("*")):
        if (
            not path.is_file()
            or path.name == "download_manifest.json"
            or ".cache" in path.parts
        ):
            continue
        files.append({
            "path": path.relative_to(output).as_posix(),
            "size_bytes": path.stat().st_size,
            "sha256": file_sha256(path),
        })
    payload = {
        "schema_version": 1,
        "model_id": MODEL_ID,
        "revision": REVISION,
        "files": files,
    }
    payload["content_hash"] = canonical_sha256(payload)
    atomic_json(output / "download_manifest.json", payload)
    print(json.dumps({
        "output_dir": str(output),
        "file_count": len(files),
        "content_hash": payload["content_hash"],
    }, sort_keys=True))


if __name__ == "__main__":
    main()