#!/usr/bin/env python3
"""
Script to upload Docker images to Cloudsmith repository.
Handles both CPU (multi-arch) and GPU (single-arch) variants.
"""

import argparse
import os
import subprocess
import sys
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 install_cloudsmith_cli():
    """Install cloudsmith CLI if not already installed."""
    print("Installing cloudsmith-cli...")
    run_command("pip install cloudsmith-cli")


def verify_environment():
    """Verify required environment variables are set."""
    required_vars = ["CLOUDSMITH_API_KEY", "CLOUDSMITH_ORG", "CLOUDSMITH_REPO"]
    missing = []

    for var in required_vars:
        if not os.environ.get(var):
            missing.append(var)

    if missing:
        print(f"❌ Error: Missing required environment variables: {', '.join(missing)}")
        sys.exit(1)

    print("✅ Environment variables verified")
    return {
        "api_key": os.environ["CLOUDSMITH_API_KEY"],
        "org": os.environ["CLOUDSMITH_ORG"],
        "repo": os.environ["CLOUDSMITH_REPO"],
    }


def upload_docker_image(config, tar_file, tags):
    """Upload a Docker image to Cloudsmith."""
    if not Path(tar_file).exists():
        print(f"❌ Error: File {tar_file} not found")
        sys.exit(1)

    print(f"Uploading {tar_file}...")

    # Build the cloudsmith push command
    cmd = [
        "cloudsmith",
        "push",
        "docker",
        "--api-key",
        f'"{config["api_key"]}"',
        "--republish",
        f'"{config["org"]}/{config["repo"]}"',
        f'"{tar_file}"',
        "--tags",
        f'"{tags}"',
    ]

    run_command(" ".join(cmd))
    print(f"✅ Successfully uploaded {tar_file}")


def upload_cpu_images(config, image_name, version):
    """Upload CPU variant images (AMD64 and ARM64)."""
    print(f"Uploading CPU Docker images to {config['org']}/{config['repo']}...")

    # Upload AMD64 image
    print("Uploading AMD64 image...")
    amd64_file = f"cyborgdb-service-cpu-{version}-amd64-docker.tar.gz"
    upload_docker_image(config, amd64_file, "amd64,cpu")

    # Upload ARM64 image
    print("Uploading ARM64 image...")
    arm64_file = f"cyborgdb-service-cpu-{version}-arm64-docker.tar.gz"
    upload_docker_image(config, arm64_file, "arm64,cpu")

    print("✅ Successfully uploaded CPU images to Cloudsmith")


def upload_gpu_images(config, image_name, version):
    """Upload GPU variant images (AMD64 only)."""
    print(f"Uploading GPU Docker images to {config['org']}/{config['repo']}...")

    # Upload AMD64 image
    print("Uploading AMD64 image...")
    amd64_file = f"cyborgdb-service-gpu-{version}-amd64-docker.tar.gz"
    upload_docker_image(config, amd64_file, "amd64,gpu")

    print("✅ Successfully uploaded GPU images to Cloudsmith")


def main():
    parser = argparse.ArgumentParser(description="Upload Docker images to Cloudsmith")
    parser.add_argument(
        "--variant",
        required=True,
        choices=["cpu", "gpu"],
        help="Variant type (cpu or gpu)",
    )
    parser.add_argument(
        "--docker-version", required=True, help="Docker version for the images"
    )
    parser.add_argument("--image-name", required=True, help="Docker image name")

    args = parser.parse_args()

    print(f"Uploading {args.variant.upper()} images to Cloudsmith")
    print(f"Image: {args.image_name}")
    print(f"Version: {args.docker_version}")
    print("-" * 50)

    try:
        # Install CLI
        install_cloudsmith_cli()

        # Verify environment
        config = verify_environment()

        # Upload based on variant
        if args.variant == "cpu":
            upload_cpu_images(config, args.image_name, args.docker_version)
        else:
            upload_gpu_images(config, args.image_name, args.docker_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()
