#!/usr/bin/env python3
import argparse
import hashlib
import io
import json
import os
import subprocess
import tarfile
import time
from collections import defaultdict
from pathlib import Path

CATALOG_FILES = {"data.pkg", "packagesite.pkg"}
LATEST_NAMES = {"latest.pkg", "os-rustproxy-latest.pkg", "rustproxy-latest.pkg"}


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


def read_manifest(pkg: Path) -> dict:
    for member in ("+COMPACT_MANIFEST", "+MANIFEST"):
        p = subprocess.run(
            ["bsdtar", "-xOf", str(pkg), member],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            check=False,
        )
        if p.returncode == 0 and p.stdout:
            return json.loads(p.stdout.decode("utf-8"))
    raise RuntimeError(f"{pkg.name}: missing +COMPACT_MANIFEST / +MANIFEST")


def is_payload_pkg(path: Path) -> bool:
    return (
        path.is_file()
        and path.suffix == ".pkg"
        and path.name not in CATALOG_FILES
        and path.name not in LATEST_NAMES
        and not path.name.endswith(".tmp.pkg")
    )


def collect_packages(repo: Path):
    packages = []
    for path in sorted(repo.rglob("*.pkg")):
        if not is_payload_pkg(path):
            continue
        manifest = read_manifest(path)
        manifest["path"] = path.relative_to(repo).as_posix()
        manifest["pkgsize"] = path.stat().st_size
        manifest["sum"] = sha256(path)
        packages.append((path, manifest))
    return packages


def enforce_retention(packages, retention: int) -> bool:
    if retention <= 0:
        return False
    groups = defaultdict(list)
    for path, manifest in packages:
        # Keep the requested number per repository subdirectory/ABI and package name.
        key = (path.parent, manifest.get("name", path.stem))
        groups[key].append((path, manifest))

    removed = False
    for (_parent, _name), entries in groups.items():
        entries.sort(key=lambda item: item[0].stat().st_mtime_ns, reverse=True)
        for path, _manifest in entries[retention:]:
            try:
                path.unlink()
                removed = True
                print(f"[opnsense-repo] retention: removed {path}", flush=True)
            except FileNotFoundError:
                pass
    return removed


def make_tar(dst: Path, member_name: str, payload: bytes):
    tmp = dst.with_name(dst.name + ".tmp")
    with tarfile.open(tmp, "w", format=tarfile.PAX_FORMAT) as tf:
        info = tarfile.TarInfo(member_name)
        info.size = len(payload)
        info.mode = 0o644
        info.mtime = int(time.time())
        tf.addfile(info, io.BytesIO(payload))
    os.replace(tmp, dst)


def rebuild(repo: Path, retention: int, repo_version: int = 1):
    packages = collect_packages(repo)
    if enforce_retention(packages, retention):
        packages = collect_packages(repo)

    manifests = [manifest for _, manifest in packages]
    packagesite = b"".join(
        json.dumps(m, separators=(",", ":"), ensure_ascii=False).encode("utf-8") + b"\n"
        for m in manifests
    )
    data = json.dumps(
        {"groups": [], "expired_packages": [], "packages": manifests},
        separators=(",", ":"),
        ensure_ascii=False,
    ).encode("utf-8")

    make_tar(repo / "packagesite.pkg", "packagesite.yaml", packagesite)
    make_tar(repo / "data.pkg", "data", data)

    meta = (
        f'version: {repo_version};\n'
        'packing_format: "tar";\n'
        'manifests: "packagesite.yaml";\n'
        'manifests_archive: "packagesite";\n'
        'data: "data";\n'
        'data_archive: "data";\n'
        'filesite: "files";\n'
        'filesite_archive: "files";\n'
    )
    tmp = repo / "meta.conf.tmp"
    tmp.write_text(meta, encoding="utf-8")
    os.replace(tmp, repo / "meta.conf")
    print(f"[opnsense-repo] rebuilt: {len(manifests)} package(s)", flush=True)


def main():
    parser = argparse.ArgumentParser(description="Build a static FreeBSD pkg repository catalogue")
    parser.add_argument("repo", type=Path, help="repository root")
    parser.add_argument("--retention", type=int, default=10, help="versions kept per package/ABI; 0 disables pruning")
    parser.add_argument("--repo-version", type=int, choices=(1, 2), default=1, help="pkg repository format; 2 for native FreeBSD 15")
    args = parser.parse_args()

    repo = args.repo.resolve()
    repo.mkdir(parents=True, exist_ok=True)
    rebuild(repo, args.retention, args.repo_version)


if __name__ == "__main__":
    main()