#!/usr/bin/env python3
"""Resolve an Ubuntu APT dependency closure into candidate payload-lock objects."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
from pathlib import Path


def filtered_lines(path: Path) -> list[str]:
    return [
        line.strip() for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.lstrip().startswith("#")
    ]


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def deb_field(path: Path, field: str) -> str:
    result = subprocess.run(
        ["dpkg-deb", "-f", str(path), field], check=True, capture_output=True, text=True
    )
    return result.stdout.strip()


def resolve(packages: list[str], work: Path) -> list[dict]:
    archives = work / "archives"
    archives.mkdir(parents=True, exist_ok=True)
    (archives / "partial").mkdir(exist_ok=True)
    status = work / "empty-dpkg-status"
    status.write_text("", encoding="utf-8")
    command = [
        "apt-get", "--download-only", "--yes", "--no-install-recommends",
        "-o", "Debug::NoLocking=1",
        "-o", f"Dir::State::status={status}",
        "-o", f"Dir::Cache::archives={archives}",
        "-o", "APT::Get::AllowUnauthenticated=false",
        "install", *packages,
    ]
    subprocess.run(command, check=True)
    debs = sorted(archives.glob("*.deb"))
    if not debs:
        raise RuntimeError("APT resolution produced no .deb files")
    objects = []
    for path in debs:
        package = deb_field(path, "Package")
        version = deb_field(path, "Version")
        architecture = deb_field(path, "Architecture")
        if architecture not in {"amd64", "all"}:
            raise RuntimeError(f"unexpected architecture {architecture} for {path.name}")
        digest = sha256_file(path)
        objects.append({
            "id": f"apt-{package}",
            "version": version,
            "architecture": architecture,
            "source": f"apt://ubuntu-22.04/{package}/{version}/{path.name}",
            "target": f"payload/apt/{path.name}",
            "sha256": digest,
            "bytes": path.stat().st_size,
            "bytes_basis": "measured",
            "license": "REVIEW-REQUIRED",
            "requested_directly": package in packages,
        })
    return objects


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--lock", type=Path, required=True)
    parser.add_argument("--packages", type=Path, required=True)
    parser.add_argument("--work", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    lock = json.loads(args.lock.read_text(encoding="utf-8"))
    packages = filtered_lines(args.packages)
    objects = resolve(packages, args.work)
    lock["objects"]["apt"] = objects
    lock["unresolved"] = [x for x in lock["unresolved"] if x != "exact APT dependency closure"]
    lock["unresolved"].extend([
        "APT repository snapshot/index digests and origin approval",
        "APT package license mapping and notice extraction",
        "approved R570/CUDA 12.8.1/cuDNN/NVIDIA Container Toolkit package closure"
    ])
    lock["unresolved"] = list(dict.fromkeys(lock["unresolved"]))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(lock, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(f"APT requested packages: {len(packages)}")
    print(f"APT resolved .deb objects: {len(objects)}")
    print(f"APT measured bytes: {sum(item['bytes'] for item in objects)}")
    print(f"Wrote candidate lock: {args.output}")
    return 0


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