#!/usr/bin/env python3
"""
Script to save Docker images as artifacts for GitHub Actions.
Handles both CPU and GPU variants with proper versioning and metadata.
"""

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime


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 get_docker_images():
    """Get list of available Docker images."""
    print("=== DEBUG: Available Docker images ===")
    run_command(
        'docker images --format "table {{.Repository}}:{{.Tag}}\t{{.Size}}\t{{.ID}}"'
    )
    print("=======================================")


def list_tmp_files():
    """List tar files in /tmp for debugging."""
    print("=== DEBUG: Files in /tmp ===")
    result = run_command(
        "ls -lh /tmp/*.tar* 2>/dev/null", check=False, capture_output=True
    )
    print(result if result else "No tar files in /tmp")
    print("=======================================")


def save_amd64_image(image_name, docker_version, variant):
    """Save AMD64 Docker image to tar.gz file."""
    tag = f"{image_name}:{docker_version}-amd64"
    output_file = f"cyborgdb-service-{variant}-{docker_version}-amd64-docker.tar.gz"

    print(f"Saving {variant.upper()} AMD64 Docker image with tag: {tag}")

    # Check if image exists
    image_id = run_command(f'docker images -q "{tag}"', capture_output=True)
    if image_id:
        run_command(f'docker save "{tag}" | gzip > {output_file}')
        print(f"✅ {variant.upper()} AMD64 image saved successfully")
        return output_file
    else:
        print(f"❌ Error: {variant.upper()} AMD64 image {tag} not found")
        sys.exit(1)


def save_arm64_image(docker_version, variant):
    """Save ARM64 Docker image from /tmp tar file (CPU only)."""
    if variant != "cpu":
        return None

    print("Compressing ARM64 Docker image...")
    tar_file = "/tmp/cyborgdb-service-cpu-arm64.tar"
    output_file = f"cyborgdb-service-cpu-{docker_version}-arm64-docker.tar.gz"

    if os.path.exists(tar_file):
        run_command(f"gzip -c {tar_file} > {output_file}")
        run_command(f"rm -f {tar_file}")
        print("✅ ARM64 image compressed successfully")
        return output_file
    else:
        print("❌ Error: ARM64 tar file not found")
        sys.exit(1)


def handle_gpu_image(docker_version):
    """Handle GPU image which is already compressed during export."""
    print("Moving GPU AMD64 Docker image...")
    tar_file = "/tmp/cyborgdb-service-gpu-amd64.tar"
    output_file = f"cyborgdb-service-gpu-{docker_version}-amd64-docker.tar.gz"

    if os.path.exists(tar_file):
        # Check disk space
        print("Disk space before move:")
        run_command("df -h /tmp")

        # Move the already-compressed file
        run_command(f"mv {tar_file} {output_file}")

        print("Disk space after move:")
        run_command("df -h /tmp")
        print("✅ GPU AMD64 image moved successfully")
        return output_file
    else:
        print(
            "❌ Error: GPU AMD64 tar file not found at /tmp/cyborgdb-service-gpu-amd64.tar"
        )
        print("Available files in /tmp:")
        run_command("ls -la /tmp/")
        sys.exit(1)


def create_metadata(
    image_name, docker_version, python_version, variant, git_ref, git_sha, platforms
):
    """Create metadata JSON file."""
    metadata = {
        "image_name": f"docker.io/{image_name}",
        "version": docker_version,
        "python_version": python_version,
        "variant": variant,
        "build_date": datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
        "git_ref": git_ref,
        "git_sha": git_sha,
        "platforms": platforms,
        "usage": {
            "pull_multiarch": f"docker pull docker.io/{image_name}:{docker_version}",
            "load_amd64": f"docker load < cyborgdb-service-{variant}-{docker_version}-amd64-docker.tar.gz",
        },
    }

    # Add ARM64 load command for CPU variant
    if variant == "cpu":
        metadata["usage"]["load_arm64"] = (
            f"docker load < cyborgdb-service-{variant}-{docker_version}-arm64-docker.tar.gz"
        )

    # Add appropriate run command
    gpu_flag = "--gpus all " if variant == "gpu" else ""
    metadata["usage"]["run"] = (
        f"docker run {gpu_flag}-p 8000:8000 "
        f"-e CYBORGDB_API_KEY=your-key "
        f"-e CYBORGDB_DB_TYPE=redis "
        f"-e CYBORGDB_CONNECTION_STRING=host:localhost,port:6379,db:0 "
        f"docker.io/{image_name}:{docker_version}"
    )

    output_file = f"docker-image-info-{variant}.json"
    with open(output_file, "w") as f:
        json.dump(metadata, f, indent=2)

    print(f"Created metadata file: {output_file}")
    return output_file


def list_artifacts(docker_version, variant):
    """List created artifact files."""
    print("Docker image sizes:")
    pattern = f"cyborgdb-service-{variant}-{docker_version}-*-docker.tar.gz"
    result = run_command(
        f"ls -lh {pattern} 2>/dev/null", check=False, capture_output=True
    )
    print(result if result else "No files found")


def main():
    parser = argparse.ArgumentParser(description="Save Docker images as artifacts")
    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 tagging"
    )
    parser.add_argument(
        "--python-version", required=True, help="Python/PEP 440 version"
    )
    parser.add_argument("--image-name", required=True, help="Docker image name")
    parser.add_argument("--git-ref", required=True, help="Git reference (branch/tag)")
    parser.add_argument("--git-sha", required=True, help="Git commit SHA")

    args = parser.parse_args()

    # Debug: Show available images and files
    get_docker_images()
    list_tmp_files()

    artifacts = []

    if args.variant == "cpu":
        # Save both AMD64 and ARM64 images
        artifacts.append(
            save_amd64_image(args.image_name, args.docker_version, args.variant)
        )
        artifacts.append(save_arm64_image(args.docker_version, args.variant))
        platforms = ["linux/amd64", "linux/arm64"]
    else:
        # GPU variant - handle pre-compressed image
        artifacts.append(handle_gpu_image(args.docker_version))
        platforms = ["linux/amd64"]

    # Create metadata file
    artifacts.append(
        create_metadata(
            args.image_name,
            args.docker_version,
            args.python_version,
            args.variant,
            args.git_ref,
            args.git_sha,
            platforms,
        )
    )

    # List created artifacts
    list_artifacts(args.docker_version, args.variant)

    print(f"\n✅ Successfully created artifacts for {args.variant.upper()} variant:")
    for artifact in artifacts:
        if artifact:
            print(f"  - {artifact}")


if __name__ == "__main__":
    main()
