#!/usr/bin/env python3
"""Rebuild a provenance-bound LeKiwi input ZIP without modifying the source repo."""

import argparse
import hashlib
import json
import urllib.request
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path, PurePosixPath

ROOT = Path(__file__).resolve().parents[2]
PROFILE = json.loads((ROOT / "robot_profiles/lekiwi-v1.json").read_text())
SOURCE = PROFILE["source"]
MAX_TOTAL = 128 * 1024 * 1024


def digest(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def prepare(source: Path | None, output: Path, download: bool = False) -> Path:
    if source is not None:
        source = source.resolve()
        if source.name == "URDF":
            source = source.parent
        if not (source / "URDF/LeKiwi.urdf").is_file():
            raise ValueError("Source must be the LeKiwi repository or its URDF directory")
        if output.resolve().is_relative_to(source):
            raise ValueError("Output must not modify the reference repository")
    elif not download:
        raise ValueError("Pass --source or explicitly opt in with --download")
    files: dict[str, bytes] = {}
    total = 0

    def read(relative: str) -> bytes:
        nonlocal total
        path = PurePosixPath(relative)
        if path.is_absolute() or ".." in path.parts or "\\" in relative:
            raise ValueError(f"Unsafe asset path: {relative}")
        if source is not None:
            file = (source / relative).resolve()
            if not file.is_relative_to(source) or file.stat().st_size > MAX_TOTAL:
                raise ValueError(f"Invalid asset: {relative}")
            data = file.read_bytes()
        else:
            url = f"https://raw.githubusercontent.com/SIGRobotics-UIUC/LeKiwi/{SOURCE['revision']}/{relative}"
            with urllib.request.urlopen(url, timeout=60) as response:
                data = response.read(MAX_TOTAL + 1)
        total += len(data)
        if total > MAX_TOTAL:
            raise ValueError("Assets exceed 128 MiB")
        files[relative] = data
        return data

    urdf = read("URDF/LeKiwi.urdf")
    if digest(urdf) != SOURCE["urdfSha256"]:
        raise ValueError("URDF revision/hash is not supported by lekiwi-v1")
    tree = ET.fromstring(urdf)
    for mesh in sorted({m.get("filename", "") for m in tree.iter("mesh")}):
        if not mesh.startswith("meshes/") or not mesh.endswith(".stl"):
            raise ValueError(f"Unexpected mesh reference: {mesh}")
        read(f"URDF/{mesh}")
    read("LICENSE.txt")
    read("CITATION.cff")
    manifest = {
        "profileId": PROFILE["id"],
        "profileVersion": PROFILE["version"],
        "source": SOURCE,
        "files": {p: digest(data) for p, data in files.items()},
    }
    files["source-manifest.json"] = json.dumps(manifest, indent=2).encode()
    files["robot-profile.json"] = json.dumps(
        {"id": PROFILE["id"], "version": PROFILE["version"]}
    ).encode()
    files["SIMULATION-NOTICE.md"] = (
        "# LeKiwi simulation input\n\nSource: "
        + SOURCE["repository"]
        + "\n\nRevision: "
        + SOURCE["revision"]
        + "\n\nOriginal URDF/STL files are unmodified, Apache-2.0. "
        "The platform applies an explicit simulation-only profile after conversion: estimated "
        "mass/inertia/limits/servos and simplified passive-roller contacts. Not a calibrated "
        "hardware model; no cameras, training or grasping guarantee.\n"
    ).encode()
    output.mkdir(parents=True, exist_ok=True)
    for relative, data in files.items():
        target = output / relative
        if target.is_symlink() or not target.resolve().is_relative_to(output.resolve()):
            raise ValueError("Unsafe output path")
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(data)
    archive = output / "lekiwi-v1.zip"
    if archive.is_symlink() or not archive.resolve().is_relative_to(output.resolve()):
        raise ValueError("Unsafe archive output path")
    with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as bundle:
        for relative, data in files.items():
            info = zipfile.ZipInfo(relative, date_time=(1980, 1, 1, 0, 0, 0))
            info.compress_type = zipfile.ZIP_DEFLATED
            info.external_attr = 0o100644 << 16
            bundle.writestr(info, data, compresslevel=9)
    return archive


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--source", type=Path)
    group.add_argument("--download", action="store_true")
    parser.add_argument("--output", type=Path, default=ROOT / "build/lekiwi")
    args = parser.parse_args()
    print(prepare(args.source, args.output, args.download))


if __name__ == "__main__":
    main()