#!/usr/bin/env python3
"""Assemble a deterministic, network-free GovOPlaN OCI build context."""

from __future__ import annotations

import argparse
import configparser
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 stat
import zipfile


NORMALIZED_PACKAGE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
VERSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._+!-]{0,127}$")
MAX_WHEELS = 512
MAX_WHEEL_BYTES = 512 * 1024 * 1024
MAX_WEB_FILES = 100_000
MAX_WEB_BYTES = 2 * 1024 * 1024 * 1024


class ContextError(ValueError):
    """The release inputs cannot form an immutable runtime context."""


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--wheelhouse", type=Path, required=True)
    parser.add_argument("--web-dist", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--required-module", action="append", default=[])
    parser.add_argument(
        "--source-date-epoch",
        type=int,
        default=int(os.environ.get("SOURCE_DATE_EPOCH", "0") or 0),
    )
    return parser


def prepare_context(
    *,
    wheelhouse: Path,
    web_dist: Path,
    output: Path,
    required_modules: tuple[str, ...] = (),
    source_date_epoch: int = 0,
) -> dict[str, object]:
    source_wheels = _regular_files(wheelhouse, suffix=".whl", maximum=MAX_WHEELS)
    if not source_wheels:
        raise ContextError("wheelhouse contains no wheel artifacts")
    if output.exists() and any(output.iterdir()):
        raise ContextError("output directory must be absent or empty")
    output.mkdir(parents=True, exist_ok=True)
    target_wheels = output / "wheelhouse"
    target_web = output / "web-dist"
    target_wheels.mkdir(mode=0o755)

    packages: list[dict[str, object]] = []
    govoplan_wheel_rows: list[dict[str, object]] = []
    roots: list[tuple[str, str]] = []
    module_ids: set[str] = set()
    seen_packages: set[str] = set()
    wheel_rows: list[dict[str, object]] = []
    for source in source_wheels:
        if source.stat().st_size > MAX_WHEEL_BYTES:
            raise ContextError(f"wheel exceeds size limit: {source.name}")
        identity = inspect_wheel(source)
        package_name = str(identity["package"])
        if package_name in seen_packages:
            raise ContextError(f"duplicate wheel distribution: {package_name}")
        seen_packages.add(package_name)
        target = target_wheels / source.name
        _copy_regular(source, target, source_date_epoch=source_date_epoch)
        row = {
            "filename": source.name,
            "sha256": _sha256_file(target),
            "size": target.stat().st_size,
        }
        wheel_rows.append(row)
        if package_name.startswith("govoplan-"):
            package_modules = tuple(str(item) for item in identity["module_ids"])
            module_ids.update(package_modules)
            package = {
                **row,
                "package": package_name,
                "version": identity["version"],
                "module_ids": list(package_modules),
            }
            packages.append(package)
            govoplan_wheel_rows.append(row)
            root = (
                f"{package_name}[server]"
                if package_name == "govoplan-core"
                else package_name
            )
            roots.append((root, str(identity["version"])))

    if not any(package["package"] == "govoplan-core" for package in packages):
        raise ContextError("wheelhouse does not contain govoplan-core")
    missing_modules = sorted(set(required_modules) - module_ids)
    if missing_modules:
        raise ContextError(
            "runtime composition is missing required modules: "
            + ", ".join(missing_modules)
        )
    requirements = "".join(
        f"{package}=={version}\n" for package, version in sorted(roots)
    )
    _write_regular(
        output / "requirements-runtime.txt",
        requirements.encode("utf-8"),
        source_date_epoch=source_date_epoch,
    )

    web_rows = _copy_web_tree(
        web_dist,
        target_web,
        source_date_epoch=source_date_epoch,
    )
    composition: dict[str, object] = {
        "schema_version": "1",
        "python": {
            "packages": sorted(packages, key=lambda item: str(item["package"])),
            "module_ids": sorted(module_ids),
            "wheelhouse_sha256": _rows_digest(govoplan_wheel_rows),
            "wheel_count": len(govoplan_wheel_rows),
        },
        "web": {
            "sha256": _rows_digest(web_rows),
            "file_count": len(web_rows),
        },
    }
    encoded = (json.dumps(composition, indent=2, sort_keys=True) + "\n").encode(
        "utf-8"
    )
    _write_regular(
        output / "composition.json",
        encoded,
        source_date_epoch=source_date_epoch,
    )
    _write_regular(
        target_web / ".well-known" / "govoplan-composition.json",
        encoded,
        source_date_epoch=source_date_epoch,
    )
    _copy_regular(
        Path(__file__).resolve().parent / "runtime" / "nginx.conf",
        output / "nginx.conf",
        source_date_epoch=source_date_epoch,
    )
    _copy_regular(
        Path(__file__).resolve().parent / "runtime" / "web-entrypoint.sh",
        output / "web-entrypoint.sh",
        source_date_epoch=source_date_epoch,
    )
    return composition


def inspect_wheel(path: Path) -> dict[str, object]:
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        descriptor = os.open(path, flags)
    except OSError as exc:
        raise ContextError(f"wheel cannot be opened safely: {path.name}") from exc
    try:
        opened = os.fstat(descriptor)
        if not stat.S_ISREG(opened.st_mode):
            raise ContextError(f"wheel is not a regular file: {path.name}")
        with os.fdopen(os.dup(descriptor), "rb") as handle:
            with zipfile.ZipFile(handle) as archive:
                metadata = [
                    member
                    for member in archive.infolist()
                    if PurePosixPath(member.filename).name == "METADATA"
                    and PurePosixPath(member.filename).parent.name.endswith(
                        ".dist-info"
                    )
                ]
                if len(metadata) != 1:
                    raise ContextError(
                        f"wheel must contain one METADATA file: {path.name}"
                    )
                parsed = BytesParser(policy=compat32).parsebytes(
                    archive.read(metadata[0])
                )
                package = _normalize_package(str(parsed.get("Name") or ""))
                version = str(parsed.get("Version") or "").strip()
                if VERSION.fullmatch(version) is None:
                    raise ContextError(f"wheel has invalid version: {path.name}")
                entry_points_name = (
                    PurePosixPath(metadata[0].filename).parent / "entry_points.txt"
                ).as_posix()
                module_ids: tuple[str, ...] = ()
                if entry_points_name in archive.namelist():
                    module_ids = _module_entry_points(
                        archive.read(entry_points_name).decode("utf-8")
                    )
        final = os.fstat(descriptor)
        if (opened.st_dev, opened.st_ino, opened.st_size, opened.st_mtime_ns) != (
            final.st_dev,
            final.st_ino,
            final.st_size,
            final.st_mtime_ns,
        ):
            raise ContextError(f"wheel changed while inspected: {path.name}")
    except (OSError, RuntimeError, zipfile.BadZipFile) as exc:
        if isinstance(exc, ContextError):
            raise
        raise ContextError(f"wheel is not a readable archive: {path.name}") from exc
    finally:
        os.close(descriptor)
    return {"package": package, "version": version, "module_ids": module_ids}


def _module_entry_points(value: str) -> tuple[str, ...]:
    parser = configparser.ConfigParser(interpolation=None, strict=True)
    try:
        parser.read_string(value)
    except configparser.Error as exc:
        raise ContextError("wheel entry_points.txt is malformed") from exc
    if not parser.has_section("govoplan.modules"):
        return ()
    values = tuple(sorted(parser.options("govoplan.modules")))
    for item in values:
        if re.fullmatch(r"[a-z][a-z0-9_]{1,63}", item) is None:
            raise ContextError(f"wheel has invalid module entry point: {item!r}")
    return values


def _normalize_package(value: str) -> str:
    normalized = re.sub(r"[-_.]+", "-", value.strip().lower())
    if NORMALIZED_PACKAGE.fullmatch(normalized) is None:
        raise ContextError("wheel has invalid package name")
    return normalized


def _regular_files(root: Path, *, suffix: str, maximum: int) -> list[Path]:
    if root.is_symlink() or not root.is_dir():
        raise ContextError(f"input directory is not a real directory: {root}")
    values = sorted(path for path in root.iterdir() if path.name.endswith(suffix))
    if len(values) > maximum:
        raise ContextError(f"input directory exceeds {maximum} files")
    for path in values:
        if path.is_symlink() or not path.is_file():
            raise ContextError(f"input artifact is not a regular file: {path.name}")
    return values


def _copy_web_tree(
    source: Path,
    target: Path,
    *,
    source_date_epoch: int,
) -> list[dict[str, object]]:
    if source.is_symlink() or not source.is_dir():
        raise ContextError("WebUI dist must be a real directory")
    rows: list[dict[str, object]] = []
    total = 0
    for path in sorted(source.rglob("*")):
        relative = path.relative_to(source)
        if path.is_symlink():
            raise ContextError(f"WebUI dist contains a symlink: {relative}")
        if path.is_dir():
            continue
        if not path.is_file():
            raise ContextError(f"WebUI dist contains a special file: {relative}")
        if len(rows) >= MAX_WEB_FILES:
            raise ContextError("WebUI dist exceeds its file-count limit")
        total += path.stat().st_size
        if total > MAX_WEB_BYTES:
            raise ContextError("WebUI dist exceeds its total-size limit")
        destination = target / relative
        _copy_regular(path, destination, source_date_epoch=source_date_epoch)
        rows.append(
            {
                "path": relative.as_posix(),
                "sha256": _sha256_file(destination),
                "size": destination.stat().st_size,
            }
        )
    if not rows:
        raise ContextError("WebUI dist contains no files")
    return rows


def _copy_regular(source: Path, target: Path, *, source_date_epoch: int) -> None:
    target.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
    with source.open("rb") as source_handle, target.open("xb") as target_handle:
        shutil.copyfileobj(source_handle, target_handle)
        target_handle.flush()
        os.fsync(target_handle.fileno())
    target.chmod(0o644)
    os.utime(target, (source_date_epoch, source_date_epoch))


def _write_regular(path: Path, value: bytes, *, source_date_epoch: int) -> None:
    path.parent.mkdir(mode=0o755, parents=True, exist_ok=True)
    path.write_bytes(value)
    path.chmod(0o644)
    os.utime(path, (source_date_epoch, source_date_epoch))


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


def _rows_digest(rows: list[dict[str, object]]) -> str:
    encoded = json.dumps(rows, separators=(",", ":"), sort_keys=True).encode(
        "utf-8"
    )
    return hashlib.sha256(encoded).hexdigest()


def main() -> int:
    args = build_parser().parse_args()
    try:
        composition = prepare_context(
            wheelhouse=args.wheelhouse.expanduser().resolve(),
            web_dist=args.web_dist.expanduser().resolve(),
            output=args.output.expanduser().resolve(),
            required_modules=tuple(args.required_module),
            source_date_epoch=args.source_date_epoch,
        )
    except (ContextError, OSError) as exc:
        print(f"error: {exc}", file=os.sys.stderr)
        return 1
    print(json.dumps(composition, indent=2, sort_keys=True))
    return 0


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