"""Update the HF Space to the newest official Aether image digest.

Why this exists:
- Aether's built-in updater cannot persist across HF Docker Space restarts,
  because the app binary/frontend live in the immutable Docker image layer.
- Rebuilding the Space from a new upstream image is the correct Docker update
  path on Hugging Face.

Behavior:
1. Check the latest digest for ghcr.io/fawney19/aether:latest.
2. Read the Space repo Dockerfile.
3. If the Dockerfile is not pinned to that digest, upload a new Dockerfile with
   `ARG AETHER_IMAGE=ghcr.io/fawney19/aether@sha256:...`.
4. HF automatically rebuilds the Space on that commit.
"""
import json
import os
import re
import sys
import tempfile
import urllib.request
from pathlib import Path

from huggingface_hub import HfApi, hf_hub_download

update_image = os.environ.get("AETHER_UPDATE_IMAGE", "ghcr.io/fawney19/aether:latest")
space_repo = os.environ.get("AETHER_HF_SPACE_REPO", "iiioooo1/aether-hf")
state_path = Path(os.environ.get(
    "AETHER_LAST_DIGEST_PATH", "/opt/aether/data/.last_image_digest"
))
hf_token = os.environ["HF_TOKEN"]


def parse_ghcr_image(ref: str):
    """Return (repo, reference) for ghcr.io/<repo>:<tag> or @<digest>."""
    if not ref.startswith("ghcr.io/"):
        raise SystemExit(f"unsupported registry: {ref}")
    rest = ref[len("ghcr.io/"):]
    if "@" in rest:
        repo, reference = rest.split("@", 1)
        return repo, reference
    if ":" in rest:
        repo, reference = rest.rsplit(":", 1)
    else:
        repo, reference = rest, "latest"
    return repo, reference


def fetch_ghcr_digest(repo: str, reference: str) -> str:
    token_url = f"https://ghcr.io/token?scope=repository:{repo}:pull"
    with urllib.request.urlopen(token_url, timeout=15) as r:
        bearer = json.loads(r.read())["token"]

    manifest_url = f"https://ghcr.io/v2/{repo}/manifests/{reference}"
    req = urllib.request.Request(
        manifest_url,
        headers={
            "Authorization": f"Bearer {bearer}",
            "Accept": (
                "application/vnd.oci.image.index.v1+json,"
                "application/vnd.docker.distribution.manifest.list.v2+json,"
                "application/vnd.oci.image.manifest.v1+json"
            ),
        },
    )
    with urllib.request.urlopen(req, timeout=20) as r:
        digest = r.headers.get("Docker-Content-Digest", "")
    if not digest.startswith("sha256:"):
        raise RuntimeError(f"no usable Docker-Content-Digest header for {repo}:{reference}: {digest!r}")
    return digest


def read_space_dockerfile(api_token: str, repo_id: str) -> str:
    local_path = hf_hub_download(
        repo_id=repo_id,
        repo_type="space",
        filename="Dockerfile",
        token=api_token,
        force_download=True,
    )
    return Path(local_path).read_text(encoding="utf-8")


def pin_dockerfile(content: str, pinned_ref: str) -> tuple[str, bool]:
    pattern = re.compile(r"^ARG\s+AETHER_IMAGE=.*$", re.MULTILINE)
    new_line = f"ARG AETHER_IMAGE={pinned_ref}"
    match = pattern.search(content)
    if match:
        if match.group(0).strip() == new_line:
            return content, False
        return pattern.sub(new_line, content, count=1), True
    # Fallback: insert after syntax line if the ARG is missing.
    lines = content.splitlines()
    if lines and lines[0].startswith("# syntax="):
        lines.insert(1, new_line)
    else:
        lines.insert(0, new_line)
    return "\n".join(lines) + "\n", True


def upload_dockerfile(api: HfApi, repo_id: str, content: str, digest: str) -> None:
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", delete=False) as tmp:
        tmp.write(content)
        tmp_path = tmp.name
    try:
        api.upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo="Dockerfile",
            repo_id=repo_id,
            repo_type="space",
            commit_message=f"chore: bump Aether upstream image {digest[:19]}",
            commit_description=(
                "Auto-update HF Space by pinning the official upstream image "
                f"ghcr.io/fawney19/aether to digest {digest}.\n\n"
                "HF rebuilds the Docker Space from this commit, which is the "
                "persistent Docker update path for Hugging Face Spaces."
            ),
        )
    finally:
        try:
            Path(tmp_path).unlink(missing_ok=True)
        except Exception:
            pass


repo, reference = parse_ghcr_image(update_image)
print(f"aether-auto-update: checking ghcr.io/{repo}:{reference}", file=sys.stderr)

digest = fetch_ghcr_digest(repo, reference)
pinned_ref = f"ghcr.io/{repo}@{digest}"
print(f"aether-auto-update: latest upstream digest {digest}", file=sys.stderr)

api = HfApi(token=hf_token)
dockerfile = read_space_dockerfile(hf_token, space_repo)
updated, changed = pin_dockerfile(dockerfile, pinned_ref)

if not changed:
    print("aether-auto-update: Space Dockerfile already pins latest digest; no rebuild needed", file=sys.stderr)
    state_path.parent.mkdir(parents=True, exist_ok=True)
    state_path.write_text(digest + "\n", encoding="utf-8")
    sys.exit(0)

print(f"aether-auto-update: updating Space Dockerfile to {pinned_ref}", file=sys.stderr)
upload_dockerfile(api, space_repo, updated, digest)
state_path.parent.mkdir(parents=True, exist_ok=True)
state_path.write_text(digest + "\n", encoding="utf-8")
print("aether-auto-update: Dockerfile uploaded; HF Space rebuild should start automatically", file=sys.stderr)