#!/usr/bin/env python3
"""
Registry Image Tag Pruner - Keeps only the latest N SHA-tag(s) per repository.

Usage:
    python3 registry_prune.py \
        --registry registry.nxtgauge.com \
        --repo nxtgauge-rust-gateway \
        --username "$REGISTRY_USERNAME" \
        --password "$REGISTRY_PASSWORD" \
        --protect "$SHA"

Environment variables can also be used:
    REGISTRY_HOST, REGISTRY_REPO, REGISTRY_USERNAME, REGISTRY_PASSWORD

SHA-like tags are identified by pattern: ^[a-f0-9]{40}$
Non-SHA tags (e.g., high-performance-latest, main-latest, latest) are NEVER deleted.
--protect tags (e.g. the SHA this very CI run just built) are NEVER deleted,
regardless of how they sort - this run's own image must survive its own prune.

Exit code: 0 on success (or if prune fails gracefully), non-zero only on critical error.
"""

import argparse
import base64
import json
import os
import sys
import time
from urllib.request import Request, urlopen
from urllib.error import URLError, HTTPError

MANIFEST_ACCEPT = ", ".join([
    "application/vnd.docker.distribution.manifest.list.v2+json",
    "application/vnd.docker.distribution.manifest.v2+json",
    "application/vnd.oci.image.index.v1+json",
    "application/vnd.oci.image.manifest.v1+json",
])


def parse_args():
    parser = argparse.ArgumentParser(
        description="Prune Docker registry tags, keeping only the latest SHA tag(s)."
    )
    parser.add_argument("--registry", default=os.environ.get("REGISTRY_HOST"))
    parser.add_argument("--repo", default=os.environ.get("REGISTRY_REPO"))
    parser.add_argument("--username", default=os.environ.get("REGISTRY_USERNAME"))
    parser.add_argument("--password", default=os.environ.get("REGISTRY_PASSWORD"))
    parser.add_argument("--keep", type=int, default=1, help="Number of SHA tags to keep (default: 1)")
    parser.add_argument(
        "--protect",
        action="append",
        default=[],
        help="Tag that must never be deleted (e.g. the SHA this CI run just built). Repeatable.",
    )
    return parser.parse_args()


def api_request(url: str, method: str, username: str, password: str, data=None, retries: int = 3) -> dict | None:
    """Make an authenticated API request with retry logic."""
    auth = base64.b64encode(f"{username}:{password}".encode()).decode()
    headers = {
        "Authorization": f"Basic {auth}",
        "Content-Type": "application/json",
    }

    for attempt in range(1, retries + 1):
        try:
            req = Request(url, method=method, headers=headers, data=data)
            with urlopen(req, timeout=30) as response:
                content = response.read()
                if content:
                    return json.loads(content)
                return {}
        except HTTPError as e:
            if e.code == 401:
                print(f"  [ERROR] Authentication failed (401)")
                return None
            if e.code == 404:
                print(f"  [WARN] Resource not found: {url}")
                return None
            print(f"  [RETRY {attempt}/{retries}] HTTP {e.code} for {url}")
        except URLError as e:
            print(f"  [RETRY {attempt}/{retries}] URL error: {e.reason}")
        except Exception as e:
            print(f"  [RETRY {attempt}/{retries}] Error: {e}")

        if attempt < retries:
            time.sleep(attempt * 2)

    print(f"  [ERROR] Failed after {retries} attempts for {url}")
    return None


def _get_manifest(registry: str, repo: str, ref: str, username: str, password: str):
    """GET a manifest by tag or digest. Returns (digest, media_type, body_json) or None."""
    url = f"https://{registry}/v2/{repo}/manifests/{ref}"
    auth = base64.b64encode(f"{username}:{password}".encode()).decode()

    for attempt in range(1, 4):
        try:
            req = Request(url, method="GET", headers={
                "Authorization": f"Basic {auth}",
                "Accept": MANIFEST_ACCEPT,
            })
            with urlopen(req, timeout=30) as response:
                digest = response.headers.get("Docker-Content-Digest", "")
                media_type = response.headers.get("Content-Type", "")
                body = json.loads(response.read())
                return digest, media_type, body
        except Exception as e:
            print(f"  [RETRY {attempt}/3] Fetching manifest {ref}: {e}")
            time.sleep(attempt)

    return None


def get_tag_created(registry: str, repo: str, tag: str, username: str, password: str) -> tuple[str, str] | None:
    """
    Get the digest and true build timestamp for a tag, by reading the
    "created" field baked into the image's config blob at build time -
    NOT the HTTP response's Date header, which only reflects the moment
    this GET request happened (every tag queried in the same prune run
    ends up with a near-identical Date, so sorting by it is close to
    random and can rank a tag pushed seconds ago as the oldest one).
    """
    result = _get_manifest(registry, repo, tag, username, password)
    if result is None:
        return None
    tag_digest, media_type, body = result

    # A tag can point at a multi-platform manifest list/index rather than
    # a single image manifest - descend into the first platform entry to
    # reach an actual image manifest with a "config" pointer.
    if "manifests" in body and body["manifests"]:
        child_digest = body["manifests"][0]["digest"]
        child = _get_manifest(registry, repo, child_digest, username, password)
        if child is None:
            return None
        _, _, body = child

    config = body.get("config")
    if not config or "digest" not in config:
        return None

    blob_url = f"https://{registry}/v2/{repo}/blobs/{config['digest']}"
    auth = base64.b64encode(f"{username}:{password}".encode()).decode()
    for attempt in range(1, 4):
        try:
            req = Request(blob_url, method="GET", headers={"Authorization": f"Basic {auth}"})
            with urlopen(req, timeout=30) as response:
                config_json = json.loads(response.read())
                created = config_json.get("created", "")
                return tag_digest, created
        except Exception as e:
            print(f"  [RETRY {attempt}/3] Fetching config blob for {tag}: {e}")
            time.sleep(attempt)

    return None


def delete_tag(registry: str, repo: str, digest: str, username: str, password: str) -> bool:
    """Delete a tag by its digest."""
    url = f"https://{registry}/v2/{repo}/manifests/{digest}"
    auth = base64.b64encode(f"{username}:{password}".encode()).decode()

    for attempt in range(1, 4):
        try:
            req = Request(url, method="DELETE", headers={
                "Authorization": f"Basic {auth}",
            })
            with urlopen(req, timeout=30) as response:
                if response.status in (200, 202, 404):
                    return True
        except HTTPError as e:
            if e.code == 404:
                return True  # Already deleted
            print(f"  [RETRY {attempt}/3] Deleting {digest[:20]}...: {e}")
        except Exception as e:
            print(f"  [RETRY {attempt}/3] Deleting {digest[:20]}...: {e}")

        time.sleep(attempt)

    return False


def is_sha_tag(tag: str) -> bool:
    """Check if tag looks like a SHA (40 hex chars)."""
    import re
    return bool(re.match(r"^[a-f0-9]{40}$", tag))


def prune_tags(registry: str, repo: str, username: str, password: str, keep: int = 1, protect: list[str] = ()) -> bool:
    """
    Main prune logic:
    - List all tags for the repo
    - Filter SHA-like tags
    - Sort by the image's actual build timestamp (newest first)
    - Keep newest `keep` tags, plus anything in `protect`
    - Delete older SHA tags by digest
    - Never delete non-SHA tags or protected tags
    """
    print(f"\n=== Pruning {registry}/{repo} ===")
    print(f"Strategy: Keep {keep} newest SHA tag(s), delete older SHA tags")
    print(f"Non-SHA tags (e.g., high-performance-latest, main-latest, latest) are preserved")
    if protect:
        print(f"Protected SHA tags (never deleted): {', '.join(protect)}")
    print()

    # Get catalog (list of repos)
    catalog_url = f"https://{registry}/v2/_catalog"
    catalog = api_request(catalog_url, "GET", username, password)
    if catalog is None:
        print("[ERROR] Failed to get repository catalog")
        return False

    if repo not in catalog.get("repositories", []):
        print(f"[INFO] Repository {repo} not found in catalog")
        return True

    # Get tags for repo
    tags_url = f"https://{registry}/v2/{repo}/tags/list"
    tags_data = api_request(tags_url, "GET", username, password)
    if tags_data is None:
        print(f"[ERROR] Failed to get tags for {repo}")
        return False

    all_tags = tags_data.get("tags", [])
    if not all_tags:
        print("[INFO] No tags found")
        return True

    # Separate SHA tags from non-SHA tags
    sha_tags = [t for t in all_tags if is_sha_tag(t)]
    non_sha_tags = [t for t in all_tags if not is_sha_tag(t)]

    print(f"Total tags: {len(all_tags)}")
    print(f"  SHA tags (candidates for pruning): {len(sha_tags)}")
    print(f"  Non-SHA tags (protected): {len(non_sha_tags)}")
    if non_sha_tags:
        print(f"  Protected tags: {', '.join(sorted(non_sha_tags))}")

    if not sha_tags:
        print("\n[INFO] No SHA tags to prune")
        return True

    # Get digest and true build time for each SHA tag
    tag_info = []
    for tag in sha_tags:
        result = get_tag_created(registry, repo, tag, username, password)
        if result:
            digest, created = result
            tag_info.append({
                "tag": tag,
                "digest": digest,
                "created": created,
                "timestamp": parse_image_date(created) if created else 0,
            })
        time.sleep(0.1)  # Be nice to the registry

    if not tag_info:
        print("\n[ERROR] Could not get info for any SHA tags")
        return False

    # Sort by timestamp (newest first)
    tag_info.sort(key=lambda x: x["timestamp"], reverse=True)

    print(f"\nSHA tags sorted by build time (newest first):")
    for i, info in enumerate(tag_info):
        is_protected = info["tag"] in protect
        will_keep = i < keep or is_protected
        marker = " [KEEP]" + (" (protected)" if is_protected else "") if will_keep else " [DELETE]"
        print(f"  {i+1}. {info['tag']} ({info['created'] or 'unknown date'}){marker}")

    # Delete older, unprotected SHA tags
    deleted_count = 0
    kept_count = 0

    for i, info in enumerate(tag_info):
        if i < keep or info["tag"] in protect:
            print(f"\n[KEEP] {info['tag']}")
            kept_count += 1
            continue

        print(f"\n[DELETE] {info['tag']} (digest: {info['digest'][:20]}...)")
        if delete_tag(registry, repo, info["digest"], username, password):
            print(f"  [OK] Deleted {info['tag']}")
            deleted_count += 1
        else:
            print(f"  [WARN] Failed to delete {info['tag']} (will retry next run)")

        time.sleep(0.2)  # Be nice to the registry

    print(f"\n=== Prune Summary ===")
    print(f"Tags kept: {kept_count}")
    print(f"Tags deleted: {deleted_count}")
    print(f"Tags protected (non-SHA): {len(non_sha_tags)}")

    return True


def parse_image_date(date_str: str) -> float:
    """Parse an OCI image config's ISO-8601 "created" field to a timestamp."""
    from datetime import datetime
    try:
        return datetime.fromisoformat(date_str.replace("Z", "+00:00")).timestamp()
    except Exception:
        return 0


def main():
    args = parse_args()

    # Validate required args
    registry = args.registry or os.environ.get("REGISTRY_HOST")
    repo = args.repo or os.environ.get("REGISTRY_REPO")
    username = args.username or os.environ.get("REGISTRY_USERNAME")
    password = args.password or os.environ.get("REGISTRY_PASSWORD")

    if not all([registry, repo, username, password]):
        print("[ERROR] Missing required arguments. Need: --registry, --repo, --username, --password")
        print("Or set environment variables: REGISTRY_HOST, REGISTRY_REPO, REGISTRY_USERNAME, REGISTRY_PASSWORD")
        sys.exit(1)

    print(f"Registry: {registry}")
    print(f"Repository: {repo}")
    print(f"Username: {username}")

    try:
        success = prune_tags(registry, repo, username, password, args.keep, args.protect)
        if success:
            print("\n[OK] Prune completed successfully")
            sys.exit(0)
        else:
            print("\n[WARN] Prune completed with some errors")
            sys.exit(0)  # Exit 0 per requirement - never fail workflow
    except Exception as e:
        print(f"\n[ERROR] Prune failed with exception: {e}")
        sys.exit(0)  # Exit 0 per requirement - never fail workflow


if __name__ == "__main__":
    main()