#!/usr/bin/env python3
"""Produce an offline release archive from an independently vetted ARM64 staging tree.

This tool does not build/download a runtime, model, native libraries, or licenses.
"""
import argparse
from datetime import datetime
import hashlib
import io
import json
import os
from pathlib import Path
import re
import sys
import tarfile
import tempfile

ARCH = "linux-aarch64"
ALLOWED = {"bin", "lib", "assets", "python", "scripts", "runtime", "model", "fonts", "licenses"}
REQUIRED = ("bin/frameyap", "bin/install.sh", "runtime/bin/python3", "python/frameyap/worker.py",
            "assets/actions.json", "fonts/font.ttf", "licenses/THIRD_PARTY_NOTICES.txt")


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--stage", type=Path, required=True, help="vetted standalone payload directory")
    p.add_argument("--output", type=Path, required=True, help="existing output directory")
    p.add_argument("--version", required=True)
    p.add_argument("--arch", choices=[ARCH], required=True)
    p.add_argument("--model-revision", required=True, help="exact vetted model revision identifier")
    p.add_argument("--external-runtime", action="store_true", help="omit ASR runtime; user must supply an independently authorized Python environment")
    args = p.parse_args(argv)
    if not re.fullmatch(r"(0|[1-9][0-9]*)\.(0|[1-9][0-9]*)\.[1-9][0-9]{11}", args.version):
        p.error("--version must be numeric MAJOR.MINOR.YYYYMMDDHHMM (no leading zeros, v prefix or git suffix)")
    try:
        datetime.strptime(args.version.rsplit(".", 1)[1], "%Y%m%d%H%M")
    except ValueError:
        p.error("--version contains an invalid UTC date/time")
    if not args.model_revision.strip():
        p.error("model revision must be nonempty")
    stage = args.stage.resolve(strict=True)
    if not stage.is_dir() or args.stage.is_symlink():
        p.error("stage must be a real directory")
    if not args.output.is_dir() or args.output.is_symlink():
        p.error("output must be an existing real directory")
    if any(item.name not in ALLOWED for item in stage.iterdir()):
        p.error("stage has unknown root entries or preexisting release.json")
    for name in REQUIRED:
        if args.external_runtime and name == "runtime/bin/python3":
            continue
        if not (stage / name).is_file() or (stage / name).is_symlink():
            p.error(f"missing file: {name}")
    if args.external_runtime and (stage / "runtime").exists():
        p.error("external-runtime package must not contain a runtime directory")
    for name in (("bin/frameyap", "bin/install.sh") if args.external_runtime else ("bin/frameyap", "bin/install.sh", "runtime/bin/python3")):
        if not os.access(stage / name, os.X_OK):
            p.error(f"not executable: {name}")
    for name in ("lib", "fonts", "licenses"):
        if not (stage / name).is_dir() or not any((stage / name).iterdir()):
            p.error(f"missing directory or empty: {name}")
    if not (stage / "licenses/THIRD_PARTY_NOTICES.txt").read_text().strip():
        p.error("third-party notices must not be empty; review redistribution licenses")
    if (stage / "model").exists() and not any((stage / "model").rglob("*")):
        p.error("model directory is empty")
    entries = sorted(stage.rglob("*"), key=lambda item: item.relative_to(stage).as_posix())
    if len(entries) > 49999:
        p.error("too many payload members")
    total = 0
    for entry in entries:
        name = entry.relative_to(stage).as_posix()
        if (len(name) > 1024 or len(name.split("/")) > 32 or "\\" in name
                or name.endswith("/") or any(part in ("", ".", "..") for part in name.split("/"))):
            p.error(f"unsupported archive member path: {name}")
        if entry.is_symlink() or not (entry.is_dir() or entry.is_file()):
            p.error(f"links and special files forbidden (copy vetted files into stage): {entry}")
        if entry.is_file():
            total += entry.stat().st_size
    if total > 12 * 1024**3:
        p.error("payload exceeds 12 GiB uncompressed limit")
    name = f"frameyap-{args.version}-{ARCH}.tar.gz"
    target = args.output / name
    sidecar = args.output / (name + ".sha256")
    if target.exists() or sidecar.exists() or target.is_symlink() or sidecar.is_symlink():
        p.error("release output exists; refusing to replace it")
    meta = {"schema": 1, "version": args.version, "arch": ARCH,
            "runtime": "external-authorized-python" if args.external_runtime else "bundled-cpu-python", "model_revision": args.model_revision,
            "includes_model": (stage / "model").is_dir()}
    fd, temp = tempfile.mkstemp(prefix=".frameyap-package-", dir=args.output)
    os.close(fd)
    try:
        with tarfile.open(temp, "w:gz", format=tarfile.PAX_FORMAT, compresslevel=1) as tar:
            payload = (json.dumps(meta, sort_keys=True, indent=2) + "\n").encode()
            info = tarfile.TarInfo("release.json")
            info.size = len(payload)
            info.mode = 0o600
            tar.addfile(info, io.BytesIO(payload))
            for entry in entries:
                tar.add(entry, arcname=entry.relative_to(stage).as_posix(), recursive=False)
        digest = hashlib.sha256()
        with open(temp, "rb") as stream:
            for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                digest.update(chunk)
        os.replace(temp, target)
        sidecar.write_text(f"{digest.hexdigest()}  {name}\n")
        print(f"Wrote {target} and {sidecar}")
    finally:
        if os.path.exists(temp):
            os.unlink(temp)


if __name__ == "__main__":
    main()