#!/usr/bin/env python3
"""
Launch GCP spot instance using Python API (no gcloud CLI needed).

Requirements:
    pip install google-cloud-compute google-auth

Setup:
    1. Create service account: https://console.cloud.google.com/iam-admin/serviceaccounts
    2. Download JSON key
    3. Set environment variable:
       export GOOGLE_APPLICATION_CREDENTIALS="/path/to/key.json"

    Or use: gcloud auth application-default login
"""

import os
import sys
import time
from google.cloud import compute_v1
from google.oauth2 import service_account
import logging

logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger(__name__)

# Configuration
PROJECT_ID = os.getenv('GCP_PROJECT_ID', None)
ZONE = 'us-central1-a'
INSTANCE_NAME = f'ensemble-test-{int(time.time())}'
MACHINE_TYPE = 'e2-medium'
IMAGE_FAMILY = 'ubuntu-2204-lts'
IMAGE_PROJECT = 'ubuntu-os-cloud'

STARTUP_SCRIPT = """#!/bin/bash
apt-get update
apt-get install -y python3-pip git

pip3 install --upgrade pip

cd /home
git clone https://huggingface.co/marcosremar2/ensemble-tts-annotation
cd ensemble-tts-annotation

pip3 install -q torch --index-url https://download.pytorch.org/whl/cpu
pip3 install -q transformers datasets librosa soundfile numpy pandas tqdm scikit-learn

python3 test_local.py > /tmp/test-results.log 2>&1

echo "✅ Test completed" >> /tmp/test-results.log
"""


def get_credentials():
    """Get GCP credentials."""
    creds_path = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')

    if creds_path and os.path.exists(creds_path):
        logger.info(f"Using credentials from: {creds_path}")
        return service_account.Credentials.from_service_account_file(creds_path)
    else:
        logger.info("Using default credentials (gcloud auth application-default login)")
        from google.auth import default
        credentials, project = default()
        return credentials


def get_project_id():
    """Get GCP project ID."""
    if PROJECT_ID:
        return PROJECT_ID

    # Try to get from gcloud config
    try:
        import subprocess
        result = subprocess.run(
            ['gcloud', 'config', 'get-value', 'project'],
            capture_output=True,
            text=True
        )
        if result.returncode == 0:
            project = result.stdout.strip()
            if project:
                return project
    except:
        pass

    # Ask user
    project = input("Enter GCP Project ID: ").strip()
    return project


def create_instance(project_id, zone, instance_name, machine_type):
    """Create spot instance."""
    logger.info("=" * 60)
    logger.info("Creating GCP Spot Instance")
    logger.info("=" * 60)
    logger.info(f"Project: {project_id}")
    logger.info(f"Zone: {zone}")
    logger.info(f"Instance: {instance_name}")
    logger.info(f"Machine Type: {machine_type}")
    logger.info(f"Estimated cost: ~$0.01/hr")
    logger.info("")

    # Get credentials
    credentials = get_credentials()

    # Create compute client
    instance_client = compute_v1.InstancesClient(credentials=credentials)

    # Get image
    logger.info("Getting Ubuntu image...")
    image_client = compute_v1.ImagesClient(credentials=credentials)
    image = image_client.get_from_family(
        project=IMAGE_PROJECT,
        family=IMAGE_FAMILY
    )

    # Configure instance
    logger.info("Configuring instance...")

    disk = compute_v1.AttachedDisk()
    disk.boot = True
    disk.auto_delete = True
    disk.initialize_params = compute_v1.AttachedDiskInitializeParams()
    disk.initialize_params.source_image = image.self_link
    disk.initialize_params.disk_size_gb = 20

    network_interface = compute_v1.NetworkInterface()
    network_interface.name = "global/networks/default"
    access_config = compute_v1.AccessConfig()
    access_config.name = "External NAT"
    access_config.type_ = "ONE_TO_ONE_NAT"
    network_interface.access_configs = [access_config]

    metadata = compute_v1.Metadata()
    metadata.items = [
        compute_v1.Items(key="startup-script", value=STARTUP_SCRIPT)
    ]

    scheduling = compute_v1.Scheduling()
    scheduling.preemptible = True
    scheduling.automatic_restart = False
    scheduling.on_host_maintenance = "TERMINATE"

    instance = compute_v1.Instance()
    instance.name = instance_name
    instance.machine_type = f"zones/{zone}/machineTypes/{machine_type}"
    instance.disks = [disk]
    instance.network_interfaces = [network_interface]
    instance.metadata = metadata
    instance.scheduling = scheduling

    # Create instance
    logger.info("Launching instance...")
    operation = instance_client.insert(
        project=project_id,
        zone=zone,
        instance_resource=instance
    )

    # Wait for operation to complete
    logger.info("Waiting for instance to be created...")
    while operation.status != compute_v1.Operation.Status.DONE:
        time.sleep(2)
        operation = compute_v1.ZoneOperationsClient(credentials=credentials).get(
            project=project_id,
            zone=zone,
            operation=operation.name
        )

    if operation.error:
        logger.error(f"Error creating instance: {operation.error}")
        return None

    # Get instance details
    instance = instance_client.get(
        project=project_id,
        zone=zone,
        instance=instance_name
    )

    external_ip = instance.network_interfaces[0].access_configs[0].nat_i_p

    logger.info("")
    logger.info("=" * 60)
    logger.info("✅ Instance created successfully!")
    logger.info("=" * 60)
    logger.info(f"Instance Name: {instance_name}")
    logger.info(f"External IP: {external_ip}")
    logger.info(f"Zone: {zone}")
    logger.info("")
    logger.info("SSH Command:")
    logger.info(f"  gcloud compute ssh {instance_name} --zone={zone} --project={project_id}")
    logger.info("")
    logger.info("Or from console:")
    logger.info(f"  https://console.cloud.google.com/compute/instances?project={project_id}")
    logger.info("")
    logger.info("Check test results (wait ~2 min for setup):")
    logger.info(f"  gcloud compute ssh {instance_name} --zone={zone} --command='cat /tmp/test-results.log'")
    logger.info("")
    logger.info("Delete instance:")
    logger.info(f"  gcloud compute instances delete {instance_name} --zone={zone} --project={project_id} --quiet")
    logger.info("")
    logger.info("=" * 60)
    logger.info("⚠️  Remember to delete the instance to avoid charges!")
    logger.info("=" * 60)

    return instance_name


def main():
    """Main function."""
    logger.info("GCP Spot Instance Launcher (Python API)")
    logger.info("")

    # Get project
    project_id = get_project_id()
    if not project_id:
        logger.error("No project ID provided")
        return 1

    # Create instance
    try:
        instance_name = create_instance(
            project_id=project_id,
            zone=ZONE,
            instance_name=INSTANCE_NAME,
            machine_type=MACHINE_TYPE
        )

        if instance_name:
            logger.info("\n✅ Success! Instance is launching.")
            logger.info("\nNext steps:")
            logger.info("1. Wait ~2 minutes for setup to complete")
            logger.info("2. SSH to instance and check results")
            logger.info("3. Delete instance when done")
            return 0
        else:
            logger.error("\n❌ Failed to create instance")
            return 1

    except Exception as e:
        logger.error(f"\n❌ Error: {e}")
        logger.error("\nMake sure you have:")
        logger.error("1. Installed: pip install google-cloud-compute")
        logger.error("2. Set credentials:")
        logger.error("   - Option A: export GOOGLE_APPLICATION_CREDENTIALS='/path/to/key.json'")
        logger.error("   - Option B: gcloud auth application-default login")
        logger.error("3. Enabled Compute Engine API in your project")
        return 1


if __name__ == "__main__":
    sys.exit(main())