#!/usr/bin/env python3
"""Refresh Galaxy's checked-in Amazon Bedrock model catalog snapshot."""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request
from pathlib import Path


UPSTREAM_GIT_URL = (
    "https://github.com/amazonbedrockmodels/amazonbedrockmodels.github.io.git"
)
UPSTREAM_RAW_URL = (
    "https://raw.githubusercontent.com/amazonbedrockmodels/"
    "amazonbedrockmodels.github.io"
)
UPSTREAM_REF = "refs/heads/main"
CATALOG_FILES = {
    "api.json": dict,
    "beta_models.json": list,
    "mantle_models.json": dict,
    "model_cards.json": dict,
    "models.json": list,
    "profiles.json": list,
}


def resolve_upstream_commit() -> str:
    result = subprocess.run(
        ["git", "ls-remote", UPSTREAM_GIT_URL, UPSTREAM_REF],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    fields = result.stdout.split()
    if len(fields) != 2 or fields[1] != UPSTREAM_REF or len(fields[0]) != 40:
        raise RuntimeError(f"unexpected git ls-remote response: {result.stdout!r}")
    return fields[0]


def download_file(commit: str, filename: str) -> bytes:
    url = f"{UPSTREAM_RAW_URL}/{commit}/data/{filename}"
    request = urllib.request.Request(url, headers={"User-Agent": "Galaxy catalog updater"})
    with urllib.request.urlopen(request, timeout=30) as response:
        if response.status != 200:
            raise RuntimeError(f"downloading {filename} returned HTTP {response.status}")
        return response.read()


def validate_document(filename: str, contents: bytes, expected_type: type) -> None:
    try:
        document = json.loads(contents)
    except (UnicodeDecodeError, json.JSONDecodeError) as error:
        raise RuntimeError(f"{filename} is not valid UTF-8 JSON: {error}") from error

    if not isinstance(document, expected_type):
        raise RuntimeError(
            f"{filename} has type {type(document).__name__}; "
            f"expected {expected_type.__name__}"
        )


def main() -> int:
    workspace_root = Path(__file__).resolve().parent.parent
    destination = workspace_root / "crates" / "galaxy_bedrock_model_catalog" / "data"
    commit_file = destination / "UPSTREAM_COMMIT"

    print("Resolving the latest Amazon Bedrock model catalog commit...")
    commit = resolve_upstream_commit()

    if commit_file.exists() and commit_file.read_text().strip() == commit:
        print(f"Catalog is already current at {commit}")
        return 0

    with tempfile.TemporaryDirectory(prefix="galaxy-bedrock-catalog-") as temp_dir:
        staged_directory = Path(temp_dir)
        for filename, expected_type in CATALOG_FILES.items():
            print(f"Downloading {filename}...")
            contents = download_file(commit, filename)
            validate_document(filename, contents, expected_type)
            (staged_directory / filename).write_bytes(contents)

        destination.mkdir(parents=True, exist_ok=True)
        for filename in CATALOG_FILES:
            (staged_directory / filename).replace(destination / filename)
        commit_file.write_text(f"{commit}\n")

    print(f"Updated Galaxy's Bedrock model catalog to {commit}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, RuntimeError, subprocess.SubprocessError, urllib.error.URLError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(1) from error