#!/usr/bin/env python3
"""Assemble a fail-closed multi-arch manifest list from two per-arch leaves.

Kaniko builds one native image per architecture (arm64 on an rpi5 pod, amd64 on
titan-24) and pushes each under an arch-suffixed candidate tag
``...-build-<N>-<arch>``.  Kaniko cannot combine, so this step:

1. Independently re-reads each per-arch candidate manifest from the registry and
   binds it to the exact Kaniko digest evidence.
2. Proves each leaf really is the architecture it claims by reading its image
   config (a swapped or cross-built leaf fails closed here).
3. Builds a Docker manifest *list* (not an OCI index) from the two verified
   leaves -- Docker manifest lists are already inside the promotion allow-list,
   so this keeps the security surface of ``hermes_oci_promote.py`` unchanged.
4. Refuses to overwrite an existing final tag, PUTs the list to the final
   ``...-build-<N>`` tag, and re-reads it to confirm the registry resolved the
   exact index digest referencing exactly the two expected leaves.

The output ``--digest-file``/``--image-file`` deliberately use the SAME format
the single-arch Kaniko step produced (``<digest>`` and ``<destination>@<digest>``
for the final, arch-less tag).  The whole downstream evidence chain --
``hermes_image_release.py`` render/verify-evidence and ``hermes_oci_promote.py``
-- therefore promotes the multi-arch INDEX with no further change.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import os
import re
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any, Callable


REGISTRY_ORIGIN = "https://registry.bstein.dev"
# The final (arch-less) Flux-visible tag; identical contract to the promoter.
# Both Hermes images that the multi-arch pipelines publish share the exact same
# arch-less final-tag contract; the repository name is the only difference and is
# captured here so the combiner stays fail-closed to just these two components.
DESTINATION_PATTERN = re.compile(
    r"^registry\.bstein\.dev/bstein/(?P<component>hermes-agent|hermes-webui):"
    r"git-(?P<revision>[0-9a-f]{40})-build-(?P<build>[1-9][0-9]*)$"
)
DIGEST_PATTERN = re.compile(r"^sha256:[0-9a-f]{64}$")
DOCKER_MANIFEST_LIST = "application/vnd.docker.distribution.manifest.list.v2+json"
# A per-arch leaf must be a single-image manifest, never itself a list/index.
LEAF_MANIFEST_TYPES = {
    "application/vnd.docker.distribution.manifest.v2+json",
    "application/vnd.oci.image.manifest.v1+json",
}
IMAGE_CONFIG_TYPES = {
    "application/vnd.docker.container.image.v1+json",
    "application/vnd.oci.image.config.v1+json",
}
# Deterministic architecture order -> deterministic manifest-list bytes/digest.
ARCHITECTURES = ("amd64", "arm64")
MAX_MANIFEST_BYTES = 4 * 1024 * 1024
MAX_CONFIG_BYTES = 1024 * 1024


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    """Never forward registry credentials to another origin."""

    def redirect_request(self, _request, _file, _code, _message, _headers, _url):
        return None


def _registry_request(request: urllib.request.Request, timeout: int) -> Any:
    """Return normal and HTTP error responses without following redirects."""
    opener = urllib.request.build_opener(_NoRedirect())
    try:
        return opener.open(request, timeout=timeout)
    except urllib.error.HTTPError as exc:
        return exc


def _status(response: Any) -> int:
    """Normalize urllib response and HTTPError status fields."""
    return int(getattr(response, "status", getattr(response, "code", 0)))


def _authorization(username: str, password: str) -> str:
    """Build a Basic authorization value without placing it in a URL."""
    if not username or not password:
        raise RuntimeError("Harbor credentials are unavailable")
    encoded = base64.b64encode(f"{username}:{password}".encode()).decode("ascii")
    return f"Basic {encoded}"


def _manifest_url(component: str, reference: str) -> str:
    """Return one same-origin, path-escaped Docker Registry manifest URL."""
    encoded = urllib.parse.quote(reference, safe="")
    return f"{REGISTRY_ORIGIN}/v2/bstein/{component}/manifests/{encoded}"


def _blob_url(component: str, digest: str) -> str:
    """Return one same-origin, path-escaped Docker Registry blob URL."""
    encoded = urllib.parse.quote(digest, safe="")
    return f"{REGISTRY_ORIGIN}/v2/bstein/{component}/blobs/{encoded}"


def _read_evidence_pair(
    *, digest_text: str, image_text: str, per_arch_tag_ref: str
) -> str:
    """Cross-check both Kaniko output files for one arch against its tag."""
    digest_lines = digest_text.splitlines()
    image_lines = image_text.splitlines()
    if len(digest_lines) != 1:
        raise ValueError("Kaniko digest evidence must contain exactly one line")
    if len(image_lines) != 1:
        raise ValueError("Kaniko image evidence must contain exactly one line")
    digest = digest_lines[0].strip()
    if not DIGEST_PATTERN.fullmatch(digest):
        raise ValueError("invalid per-arch image digest")
    if image_lines[0].strip() != f"{per_arch_tag_ref}@{digest}":
        raise ValueError("per-arch image evidence does not match tag and digest")
    return digest


def _verified_leaf(
    *,
    component: str,
    per_arch_tag: str,
    architecture: str,
    expected_digest: str,
    authorization: str,
    opener: Callable[[urllib.request.Request, int], Any],
) -> dict[str, Any]:
    """Re-read one per-arch leaf and prove its digest, type, and architecture."""
    accept = ", ".join(sorted(LEAF_MANIFEST_TYPES))
    request = urllib.request.Request(
        _manifest_url(component, per_arch_tag),
        headers={"Accept": accept, "Authorization": authorization},
        method="GET",
    )
    with opener(request, 30) as response:
        if _status(response) != 200:
            raise RuntimeError(
                f"{architecture} leaf manifest returned HTTP {_status(response)}"
            )
        body = response.read(MAX_MANIFEST_BYTES + 1)
        if len(body) > MAX_MANIFEST_BYTES:
            raise RuntimeError(f"{architecture} leaf manifest exceeded the size limit")
        observed_digest = response.headers.get("Docker-Content-Digest", "")
        content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
    if observed_digest != expected_digest:
        raise RuntimeError(f"{architecture} leaf digest does not match build evidence")
    # Defence in depth: the digest header is registry-asserted; recompute it too.
    if f"sha256:{hashlib.sha256(body).hexdigest()}" != expected_digest:
        raise RuntimeError(f"{architecture} leaf bytes do not hash to its digest")
    if content_type not in LEAF_MANIFEST_TYPES:
        raise RuntimeError(f"{architecture} leaf is not a single-image manifest")
    try:
        manifest = json.loads(body.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RuntimeError(f"{architecture} leaf manifest is not valid JSON") from exc
    config = manifest.get("config")
    if not isinstance(config, dict):
        raise RuntimeError(f"{architecture} leaf manifest omits its config descriptor")
    config_digest = str(config.get("digest") or "")
    config_type = str(config.get("mediaType") or "")
    if not DIGEST_PATTERN.fullmatch(config_digest):
        raise RuntimeError(f"{architecture} leaf config digest is invalid")
    if config_type not in IMAGE_CONFIG_TYPES:
        raise RuntimeError(f"{architecture} leaf config media type is unsupported")
    config_request = urllib.request.Request(
        _blob_url(component, config_digest),
        headers={"Accept": config_type, "Authorization": authorization},
        method="GET",
    )
    with opener(config_request, 30) as response:
        if _status(response) != 200:
            raise RuntimeError(
                f"{architecture} leaf config returned HTTP {_status(response)}"
            )
        config_body = response.read(MAX_CONFIG_BYTES + 1)
        if len(config_body) > MAX_CONFIG_BYTES:
            raise RuntimeError(f"{architecture} leaf config exceeded the size limit")
    if f"sha256:{hashlib.sha256(config_body).hexdigest()}" != config_digest:
        raise RuntimeError(f"{architecture} leaf config bytes do not hash to its digest")
    try:
        config_json = json.loads(config_body.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RuntimeError(f"{architecture} leaf config is not valid JSON") from exc
    if config_json.get("architecture") != architecture:
        raise RuntimeError(
            f"{architecture} leaf config reports architecture "
            f"{config_json.get('architecture')!r}"
        )
    if config_json.get("os") != "linux":
        raise RuntimeError(f"{architecture} leaf config reports a non-linux os")
    return {
        "mediaType": content_type,
        "size": len(body),
        "digest": expected_digest,
        "platform": {"architecture": architecture, "os": "linux"},
    }


def _manifest_list_bytes(descriptors: list[dict[str, Any]]) -> bytes:
    """Serialize the manifest list deterministically for a stable index digest."""
    document = {
        "schemaVersion": 2,
        "mediaType": DOCKER_MANIFEST_LIST,
        "manifests": descriptors,
    }
    return json.dumps(document, sort_keys=True, separators=(",", ":")).encode("utf-8")


def combine_multiarch_index(
    *,
    destination: str,
    arch_digests: dict[str, str],
    username: str,
    password: str,
    opener: Callable[[urllib.request.Request, int], Any] = _registry_request,
) -> dict[str, str]:
    """Verify both leaves, publish, and re-verify one multi-arch index tag."""
    match = DESTINATION_PATTERN.fullmatch(destination.strip())
    if not match:
        raise ValueError("invalid multi-arch destination")
    if set(arch_digests) != set(ARCHITECTURES):
        raise ValueError("expected exactly the arm64 and amd64 per-arch digests")
    component = match.group("component")
    index_tag = destination.rsplit(":", 1)[1]
    authorization = _authorization(username, password)

    descriptors = [
        _verified_leaf(
            component=component,
            per_arch_tag=f"{index_tag}-{architecture}",
            architecture=architecture,
            expected_digest=arch_digests[architecture],
            authorization=authorization,
            opener=opener,
        )
        for architecture in ARCHITECTURES
    ]
    manifest_list = _manifest_list_bytes(descriptors)
    if len(manifest_list) > MAX_MANIFEST_BYTES:
        raise RuntimeError("assembled manifest list exceeded the size limit")
    index_digest = f"sha256:{hashlib.sha256(manifest_list).hexdigest()}"

    index_url = _manifest_url(component, index_tag)
    head_request = urllib.request.Request(
        index_url,
        headers={"Accept": DOCKER_MANIFEST_LIST, "Authorization": authorization},
        method="HEAD",
    )
    with opener(head_request, 20) as response:
        head_status = _status(response)
        existing_digest = response.headers.get("Docker-Content-Digest", "")
    if head_status == 200:
        if existing_digest != index_digest:
            raise RuntimeError("final tag already exists with another index digest")
        result = "already-present"
    elif head_status == 404:
        put_request = urllib.request.Request(
            index_url,
            data=manifest_list,
            headers={
                "Authorization": authorization,
                "Content-Type": DOCKER_MANIFEST_LIST,
            },
            method="PUT",
        )
        with opener(put_request, 30) as response:
            put_status = _status(response)
            put_digest = response.headers.get("Docker-Content-Digest", "")
        if put_status not in {201, 202}:
            raise RuntimeError(f"index manifest returned HTTP {put_status}")
        if put_digest and put_digest != index_digest:
            raise RuntimeError("index manifest digest changed during publish")
        result = "published"
    else:
        raise RuntimeError(f"final tag preflight returned HTTP {head_status}")

    verify_request = urllib.request.Request(
        index_url,
        headers={"Accept": DOCKER_MANIFEST_LIST, "Authorization": authorization},
        method="GET",
    )
    with opener(verify_request, 30) as response:
        if _status(response) != 200:
            raise RuntimeError(f"index verification returned HTTP {_status(response)}")
        verify_body = response.read(MAX_MANIFEST_BYTES + 1)
        if len(verify_body) > MAX_MANIFEST_BYTES:
            raise RuntimeError("index verification exceeded the size limit")
        verify_digest = response.headers.get("Docker-Content-Digest", "")
        verify_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
    if verify_digest != index_digest:
        raise RuntimeError("registry resolved the final tag to another index digest")
    if verify_type != DOCKER_MANIFEST_LIST:
        raise RuntimeError("registry did not store a Docker manifest list")
    try:
        published = json.loads(verify_body.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise RuntimeError("registry returned invalid index JSON") from exc
    published_leaves = {
        (
            str((item or {}).get("platform", {}).get("architecture")),
            str((item or {}).get("digest")),
        )
        for item in published.get("manifests") or []
    }
    expected_leaves = {
        (architecture, arch_digests[architecture]) for architecture in ARCHITECTURES
    }
    if published_leaves != expected_leaves:
        raise RuntimeError("published index does not reference the exact two leaves")

    return {
        "component": component,
        "index_digest": index_digest,
        "index_tag": index_tag,
        "result": result,
        **{f"{architecture}_digest": arch_digests[architecture] for architecture in ARCHITECTURES},
    }


def _load_arch_digest(
    *, destination: str, architecture: str, digest_file: Path, image_file: Path
) -> str:
    """Bind one arch's two Kaniko evidence files to its arch-suffixed tag."""
    per_arch_tag_ref = f"{destination}-{architecture}"
    return _read_evidence_pair(
        digest_text=digest_file.read_text(encoding="utf-8"),
        image_text=image_file.read_text(encoding="utf-8"),
        per_arch_tag_ref=per_arch_tag_ref,
    )


def main() -> int:
    """Combine two verified per-arch leaves and emit index digest evidence."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--destination", required=True)
    parser.add_argument("--source-revision", required=True)
    parser.add_argument("--build-number", required=True)
    parser.add_argument("--arm64-digest-file", required=True, type=Path)
    parser.add_argument("--arm64-image-file", required=True, type=Path)
    parser.add_argument("--amd64-digest-file", required=True, type=Path)
    parser.add_argument("--amd64-image-file", required=True, type=Path)
    parser.add_argument("--digest-file", required=True, type=Path)
    parser.add_argument("--image-file", required=True, type=Path)
    args = parser.parse_args()
    try:
        match = DESTINATION_PATTERN.fullmatch(args.destination.strip())
        if not match:
            raise ValueError("invalid multi-arch destination")
        if match.group("revision") != args.source_revision.strip():
            raise ValueError("destination revision does not match evidence")
        if match.group("build") != args.build_number.strip():
            raise ValueError("destination build number does not match evidence")
        destination = args.destination.strip()
        arch_digests = {
            "arm64": _load_arch_digest(
                destination=destination,
                architecture="arm64",
                digest_file=args.arm64_digest_file,
                image_file=args.arm64_image_file,
            ),
            "amd64": _load_arch_digest(
                destination=destination,
                architecture="amd64",
                digest_file=args.amd64_digest_file,
                image_file=args.amd64_image_file,
            ),
        }
        result = combine_multiarch_index(
            destination=destination,
            arch_digests=arch_digests,
            username=os.environ.get("HARBOR_USER", ""),
            password=os.environ.get("HARBOR_PASSWORD", ""),
        )
        args.digest_file.write_text(result["index_digest"] + "\n", encoding="utf-8")
        args.image_file.write_text(
            f"{destination}@{result['index_digest']}\n", encoding="utf-8"
        )
    except (OSError, ValueError, RuntimeError, urllib.error.URLError) as exc:
        print(json.dumps({"error": str(exc)}, sort_keys=True))
        return 1
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":  # pragma: no cover - exercised through main()
    raise SystemExit(main())