#!/usr/bin/env python3
"""Download, verify, and lock exact GovOPlaN registry package artifacts."""

from __future__ import annotations

import argparse
import base64
from email.parser import BytesParser
from email.policy import compat32
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
from urllib.parse import quote, urlsplit, urlunsplit
import zipfile


NAME = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
WEBUI_NAME = re.compile(r"^@govoplan/[a-z0-9]+(?:-[a-z0-9]+)*-webui$")
VERSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._+!-]{0,127}$")
COMMIT = re.compile(r"^[0-9a-f]{40}$")
MAX_ARTIFACT_BYTES = 512 * 1024 * 1024


class PackageArtifactError(ValueError):
    """Registry artifacts do not match the selected package set."""


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--package-set", type=Path, required=True)
    parser.add_argument("--wheelhouse", type=Path, required=True)
    parser.add_argument("--webui-packages", type=Path, required=True)
    parser.add_argument("--lock-output", type=Path, required=True)
    parser.add_argument("--requirements-output", type=Path)
    parser.add_argument("--python", default=sys.executable)
    parser.add_argument("--npm", default="npm")
    return parser


def resolve(args: argparse.Namespace) -> dict[str, object]:
    package_set = _load_package_set(args.package_set)
    wheelhouse = args.wheelhouse.expanduser().resolve()
    webui_packages = args.webui_packages.expanduser().resolve()
    _require_empty_destination(wheelhouse)
    _require_empty_destination(webui_packages)
    wheelhouse.parent.mkdir(parents=True, exist_ok=True)
    webui_packages.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="govoplan-package-resolution-") as value:
        temporary = Path(value)
        wheels = temporary / "wheels"
        webui = temporary / "webui"
        wheels.mkdir()
        webui.mkdir()
        _download_wheels(
            packages=tuple(package_set["python"]),
            destination=wheels,
            python=args.python,
            index_url=str(package_set["registries"]["python"]),
        )
        _download_webui(
            packages=tuple(package_set["webui"]),
            destination=webui,
            npm=args.npm,
            registry=str(package_set["registries"]["npm"]),
        )
        python_rows = _verify_wheels(tuple(package_set["python"]), wheels)
        webui_rows = _verify_webui(tuple(package_set["webui"]), webui)
        lock: dict[str, object] = {
            "schema_version": "1",
            "release_version": package_set["release_version"],
            "package_set_sha256": package_set["package_set_sha256"],
            "registries": package_set["registries"],
            "python": python_rows,
            "webui": webui_rows,
        }
        lock["lock_sha256"] = _canonical_sha256(lock)
        shutil.copytree(wheels, wheelhouse, dirs_exist_ok=True)
        shutil.copytree(webui, webui_packages, dirs_exist_ok=True)
    args.lock_output.parent.mkdir(parents=True, exist_ok=True)
    args.lock_output.write_text(json.dumps(lock, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    if args.requirements_output is not None:
        _write_requirements(args.requirements_output, python_rows)
    return lock


def _load_package_set(path: Path) -> dict[str, object]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict) or value.get("schema_version") != "1":
        raise PackageArtifactError("package set has an unsupported shape")
    expected_hash = value.get("package_set_sha256")
    unsigned = dict(value)
    unsigned.pop("package_set_sha256", None)
    if expected_hash != _canonical_sha256(unsigned):
        raise PackageArtifactError("package set hash does not match its contents")
    registries = value.get("registries")
    if not isinstance(registries, dict) or set(registries) != {"python", "npm"}:
        raise PackageArtifactError("package set registries are invalid")
    for registry in registries.values():
        parsed = urlsplit(str(registry))
        if parsed.scheme != "https" or not parsed.netloc or parsed.username or parsed.password:
            raise PackageArtifactError("package registries must use credential-free HTTPS URLs")
    for group in ("python", "webui"):
        packages = value.get(group)
        if not isinstance(packages, list) or not packages:
            raise PackageArtifactError(f"package set {group} entries are missing")
        _validate_package_entries(group, packages)
    return value


def _validate_package_entries(group: str, packages: list[object]) -> None:
    names: set[str] = set()
    for raw in packages:
        if not isinstance(raw, dict):
            raise PackageArtifactError(f"package set {group} entry is malformed")
        name = raw.get("name")
        version = raw.get("version")
        repository = raw.get("repository")
        tag = raw.get("tag")
        commit = raw.get("commit")
        name_valid = (
            isinstance(name, str)
            and (NAME.fullmatch(name) if group == "python" else WEBUI_NAME.fullmatch(name))
        )
        if (
            not name_valid
            or name in names
            or not isinstance(version, str)
            or VERSION.fullmatch(version) is None
            or not isinstance(repository, str)
            or NAME.fullmatch(repository) is None
            or tag != f"v{version}"
            or not isinstance(commit, str)
            or COMMIT.fullmatch(commit) is None
        ):
            raise PackageArtifactError(f"package set {group} entry has an invalid identity")
        names.add(name)
        if group == "python":
            extras = raw.get("extras")
            if not isinstance(extras, list) or any(
                not isinstance(item, str)
                or re.fullmatch(r"[a-z][a-z0-9_-]*", item) is None
                for item in extras
            ):
                raise PackageArtifactError("package set Python extras are invalid")


def _download_wheels(
    *, packages: tuple[dict[str, object], ...], destination: Path, python: str, index_url: str
) -> None:
    requirements = [_python_requirement(item) for item in packages]
    environment = dict(os.environ)
    environment["PIP_INDEX_URL"] = _authenticated_url(index_url)
    environment["PIP_EXTRA_INDEX_URL"] = ""
    environment["PIP_CONFIG_FILE"] = os.devnull
    environment["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
    subprocess.run(
        [python, "-m", "pip", "download", "--no-deps", "--only-binary=:all:", "--dest", str(destination), *requirements],
        check=True,
        env=environment,
    )


def _download_webui(
    *, packages: tuple[dict[str, object], ...], destination: Path, npm: str, registry: str
) -> None:
    environment = dict(os.environ)
    npmrc: tempfile.NamedTemporaryFile[bytes] | None = None
    token = os.environ.get("GOVOPLAN_PACKAGE_TOKEN", "")
    if token:
        parsed = urlsplit(registry)
        auth_path = f"//{parsed.netloc}{parsed.path}:_authToken={token}\n"
        npmrc = tempfile.NamedTemporaryFile(prefix="govoplan-npmrc-", delete=False)
        npmrc.write(f"@govoplan:registry={registry}\n{auth_path}".encode("utf-8"))
        npmrc.close()
        os.chmod(npmrc.name, 0o600)
        environment["NPM_CONFIG_USERCONFIG"] = npmrc.name
    try:
        for item in packages:
            subprocess.run(
                [npm, "pack", f"{item['name']}@{item['version']}", "--ignore-scripts", "--pack-destination", str(destination), "--registry", registry],
                check=True,
                env=environment,
            )
    finally:
        if npmrc is not None:
            Path(npmrc.name).unlink(missing_ok=True)


def _verify_wheels(packages: tuple[dict[str, object], ...], root: Path) -> list[dict[str, object]]:
    expected = {_normalize(str(item["name"])): item for item in packages}
    rows: list[dict[str, object]] = []
    seen: set[str] = set()
    for path in sorted(root.glob("*.whl")):
        identity = _wheel_identity(path)
        name = str(identity["name"])
        package = expected.get(name)
        if package is None or identity["version"] != package["version"] or name in seen:
            raise PackageArtifactError(f"unexpected wheel artifact: {path.name}")
        seen.add(name)
        rows.append(_artifact_row(path, package))
    if seen != set(expected):
        raise PackageArtifactError("registry did not return every selected Python wheel")
    return sorted(rows, key=lambda item: str(item["name"]))


def _verify_webui(packages: tuple[dict[str, object], ...], root: Path) -> list[dict[str, object]]:
    expected = {str(item["name"]): item for item in packages}
    rows: list[dict[str, object]] = []
    seen: set[str] = set()
    for path in sorted(root.glob("*.tgz")):
        identity = _npm_identity(path)
        name = str(identity["name"])
        package = expected.get(name)
        if package is None or identity["version"] != package["version"] or name in seen:
            raise PackageArtifactError(f"unexpected WebUI artifact: {path.name}")
        seen.add(name)
        row = _artifact_row(path, package)
        row["integrity"] = "sha512-" + base64.b64encode(hashlib.sha512(path.read_bytes()).digest()).decode("ascii")
        rows.append(row)
    if seen != set(expected):
        raise PackageArtifactError("registry did not return every selected WebUI package")
    return sorted(rows, key=lambda item: str(item["name"]))


def _wheel_identity(path: Path) -> dict[str, str]:
    _bounded(path)
    with zipfile.ZipFile(path) as archive:
        metadata = [
            item for item in archive.infolist()
            if PurePosixPath(item.filename).name == "METADATA"
            and PurePosixPath(item.filename).parent.name.endswith(".dist-info")
        ]
        if len(metadata) != 1 or metadata[0].file_size > 1024 * 1024:
            raise PackageArtifactError(f"wheel metadata is invalid: {path.name}")
        parsed = BytesParser(policy=compat32).parsebytes(archive.read(metadata[0]))
    name = _normalize(str(parsed.get("Name") or ""))
    version = str(parsed.get("Version") or "")
    if NAME.fullmatch(name) is None or VERSION.fullmatch(version) is None:
        raise PackageArtifactError(f"wheel identity is invalid: {path.name}")
    return {"name": name, "version": version}


def _npm_identity(path: Path) -> dict[str, str]:
    _bounded(path)
    with tarfile.open(path, mode="r:gz") as archive:
        try:
            member = archive.getmember("package/package.json")
        except KeyError as exc:
            raise PackageArtifactError(f"npm package metadata is missing: {path.name}") from exc
        if not member.isfile() or member.size > 1024 * 1024:
            raise PackageArtifactError(f"npm package metadata is invalid: {path.name}")
        extracted = archive.extractfile(member)
        if extracted is None:
            raise PackageArtifactError(f"npm package metadata cannot be read: {path.name}")
        value = json.load(extracted)
    name = value.get("name")
    version = value.get("version")
    if not isinstance(name, str) or not name.startswith("@govoplan/") or not isinstance(version, str) or VERSION.fullmatch(version) is None:
        raise PackageArtifactError(f"npm package identity is invalid: {path.name}")
    return {"name": name, "version": version}


def _artifact_row(path: Path, package: dict[str, object]) -> dict[str, object]:
    row = {
        "name": package["name"],
        "version": package["version"],
        "repository": package["repository"],
        "tag": package["tag"],
        "commit": package["commit"],
        "filename": path.name,
        "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
        "size": path.stat().st_size,
    }
    if "extras" in package:
        row["extras"] = package["extras"]
    return row


def _write_requirements(path: Path, rows: list[dict[str, object]]) -> None:
    lines = ["--no-index", "--find-links ./local-wheels", "--require-hashes"]
    for row in rows:
        selected_extras = row.get("extras") or []
        extras = (
            f"[{','.join(str(value) for value in selected_extras)}]"
            if selected_extras
            else ""
        )
        lines.append(
            f"{row['name']}{extras}=={row['version']} "
            f"--hash=sha256:{row['sha256']}"
        )
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _python_requirement(item: dict[str, object]) -> str:
    extras = item.get("extras") or []
    suffix = f"[{','.join(str(value) for value in extras)}]" if extras else ""
    return f"{item['name']}{suffix}=={item['version']}"


def _authenticated_url(url: str) -> str:
    token = os.environ.get("GOVOPLAN_PACKAGE_TOKEN", "")
    username = os.environ.get("GOVOPLAN_PACKAGE_USERNAME", "")
    if not token:
        return url
    if not username:
        raise PackageArtifactError("GOVOPLAN_PACKAGE_USERNAME is required with a package token")
    parsed = urlsplit(url)
    return urlunsplit(
        (
            parsed.scheme,
            f"{quote(username, safe='')}:{quote(token, safe='')}@{parsed.netloc}",
            parsed.path,
            parsed.query,
            "",
        )
    )


def _require_empty_destination(path: Path) -> None:
    if path.exists() and (not path.is_dir() or any(path.iterdir())):
        raise PackageArtifactError(f"output directory must be absent or empty: {path}")
    if path.is_symlink():
        raise PackageArtifactError(f"output directory must not be a symlink: {path}")


def _bounded(path: Path) -> None:
    if path.is_symlink() or not path.is_file() or path.stat().st_size > MAX_ARTIFACT_BYTES:
        raise PackageArtifactError(f"package artifact is invalid or too large: {path.name}")


def _normalize(value: str) -> str:
    return re.sub(r"[-_.]+", "-", value.strip().lower())


def _canonical_sha256(value: object) -> str:
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest()


def main() -> int:
    args = build_parser().parse_args()
    try:
        lock = resolve(args)
    except (PackageArtifactError, OSError, ValueError, subprocess.CalledProcessError, zipfile.BadZipFile, tarfile.TarError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    print(f"Resolved {len(lock['python'])} Python and {len(lock['webui'])} WebUI packages.")
    print(f"Package artifact lock written to {args.lock_output}")
    return 0


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