#!/usr/bin/env python3
from __future__ import annotations

import argparse
from collections import deque
from dataclasses import dataclass
import hashlib
from io import BytesIO
import json
import lzma
from pathlib import Path
import re
import tarfile
from urllib.parse import urljoin
from urllib.request import urlopen


PACKAGE_NAME = re.compile(r"^([a-z0-9][a-z0-9+.-]*(?::(?:any|native|arm64))?)")
ARCH_FILTER = re.compile(r"\[([^]]+)]")
VERSION_FILTER = re.compile(r"\((=|>=|<=|>>|<<)\s*([^)\s]+)\)")
IMAGE_DEBIAN_ROOTS = ("ffmpeg", "libheif-examples", "python3.11-venv")


@dataclass(frozen=True)
class Dependency:
    name: str
    operator: str = ""
    version: str = ""


@dataclass(frozen=True)
class PackageRecord:
    name: str
    version: str
    architecture: str
    dependencies: tuple[tuple[Dependency, ...], ...]
    provides: tuple[str, ...]
    filename: str = ""
    sha256: str = ""
    size: int = 0
    repository_base: str = ""
    local_path: Path | None = None


def normalize_package_name(value: str) -> str:
    value = value.strip()
    if ":" in value:
        name, qualifier = value.rsplit(":", 1)
        if qualifier in {"any", "native", "arm64"}:
            return name
    return value


def _arch_applies(value: str, architecture: str = "arm64") -> bool:
    match = ARCH_FILTER.search(value)
    if not match:
        return True
    filters = match.group(1).split()
    positives = {item for item in filters if not item.startswith("!")}
    negatives = {item[1:] for item in filters if item.startswith("!")}
    return architecture not in negatives and (not positives or architecture in positives)


def parse_dependency_field(value: str) -> tuple[tuple[Dependency, ...], ...]:
    groups: list[tuple[Dependency, ...]] = []
    for raw_group in value.split(","):
        alternatives: list[Dependency] = []
        for raw_alternative in raw_group.split("|"):
            if not _arch_applies(raw_alternative):
                continue
            match = PACKAGE_NAME.match(raw_alternative.strip())
            if not match:
                continue
            version_match = VERSION_FILTER.search(raw_alternative)
            alternatives.append(
                Dependency(
                    normalize_package_name(match.group(1)),
                    version_match.group(1) if version_match else "",
                    version_match.group(2) if version_match else "",
                )
            )
        if alternatives:
            groups.append(tuple(alternatives))
    return tuple(groups)


def parse_control_stanzas(text: str) -> list[dict[str, str]]:
    stanzas: list[dict[str, str]] = []
    current: dict[str, str] = {}
    current_key = ""
    for line in text.splitlines() + [""]:
        if not line:
            if current:
                stanzas.append(current)
                current = {}
                current_key = ""
            continue
        if line[0].isspace() and current_key:
            current[current_key] += " " + line.strip()
            continue
        key, separator, value = line.partition(":")
        if separator:
            current_key = key
            current[key] = value.strip()
    return stanzas


def record_from_fields(
    fields: dict[str, str], *, repository_base: str = "", local_path: Path | None = None
) -> PackageRecord:
    dependencies = ", ".join(
        value for key in ("Pre-Depends", "Depends") if (value := fields.get(key, ""))
    )
    provides = tuple(
        normalize_package_name(item.strip().split()[0])
        for item in fields.get("Provides", "").split(",")
        if item.strip()
    )
    return PackageRecord(
        name=normalize_package_name(fields["Package"]),
        version=fields["Version"],
        architecture=fields.get("Architecture", "arm64"),
        dependencies=parse_dependency_field(dependencies),
        provides=provides,
        filename=fields.get("Filename", ""),
        sha256=fields.get("SHA256", "").lower(),
        size=int(fields.get("Size", "0")),
        repository_base=repository_base,
        local_path=local_path,
    )


def load_package_index(path: Path, repository_base: str) -> dict[str, PackageRecord]:
    path = Path(path)
    if path.suffix == ".xz":
        with lzma.open(path, "rt", encoding="utf-8") as handle:
            text = handle.read()
    else:
        text = path.read_text(encoding="utf-8")
    records: dict[str, PackageRecord] = {}
    for fields in parse_control_stanzas(text):
        if fields.get("Architecture") not in {"arm64", "all"}:
            continue
        record = record_from_fields(fields, repository_base=repository_base)
        records.setdefault(record.name, record)
    return records


def _read_ar_member(path: Path, wanted_prefix: str) -> bytes:
    data = Path(path).read_bytes()
    if not data.startswith(b"!<arch>\n"):
        raise ValueError(f"not a Debian ar archive: {path}")
    offset = 8
    while offset + 60 <= len(data):
        header = data[offset : offset + 60]
        name = header[:16].decode("ascii").strip().rstrip("/")
        size = int(header[48:58].decode("ascii").strip())
        start = offset + 60
        end = start + size
        if name.startswith(wanted_prefix):
            return data[start:end]
        offset = end + (size % 2)
    raise ValueError(f"{wanted_prefix} is missing from {path}")


def read_deb_control(path: Path) -> dict[str, str]:
    control_archive = _read_ar_member(path, "control.tar")
    with tarfile.open(fileobj=BytesIO(control_archive), mode="r:*") as archive:
        member = next(
            (item for item in archive.getmembers() if item.name.lstrip("./") == "control"),
            None,
        )
        if member is None:
            raise ValueError(f"control file is missing from {path}")
        handle = archive.extractfile(member)
        if handle is None:
            raise ValueError(f"control file cannot be read from {path}")
        stanzas = parse_control_stanzas(handle.read().decode("utf-8"))
    if len(stanzas) != 1:
        raise ValueError(f"unexpected control data in {path}")
    return stanzas[0]


def load_local_packages(directory: Path) -> dict[str, PackageRecord]:
    records: dict[str, PackageRecord] = {}
    for path in sorted(Path(directory).glob("*.deb")):
        record = record_from_fields(read_deb_control(path), local_path=path.resolve())
        if record.name in records:
            raise ValueError(f"duplicate package {record.name} in {directory}")
        records[record.name] = record
    return records


def load_installed_inventory(path: Path) -> dict[str, str]:
    installed: dict[str, str] = {}
    for number, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), 1):
        if not line:
            continue
        parts = line.split("\t")
        if len(parts) != 3:
            raise ValueError(f"invalid inventory line {number}")
        status, name, version = parts
        if status == "ii ":
            installed[normalize_package_name(name)] = version
    return installed


def _exact_version_matches(dependency: Dependency, version: str) -> bool:
    return dependency.operator != "=" or dependency.version == version


def resolve_closure(
    roots: list[str],
    installed: dict[str, str],
    local: dict[str, PackageRecord],
    available: dict[str, PackageRecord],
) -> tuple[dict[str, PackageRecord], list[str]]:
    installed = dict(installed)
    for name in local:
        installed.pop(name, None)
    providers: dict[str, list[PackageRecord]] = {}
    for record in [*local.values(), *available.values()]:
        for provided in record.provides:
            providers.setdefault(provided, []).append(record)
    selected: dict[str, PackageRecord] = {}
    queue: deque[str] = deque(normalize_package_name(item) for item in roots)
    unresolved: list[str] = []

    def candidate_for(dependency: Dependency) -> PackageRecord | None:
        direct = local.get(dependency.name) or available.get(dependency.name)
        if direct and _exact_version_matches(dependency, direct.version):
            return direct
        for provider in providers.get(dependency.name, []):
            if _exact_version_matches(dependency, provider.version):
                return provider
        return None

    while queue:
        name = queue.popleft()
        if name in selected:
            continue
        record = local.get(name) or available.get(name)
        if record is None:
            unresolved.append(f"required package is unavailable: {name}")
            continue
        selected[name] = record
        for group in record.dependencies:
            if any(
                dependency.name in installed
                and _exact_version_matches(dependency, installed[dependency.name])
                for dependency in group
            ):
                continue
            if any(
                dependency.name in selected
                and _exact_version_matches(dependency, selected[dependency.name].version)
                for dependency in group
            ):
                continue
            choice = next((candidate_for(dependency) for dependency in group if candidate_for(dependency)), None)
            if choice is None:
                unresolved.append(
                    f"{record.name} cannot satisfy: "
                    + " | ".join(
                        f"{item.name} {item.operator} {item.version}".strip() for item in group
                    )
                )
                continue
            queue.append(choice.name)
    return selected, sorted(set(unresolved))


def verify_local_closure(directory: Path, inventory: Path, roots: list[str]) -> dict[str, PackageRecord]:
    local = load_local_packages(directory)
    selected, unresolved = resolve_closure(roots, load_installed_inventory(inventory), local, {})
    missing = sorted(set(roots) - set(selected))
    if missing or unresolved:
        detail = "; ".join([*(f"missing root: {item}" for item in missing), *unresolved])
        raise ValueError(f"offline Debian dependency closure is incomplete: {detail}")
    return selected


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


def download_records(records: dict[str, PackageRecord], output: Path) -> list[Path]:
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    downloaded: list[Path] = []
    for record in sorted(records.values(), key=lambda item: item.name):
        if record.local_path is not None:
            continue
        if not record.filename or not record.sha256 or not record.size:
            raise ValueError(f"repository metadata is incomplete for {record.name}")
        target = output / Path(record.filename).name
        if target.exists():
            if target.stat().st_size != record.size or _sha256_file(target) != record.sha256:
                raise ValueError(f"existing download does not match repository metadata: {target}")
            downloaded.append(target)
            continue
        url = urljoin(record.repository_base.rstrip("/") + "/", record.filename)
        partial = target.with_suffix(target.suffix + ".partial")
        with urlopen(url, timeout=60) as response, partial.open("wb") as handle:
            while chunk := response.read(1024 * 1024):
                handle.write(chunk)
        if partial.stat().st_size != record.size or _sha256_file(partial) != record.sha256:
            partial.unlink(missing_ok=True)
            raise ValueError(f"download verification failed: {record.name}")
        partial.replace(target)
        downloaded.append(target)
    return downloaded


def main() -> int:
    parser = argparse.ArgumentParser(description="Resolve a minimal Debian 12/AArch64 offline closure")
    parser.add_argument("--inventory", type=Path, required=True)
    parser.add_argument("--local", type=Path, required=True)
    parser.add_argument("--index", action="append", default=[], metavar="PATH=REPOSITORY_BASE")
    parser.add_argument("--root", action="append", required=True)
    parser.add_argument("--download-to", type=Path)
    parser.add_argument("--report", type=Path)
    args = parser.parse_args()

    available: dict[str, PackageRecord] = {}
    index_info: list[dict[str, str]] = []
    for value in args.index:
        path_text, separator, repository_base = value.partition("=")
        if not separator:
            parser.error("--index must be PATH=REPOSITORY_BASE")
        path = Path(path_text).resolve()
        for name, record in load_package_index(path, repository_base).items():
            available.setdefault(name, record)
        index_info.append({"file": path.name, "sha256": _sha256_file(path), "base": repository_base})

    local = load_local_packages(args.local)
    selected, unresolved = resolve_closure(
        args.root, load_installed_inventory(args.inventory), local, available
    )
    if unresolved:
        raise SystemExit("\n".join(unresolved))
    downloaded = download_records(selected, args.download_to) if args.download_to else []
    report = {
        "schema_version": 1,
        "architecture": "arm64",
        "roots": args.root,
        "inventory_sha256": _sha256_file(args.inventory),
        "indices": index_info,
        "selected": [
            {
                "package": item.name,
                "version": item.version,
                "source": "local" if item.local_path else item.filename,
                "sha256": _sha256_file(item.local_path) if item.local_path else item.sha256,
                "size": item.local_path.stat().st_size if item.local_path else item.size,
            }
            for item in sorted(selected.values(), key=lambda record: record.name)
        ],
        "downloaded_files": [item.name for item in downloaded],
    }
    rendered = json.dumps(report, ensure_ascii=False, indent=2) + "\n"
    if args.report:
        args.report.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")
    return 0


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