#!/usr/bin/env python3
"""Build-time only: bundle ELF dependencies into a static launcher on RK3588.

Deployment needs no Python, SDK, linker, or external tokenizer/graph files;
the default portable package also carries the selected RKNN runtime dependencies.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import re
import shutil
import struct
import subprocess
import tempfile


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--engine", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--notices", type=Path, required=True)
    parser.add_argument("--without-rknn-runtime", action="store_true",
                        help="Use the device-provided librknnrt.so instead of bundling it")
    args = parser.parse_args()
    if platform.machine() != "aarch64":
        raise SystemExit("Run packaging on the aarch64 build machine")
    source = Path(__file__).resolve().parent / "portable_launcher.c"
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="ling3-package-build-") as temp:
        stage = Path(temp)
        launcher = stage / "launcher"
        subprocess.run(["cc", "-O2", "-static", "-o", str(launcher), str(source)], check=True)
        subprocess.run(["strip", str(launcher)], check=True)
        files = {"mindnano-infer-engine": args.engine.resolve()}
        pending = [args.engine.resolve()]
        while pending:
            binary = pending.pop()
            output = subprocess.check_output(["ldd", str(binary)], text=True)
            if "not found" in output:
                raise RuntimeError("Unresolved dependency: " + output)
            for line in output.splitlines():
                match = re.search(r"(?:=>\s*)?(/[^\s]+)\s+\(", line)
                if not match:
                    continue
                path = Path(match.group(1))
                if path.name == "librknnrt.so" and args.without_rknn_runtime:
                    continue
                if path.name not in files:
                    files[path.name] = path.resolve()
                    pending.append(path.resolve())
        if "ld-linux-aarch64.so.1" not in files:
            raise RuntimeError("Missing matching private dynamic loader")
        for path in sorted(args.notices.iterdir()):
            if path.is_file():
                files["NOTICE." + path.name] = path
        # Keep the exact dependency versions and hashes inspectable inside the binary.
        manifest = {
            "format": "L3PORT01", "model": "mindnano-ling3-tiny",
            "minimum_qualified_rknpu_driver": "0.9.8", "native_context": 131072,
            "default_context": 4096,
            "files": {name: {"bytes": path.stat().st_size,
                              "sha256": hashlib.file_digest(path.open("rb"), "sha256").hexdigest()}
                      for name, path in files.items()},
        }
        manifest_path = stage / "MANIFEST.json"
        manifest_path.write_text(json.dumps(manifest, indent=2) + "\n")
        files["MANIFEST.json"] = manifest_path
        temporary = args.output.with_suffix(".building")
        try:
            with temporary.open("wb") as out:
                with launcher.open("rb") as inp:
                    shutil.copyfileobj(inp, out)
                offset = out.tell()
                out.write(struct.pack("<I", len(files)))
                for name, path in files.items():
                    encoded = name.encode("ascii")
                    out.write(struct.pack("<IQ", len(encoded), path.stat().st_size))
                    out.write(encoded)
                    with path.open("rb") as inp:
                        shutil.copyfileobj(inp, out)
                out.write(b"L3PORT01" + struct.pack("<Q", offset))
            temporary.chmod(0o755)
            os.replace(temporary, args.output)
        finally:
            temporary.unlink(missing_ok=True)
        shutil.copyfile(manifest_path, args.output.parent / "MANIFEST.json")
    print(json.dumps({"binary": str(args.output), "bytes": args.output.stat().st_size,
                      "embedded_files": len(files)}, indent=2))


if __name__ == "__main__":
    main()