#!/usr/bin/env python3

"""Update pinned LLVM release asset names and SHA-256 digests in MODULE.bazel."""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
import urllib.request
from pathlib import Path
from typing import Any


LLVM_VERSION_PATTERN = re.compile(r'^LLVM_VERSION = "(?P<version>\d+\.\d+\.\d+)"$', re.MULTILINE)
LLVM_DISTRIBUTIONS_PATTERN = re.compile(
    r"^LLVM_DISTRIBUTIONS = \{\n.*?^\}$",
    re.MULTILINE | re.DOTALL,
)
LLVM_DISTRIBUTION_PLATFORMS = (
    "Linux-ARM64",
    "Linux-X64",
    "macOS-ARM64",
)


def fetch_release(version: str, token: str | None) -> dict[str, Any]:
    url = f"https://api.github.com/repos/llvm/llvm-project/releases/tags/llvmorg-{version}"
    headers = {
        "Accept": "application/vnd.github+json",
        "User-Agent": "eagle0-llvm-distribution-updater",
        "X-GitHub-Api-Version": "2022-11-28",
    }
    if token:
        headers["Authorization"] = f"Bearer {token}"

    request = urllib.request.Request(url, headers=headers)
    with urllib.request.urlopen(request, timeout=30) as response:
        return json.load(response)


def load_release(path: Path) -> dict[str, Any]:
    with path.open(encoding="utf-8") as release_file:
        return json.load(release_file)


def llvm_distributions(version: str, release: dict[str, Any]) -> dict[str, str]:
    expected_tag = f"llvmorg-{version}"
    if release.get("tag_name") != expected_tag:
        raise ValueError(
            f"Expected release tag {expected_tag}, got {release.get('tag_name')!r}"
        )

    assets = {asset["name"]: asset for asset in release.get("assets", [])}
    distributions: dict[str, str] = {}
    for platform in LLVM_DISTRIBUTION_PLATFORMS:
        name = f"LLVM-{version}-{platform}.tar.xz"
        asset = assets.get(name)
        if asset is None:
            raise ValueError(f"LLVM release {version} is missing required asset {name}")

        digest = asset.get("digest")
        if not isinstance(digest, str) or not digest.startswith("sha256:"):
            raise ValueError(f"LLVM release asset {name} has no SHA-256 digest")
        distributions[name] = digest.removeprefix("sha256:")

    return distributions


def render_distributions(distributions: dict[str, str]) -> str:
    lines = ["LLVM_DISTRIBUTIONS = {"]
    lines.extend(f'    "{name}": "{digest}",' for name, digest in distributions.items())
    lines.append("}")
    return "\n".join(lines)


def update_module(module_path: Path, release: dict[str, Any]) -> bool:
    original = module_path.read_text(encoding="utf-8")
    version_match = LLVM_VERSION_PATTERN.search(original)
    if version_match is None:
        raise ValueError(f"Could not find LLVM_VERSION in {module_path}")

    distributions = llvm_distributions(version_match.group("version"), release)
    replacement = render_distributions(distributions)
    updated, replacements = LLVM_DISTRIBUTIONS_PATTERN.subn(replacement, original, count=1)
    if replacements != 1:
        raise ValueError(f"Could not find exactly one LLVM_DISTRIBUTIONS block in {module_path}")

    if updated == original:
        return False

    module_path.write_text(updated, encoding="utf-8")
    return True


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--module", type=Path, default=Path("MODULE.bazel"))
    parser.add_argument(
        "--release-json",
        type=Path,
        help="Read GitHub release JSON from a file instead of the GitHub API",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    module_text = args.module.read_text(encoding="utf-8")
    version_match = LLVM_VERSION_PATTERN.search(module_text)
    if version_match is None:
        print(f"ERROR: Could not find LLVM_VERSION in {args.module}", file=sys.stderr)
        return 1

    version = version_match.group("version")
    try:
        release = (
            load_release(args.release_json)
            if args.release_json
            else fetch_release(version, os.environ.get("GITHUB_TOKEN"))
        )
        changed = update_module(args.module, release)
    except (OSError, ValueError) as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 1

    if changed:
        print(f"Updated LLVM {version} distribution metadata in {args.module}")
    else:
        print(f"LLVM {version} distribution metadata is current")
    return 0


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