#!/usr/bin/env python3
"""
Script to load Docker images from artifacts and push them to Docker Hub.
Handles both CPU (multi-arch) and GPU (single-arch) variants.
"""

import argparse
import subprocess
import sys
import re
from pathlib import Path


def run_command(cmd, check=True, capture_output=False):
    """Run a shell command and return the result."""
    if capture_output:
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True, check=check
        )
        return result.stdout.strip()
    else:
        return subprocess.run(cmd, shell=True, check=check)


def load_docker_image(tar_file):
    """Load a Docker image from a tar.gz file."""
    if not Path(tar_file).exists():
        print(f"❌ Error: File {tar_file} not found")
        sys.exit(1)

    print(f"Loading {tar_file}...")
    run_command(f"docker load < {tar_file}")
    print(f"✅ Loaded {tar_file}")


def push_image(image_tag):
    """Push a Docker image to the registry."""
    print(f"Pushing {image_tag}...")
    run_command(f"docker push {image_tag}")
    print(f"✅ Pushed {image_tag}")


def create_and_push_manifest(image_name, tag, platforms, version=None):
    """Create and push a multi-arch manifest."""
    print(f"Creating multi-arch manifest for {image_name}:{tag}")

    # Build the manifest create command
    if version and tag != version:
        # For non-version tags (like latest), use version-specific platform images
        platform_images = [
            f"{image_name}:{version}-{platform}" for platform in platforms
        ]
    else:
        # For version tags, use tag-specific platform images
        platform_images = [f"{image_name}:{tag}-{platform}" for platform in platforms]

    manifest_cmd = f"docker manifest create {image_name}:{tag} " + " ".join(
        platform_images
    )

    run_command(manifest_cmd)
    run_command(f"docker manifest push {image_name}:{tag}")
    print(f"✅ Pushed manifest {image_name}:{tag}")


def extract_version_tags(version):
    """Extract major and major.minor version tags from a full version."""
    tags = []

    # Match semantic version pattern
    match = re.match(r"^(\d+)\.(\d+)\.(\d+)", version)
    if match:
        major = match.group(1)
        minor = match.group(2)

        # Add major.minor tag
        tags.append(f"{major}.{minor}")
        # Add major tag
        tags.append(major)

    return tags


def push_cpu_images(image_name, version):
    """Load and push CPU images (AMD64 and ARM64)."""
    print("Loading Docker images...")

    # Load platform-specific images
    amd64_file = f"cyborgdb-service-cpu-{version}-amd64-docker.tar.gz"
    arm64_file = f"cyborgdb-service-cpu-{version}-arm64-docker.tar.gz"

    load_docker_image(amd64_file)
    load_docker_image(arm64_file)

    # Tag images if needed (they should already be tagged from load)
    print("Tagging images...")
    run_command(f"docker tag {image_name}:{version}-amd64 {image_name}:{version}-amd64")
    run_command(f"docker tag {image_name}:{version}-arm64 {image_name}:{version}-arm64")

    # Push platform-specific images
    print("Pushing platform-specific images...")
    push_image(f"{image_name}:{version}-amd64")
    push_image(f"{image_name}:{version}-arm64")

    # Create and push multi-arch manifests
    platforms = ["amd64", "arm64"]

    # Version tag
    create_and_push_manifest(image_name, version, platforms)

    # Latest tag
    print("Creating latest tag...")
    create_and_push_manifest(image_name, "latest", platforms, version)

    # Additional version tags (major.minor, major)
    for tag in extract_version_tags(version):
        print(f"Creating version tag: {tag}")
        create_and_push_manifest(image_name, tag, platforms, version)

    print("✅ Successfully published CPU images to Docker Hub")


def push_gpu_images(image_name, version):
    """Load and push GPU images (AMD64 only)."""
    print("Loading Docker images...")

    # Load AMD64 image
    amd64_file = f"cyborgdb-service-gpu-{version}-amd64-docker.tar.gz"
    load_docker_image(amd64_file)

    # Tag image for Docker Hub
    print("Tagging image...")
    run_command(f"docker tag {image_name}:{version}-amd64 {image_name}:{version}")

    # Push versioned image
    print("Pushing versioned image...")
    push_image(f"{image_name}:{version}")

    # Create additional version tags
    print("Creating additional version tags...")

    # Latest tag
    run_command(f"docker tag {image_name}:{version} {image_name}:latest")
    push_image(f"{image_name}:latest")

    # Additional version tags (major.minor, major)
    for tag in extract_version_tags(version):
        print(f"Creating version tag: {tag}")
        run_command(f"docker tag {image_name}:{version} {image_name}:{tag}")
        push_image(f"{image_name}:{tag}")

    print("✅ Successfully published GPU image to Docker Hub")


def main():
    parser = argparse.ArgumentParser(description="Push Docker images to Docker Hub")
    parser.add_argument(
        "--variant",
        required=True,
        choices=["cpu", "gpu"],
        help="Variant type (cpu or gpu)",
    )
    parser.add_argument("--version", required=True, help="Docker version for tagging")
    parser.add_argument(
        "--image-name", required=True, help="Docker image name (without tag)"
    )

    args = parser.parse_args()

    print(f"Publishing {args.variant.upper()} images to Docker Hub")
    print(f"Image: {args.image_name}")
    print(f"Version: {args.version}")
    print("-" * 50)

    try:
        if args.variant == "cpu":
            push_cpu_images(args.image_name, args.version)
        else:
            push_gpu_images(args.image_name, args.version)
    except subprocess.CalledProcessError as e:
        print(f"❌ Error: Command failed: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()
