#!/usr/bin/env python3
"""Resolve amd64/arm64 child digests from an OCI image index."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import re
import sys


DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$")


def resolve_platforms(
    payload: object,
    *,
    repository: str,
    index_digest: str,
) -> dict[str, object]:
    if not repository or "@" in repository or any(value.isspace() for value in repository):
        raise ValueError("repository must be an unpinned OCI repository name")
    if DIGEST.fullmatch(index_digest) is None:
        raise ValueError("index digest must be sha256:<hex>")
    if not isinstance(payload, dict) or not isinstance(payload.get("manifests"), list):
        raise ValueError("OCI index must contain manifests")
    platforms: dict[str, str] = {}
    for item in payload["manifests"]:
        if not isinstance(item, dict) or not isinstance(item.get("platform"), dict):
            continue
        platform = item["platform"]
        key = f"{platform.get('os')}/{platform.get('architecture')}"
        if key not in {"linux/amd64", "linux/arm64"}:
            continue
        digest = item.get("digest")
        if not isinstance(digest, str) or DIGEST.fullmatch(digest) is None:
            raise ValueError(f"OCI index has an invalid {key} digest")
        if key in platforms:
            raise ValueError(f"OCI index has duplicate {key} manifests")
        platforms[key] = f"{repository}@{digest}"
    if set(platforms) != {"linux/amd64", "linux/arm64"}:
        raise ValueError("OCI index must contain linux/amd64 and linux/arm64")
    return {
        "index": f"{repository}@{index_digest}",
        "platforms": dict(sorted(platforms.items())),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repository", required=True)
    parser.add_argument("--index-digest", required=True)
    parser.add_argument("--index", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    try:
        payload = json.loads(args.index.read_text(encoding="utf-8"))
        result = resolve_platforms(
            payload,
            repository=args.repository,
            index_digest=args.index_digest,
        )
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(
            json.dumps(result, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    return 0


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