#!/usr/bin/env python3
"""Download and serve the public, verified Qwen3.5-9B TT-native P150 release.

Requires Python 3.11+, Docker, and an exclusively caller-owned P150 to serve.
No Hugging Face login, Python ML packages, or original model weights are needed.
Downloads are anonymous; serving is offline inside the digest-pinned container.
"""

import argparse
import hashlib
import json
import os
from pathlib import Path
import re
import stat
import sys
import tempfile
from urllib.error import URLError
from urllib.parse import quote
from urllib.request import Request, urlopen


MODEL_REPO = "Lottolabs/Qwen3.5-9B-TT-Mixed-BFP4-BFP8-P150"
RELEASE_ID = "prefill-mtp-20260907"
IMAGE_REPO = "ghcr.io/lottolottolotto/qwen35-tt-p150"
HF_BASE = "https://huggingface.co"
SHA256 = re.compile(r"[0-9a-f]{64}")
CHUNK_SIZE = 1024 * 1024


def request(url):
    # urllib does not load HF tokens, netrc credentials, or ML dependencies.
    return urlopen(Request(url, headers={"User-Agent": "qwen9b-native-launch/1"}), timeout=120)


def fetch_json(url):
    with request(url) as response:
        raw = response.read(16 * CHUNK_SIZE + 1)
    if len(raw) > 16 * CHUNK_SIZE:
        raise ValueError("Remote JSON exceeds the 16 MiB metadata limit")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("Expected a JSON object")
    return value, raw


def safe_relative(value):
    if (not isinstance(value, str) or not value
            or any(part in ("", ".", "..") for part in value.split("/"))
            or "\\" in value or ":" in value
            or any(ord(char) < 32 or ord(char) == 127 for char in value)):
        raise ValueError(f"Unsafe release path: {value!r}")
    return value


def file_records(value):
    if not isinstance(value, dict) or not value:
        raise ValueError("Manifest must contain a nonempty files object")
    for relative, record in value.items():
        safe_relative(relative)
        if (not isinstance(record, dict)
                or not isinstance(record.get("sha256"), str)
                or SHA256.fullmatch(record["sha256"]) is None
                or type(record.get("bytes")) is not int or record["bytes"] < 0):
            raise ValueError(f"Invalid file checksum/size: {relative}")
        parts = relative.split("/")
        if any("/".join(parts[:index]) in value for index in range(1, len(parts))):
            raise ValueError(f"File/directory collision: {relative}")
    return value


def validate_release(release):
    if (type(release.get("schema_version")) is not int or release["schema_version"] != 1
            or release.get("model_repo") != MODEL_REPO
            or release.get("release_id") != RELEASE_ID
            or release.get("checkpoint_subdir") != "checkpoint"):
        raise ValueError("Unsupported runtime release identity or schema")
    image = release.get("image")
    if (not isinstance(image, str) or re.fullmatch(
            re.escape(IMAGE_REPO) + r"(?::[A-Za-z0-9_][A-Za-z0-9_.-]*)?@sha256:[0-9a-f]{64}",
            image) is None):
        raise ValueError("Release image must pin the official GHCR repository by SHA256 digest")
    files = file_records(release.get("files"))
    required = {"serve_native.py", "checkpoint/native_manifest.json",
                "checkpoint/equivalence.json", "checkpoint/equivalence-mtp.json"}
    if not required.issubset(files):
        raise ValueError("Release is missing the server, native manifest, or MTP mode proofs")
    if "runtime-release.json" in files or "launch.py" in files:
        raise ValueError("Release inventory must exclude runtime-release.json and launch.py")
    if release.get("checkpoint_manifest_sha256") != files["checkpoint/native_manifest.json"]["sha256"]:
        raise ValueError("Release checkpoint manifest checksum does not match its file inventory")
    return files


def cache_path(root, relative, directory=False):
    """Do not follow pre-existing links or special files within a snapshot."""
    current = root
    parts = safe_relative(relative).split("/")
    for index, part in enumerate(parts):
        current = current / part
        if current.is_symlink():
            raise ValueError(f"Refusing a symlink in the release cache: {current}")
        if current.exists():
            mode = current.stat().st_mode
            expected = stat.S_ISDIR if directory or index < len(parts) - 1 else stat.S_ISREG
            if not expected(mode):
                raise ValueError(f"Unexpected cache entry type: {current}")
    return current


def verified(path, record):
    if not path.is_file() or path.stat().st_size != record["bytes"]:
        return False
    with path.open("rb") as source:
        return hashlib.file_digest(source, "sha256").hexdigest() == record["sha256"]


def atomic_write(destination, chunks, record=None):
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = None
    try:
        with tempfile.NamedTemporaryFile(prefix=".download-", dir=destination.parent, delete=False) as target:
            temporary = Path(target.name)
            digest = hashlib.sha256()
            size = 0
            for chunk in chunks:
                size += len(chunk)
                if record is not None and size > record["bytes"]:
                    raise ValueError(f"Download exceeds expected size: {destination.name}")
                target.write(chunk)
                digest.update(chunk)
            if record is not None and (size != record["bytes"] or digest.hexdigest() != record["sha256"]):
                raise ValueError(f"Download failed SHA256/size verification: {destination}")
            target.flush()
            os.fsync(target.fileno())
            # The container must be able to read the read-only checkpoint mount.
            os.fchmod(target.fileno(), 0o644)
        os.replace(temporary, destination)
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def download(snapshot, base_url, relative, record):
    destination = cache_path(snapshot, relative)
    if verified(destination, record):
        return
    print(f"Downloading {relative} ({record['bytes']:,} bytes)", file=sys.stderr, flush=True)
    with request(base_url + quote(relative, safe="/")) as response:
        atomic_write(destination, iter(lambda: response.read(CHUNK_SIZE), b""), record)


def independent_tensor_cache(snapshot, tensor_cache):
    if (tensor_cache == snapshot or tensor_cache.is_relative_to(snapshot)
            or snapshot.is_relative_to(tensor_cache)):
        raise ValueError("Tensor cache must be independent of the downloaded release directory")
    if ":" in str(snapshot) or ":" in str(tensor_cache):
        raise ValueError("Docker bind-mount paths cannot contain ':'")
    tensor_cache.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryFile(dir=tensor_cache) as probe:
        probe.write(b"writable")
        probe.flush()


def parser():
    cache_home = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
    result = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""Examples:
  python3 launch.py --download-only --cache-root /large-disk/qwen9b
  python3 launch.py --cache-root /large-disk/qwen9b --print-command
  python3 launch.py --cache-root /large-disk/qwen9b --device-ownership-confirmed
  python3 launch.py --mtp off --port 8001 --device-ownership-confirmed

--print-command downloads/verifies the release but does not run Docker.
--download-only downloads/verifies without requiring Docker or device ownership.
Serving never stops containers, resets the device, or takes ownership for you.
Stop your other accelerator workloads yourself before confirming ownership.
The API binds to 127.0.0.1 only. Keep the printed immutable revision to reproduce.
""")
    result.add_argument("--cache-root", type=Path, default=cache_home / "qwen9b-tt-native",
                        help="Download/cache root (default: $XDG_CACHE_HOME/qwen9b-tt-native or ~/.cache/qwen9b-tt-native)")
    result.add_argument("--tensor-cache", type=Path,
                        help="Independent writable tensor cache (default: CACHE_ROOT/tensor-cache/COMMIT)")
    result.add_argument("--port", type=int, default=8000, help="Loopback API port (default: 8000)")
    result.add_argument("--mtp", choices=("on", "off"), default="on", help="MTP-1 mode (default: on)")
    result.add_argument("--max-model-len", type=int, default=8192, help="Context limit, 128–8192 (default: 8192)")
    result.add_argument("--name", default="qwen9b-native-api", help="Docker container name (default: qwen9b-native-api)")
    result.add_argument("--revision", default="main", help="Public HF branch, tag, or commit (default: main)")
    result.add_argument("--device-ownership-confirmed", action="store_true",
                        help="Confirm you exclusively own the P150 and have stopped your other workloads")
    mode = result.add_mutually_exclusive_group()
    mode.add_argument("--print-command", action="store_true", help="Download, verify, and print the Docker command without running it")
    mode.add_argument("--download-only", action="store_true", help="Download and verify the native package without serving")
    return result


def main():
    cli = parser()
    args = cli.parse_args()
    if sys.version_info < (3, 11):
        cli.error("Python 3.11 or newer is required")
    if not 1 <= args.port <= 65535 or not 128 <= args.max_model_len <= 8192:
        cli.error("Port must be 1–65535 and max-model-len must be 128–8192")
    if re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", args.name) is None:
        cli.error("Invalid Docker container name")
    if not args.revision.strip():
        cli.error("Revision must not be empty")
    if not (args.print_command or args.download_only or args.device_ownership_confirmed):
        cli.error("Confirm exclusive P150 ownership with --device-ownership-confirmed, or use --print-command/--download-only")
    try:
        info, _ = fetch_json(f"{HF_BASE}/api/models/{MODEL_REPO}/revision/{quote(args.revision, safe='')}")
        commit = info.get("sha")
        if not isinstance(commit, str) or re.fullmatch(r"[0-9a-f]{40}", commit) is None:
            raise ValueError("Hugging Face did not resolve the revision to an immutable commit")
        base_url = f"{HF_BASE}/{MODEL_REPO}/resolve/{commit}/"
        release, release_raw = fetch_json(base_url + "runtime-release.json")
        files = validate_release(release)
        root = args.cache_root.expanduser().resolve()
        root.mkdir(parents=True, exist_ok=True)
        snapshot = cache_path(root, f"snapshots/{commit}", directory=True)
        snapshot.mkdir(parents=True, exist_ok=True)
        tensor_cache = (args.tensor_cache.expanduser() if args.tensor_cache else root / "tensor-cache" / commit).resolve()
        independent_tensor_cache(snapshot, tensor_cache)
        print(f"Revision: {commit}\nImage: {release['image']}\nCheckpoint: {snapshot / 'checkpoint'}\nTensor cache: {tensor_cache}",
              file=sys.stderr, flush=True)
        # Check the native inventory before downloading the multi-gigabyte tensors.
        manifest_name = "checkpoint/native_manifest.json"
        download(snapshot, base_url, manifest_name, files[manifest_name])
        native = json.loads((snapshot / manifest_name).read_bytes())
        if not isinstance(native, dict):
            raise ValueError("Native manifest must be a JSON object")
        for relative, record in file_records(native.get("files")).items():
            published = files.get("checkpoint/" + relative)
            if published is None or any(published[key] != record[key] for key in ("sha256", "bytes")):
                raise ValueError(f"Release inventory does not cover the native checkpoint: {relative}")
        for relative, record in files.items():
            if relative != manifest_name:
                download(snapshot, base_url, relative, record)
        atomic_write(cache_path(snapshot, "runtime-release.json"), (release_raw,))
        print("All release files verified.", file=sys.stderr, flush=True)
        if args.download_only:
            return 0
        command = [sys.executable, "-B", str(snapshot / "serve_native.py"),
                   "--checkpoint", str(snapshot / "checkpoint"),
                   "--cache-root", str(tensor_cache), "--image", release["image"],
                   "--port", str(args.port), "--name", args.name,
                   "--max-model-len", str(args.max_model_len), "--mtp", args.mtp]
        if args.print_command:
            command.append("--print-command")
        if args.device_ownership_confirmed:
            command.append("--device-ownership-confirmed")
        os.execv(sys.executable, command)
    except (OSError, URLError, ValueError) as error:
        cli.exit(1, f"launch.py: {error}\n")
    return 0


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