#!/usr/bin/env python3
"""Resolve one Python 3.11 wheel closure into candidate payload-lock objects."""

from __future__ import annotations

import argparse
import email
import hashlib
import json
import re
import subprocess
import zipfile
from pathlib import Path


def requirements(path: Path) -> list[str]:
    result = []
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.split("#", 1)[0].strip()
        if line:
            result.append(line)
    return result


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 metadata(path: Path) -> email.message.Message:
    with zipfile.ZipFile(path) as archive:
        # Some wheels vendor other distributions (for example setuptools). Only the
        # wheel's root dist-info directory describes the payload being locked.
        names = [
            name for name in archive.namelist()
            if name.endswith(".dist-info/METADATA") and name.count("/") == 1
        ]
        if len(names) != 1:
            raise RuntimeError(f"{path.name}: expected one dist-info/METADATA, found {len(names)}")
        return email.message_from_bytes(archive.read(names[0]))


def normalize(name: str) -> str:
    return re.sub(r"[-_.]+", "-", name).lower()


def resolve_group(group: str, requested: list[str], work: Path, resolver_pip: str) -> list[dict]:
    group_work = work / group
    wheels = group_work / "wheels"
    wheels.mkdir(parents=True, exist_ok=True)
    combined = group_work / "requirements.in"
    combined.write_text("\n".join(dict.fromkeys(requested)) + "\n", encoding="utf-8")

    venv = work / "resolver-venv"
    if not (venv / "bin/python").exists():
        subprocess.run(["python3", "-m", "venv", str(venv)], check=True)
    pip = venv / "bin/pip"
    subprocess.run([str(pip), "install", "--disable-pip-version-check", f"pip=={resolver_pip}"], check=True)
    command = [
        str(pip), "download", "--disable-pip-version-check", "--only-binary=:all:",
        "--python-version", "3.11", "--implementation", "cp", "--abi", "cp311",
        "--dest", str(wheels), "--requirement", str(combined),
    ]
    subprocess.run(command, check=True)

    objects = []
    versions: dict[str, str] = {}
    for path in sorted(wheels.glob("*.whl")):
        meta = metadata(path)
        name = meta.get("Name")
        version = meta.get("Version")
        if not name or not version:
            raise RuntimeError(f"{path.name}: missing Name or Version metadata")
        key = normalize(name)
        previous = versions.setdefault(key, version)
        if previous != version:
            raise RuntimeError(f"multiple versions resolved for {name}: {previous}, {version}")
        license_value = str(
            meta.get("License-Expression") or meta.get("License") or "REVIEW-REQUIRED"
        )
        if len(license_value) > 160 or "\n" in license_value or not license_value.strip():
            license_value = "REVIEW-REQUIRED"
        architecture = "all" if path.name.endswith("-none-any.whl") else "amd64"
        objects.append({
            "id": f"python-{group}-{key}",
            "version": version,
            "architecture": architecture,
            "source": f"pypi://{key}/{version}/{path.name}",
            "target": f"payload/python/{group}/{path.name}",
            "sha256": sha256_file(path),
            "bytes": path.stat().st_size,
            "bytes_basis": "measured",
            "license": license_value.strip(),
        })
    if not objects:
        raise RuntimeError(f"Python resolution produced no wheels for {group}")
    return objects


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--lock", type=Path, required=True)
    parser.add_argument("--backend-requirements", type=Path, required=True)
    parser.add_argument("--platform-requirements", type=Path, required=True)
    parser.add_argument("--work", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--resolver-pip", default="25.1.1")
    args = parser.parse_args()

    lock = json.loads(args.lock.read_text(encoding="utf-8"))
    platform_requested = requirements(args.platform_requirements)
    chromadb_requested = [item for item in platform_requested if normalize(item.split("=", 1)[0]) == "chromadb"]
    remaining_platform = [item for item in platform_requested if item not in chromadb_requested]
    groups = {
        "backend": requirements(args.backend_requirements),
        "chromadb": chromadb_requested,
        "platform": remaining_platform,
    }
    if not groups["chromadb"]:
        raise RuntimeError("platform requirements do not declare the ChromaDB wheelhouse")
    objects = []
    for group, requested in groups.items():
        objects.extend(resolve_group(group, requested, args.work, args.resolver_pip))
    lock["objects"]["python"] = objects
    lock["unresolved"] = [x for x in lock["unresolved"] if x != "hash-locked Python dependency closure"]
    lock["unresolved"].extend([
        "Python index snapshot/origin and resolver report approval",
        "Python wheel license/notice review",
        "Python 3.11 clean-environment import and offline installation evidence",
        "GPU framework compatibility with approved CUDA 12.8.1 runtime"
    ])
    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"Python wheel objects: {len(objects)}")
    for group in groups:
        print(f"Python {group} wheel objects: {sum(item['id'].startswith(f'python-{group}-') for item in objects)}")
    print(f"Python measured bytes: {sum(item['bytes'] for item in objects)}")
    print(f"License review required: {sum(item['license'] == 'REVIEW-REQUIRED' for item in objects)}")
    print(f"Wrote candidate lock: {args.output}")
    return 0


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