#!/usr/bin/env python3
"""
RunPod Full Deployment Orchestrator

All-in-one script that handles the complete deployment workflow:
  1. Check/upload binaries and test data to volume
  2. Deploy pod with training job
  3. Monitor training progress
  4. Download trained models from volume
  5. Cleanup pod (optional)

This script orchestrates `upload_to_runpod_volume.py` and
`runpod_deploy_production.py` to provide a seamless deployment experience.

Prerequisites:
  1. Set RUNPOD_API_KEY environment variable
  2. Ensure .env.runpod exists with S3 credentials
  3. Build release binaries: cargo build --release --examples

Usage:
  # Smoke test (upload if needed, deploy, monitor, download)
  ./scripts/runpod_full_deploy.py --smoke-test

  # Full training with auto-cleanup after download
  ./scripts/runpod_full_deploy.py --full-training --cleanup

  # Skip upload (binaries already on volume)
  ./scripts/runpod_full_deploy.py --smoke-test --skip-upload

  # Force re-upload binaries even if unchanged
  ./scripts/runpod_full_deploy.py --full-training --force-upload

  # Custom configuration
  ./scripts/runpod_full_deploy.py --model tft --epochs 25 --batch-size 16

  # Dry run (show plan without executing)
  ./scripts/runpod_full_deploy.py --smoke-test --dry-run
"""

import argparse
import os
import subprocess
import sys
import time
from pathlib import Path
from typing import Optional

from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table

console = Console()


def run_command(cmd: list, description: str, dry_run: bool = False) -> bool:
    """
    Run a command and return success status.

    Args:
        cmd: Command to run (list of args)
        description: Description for logging
        dry_run: If True, print command without executing

    Returns:
        True if command succeeded, False otherwise
    """
    if dry_run:
        console.print(f"\n[dim]Would run: {' '.join(cmd)}[/dim]")
        return True

    console.print(f"\n[bold cyan]→ {description}[/bold cyan]")
    console.print(f"[dim]  Command: {' '.join(cmd)}[/dim]\n")

    try:
        result = subprocess.run(cmd, check=True)
        console.print(f"[green]✓ {description} completed[/green]")
        return True
    except subprocess.CalledProcessError as e:
        console.print(f"[red]✗ {description} failed with exit code {e.returncode}[/red]")
        return False
    except FileNotFoundError:
        console.print(f"[red]✗ Command not found: {cmd[0]}[/red]")
        return False


def check_prerequisites() -> bool:
    """Check if all prerequisites are met."""
    console.print("\n[bold cyan]Checking Prerequisites...[/bold cyan]")

    checks = []

    # Check RUNPOD_API_KEY
    if os.getenv("RUNPOD_API_KEY"):
        console.print("[green]✓ RUNPOD_API_KEY is set[/green]")
        checks.append(True)
    else:
        console.print("[red]✗ RUNPOD_API_KEY not set[/red]")
        console.print("  Set with: export RUNPOD_API_KEY='your-key-here'")
        checks.append(False)

    # Check .env.runpod
    env_file = Path(__file__).parent.parent / ".env.runpod"
    if env_file.exists():
        console.print(f"[green]✓ .env.runpod exists[/green]")
        checks.append(True)
    else:
        console.print(f"[red]✗ .env.runpod not found[/red]")
        console.print("  Create it first (see RUNPOD_S3_DEPLOYMENT_GUIDE.md)")
        checks.append(False)

    # Check upload script
    upload_script = Path(__file__).parent / "upload_to_runpod_volume.py"
    if upload_script.exists():
        console.print(f"[green]✓ upload_to_runpod_volume.py exists[/green]")
        checks.append(True)
    else:
        console.print(f"[red]✗ upload_to_runpod_volume.py not found[/red]")
        checks.append(False)

    # Check deployment script
    deploy_script = Path(__file__).parent / "runpod_deploy_production.py"
    if deploy_script.exists():
        console.print(f"[green]✓ runpod_deploy_production.py exists[/green]")
        checks.append(True)
    else:
        console.print(f"[red]✗ runpod_deploy_production.py not found[/red]")
        checks.append(False)

    # Check binaries
    binaries_dir = Path("target/release/examples")
    if binaries_dir.exists():
        binaries = list(binaries_dir.glob("train_*"))
        if binaries:
            console.print(f"[green]✓ Found {len(binaries)} training binaries[/green]")
            checks.append(True)
        else:
            console.print(f"[yellow]⚠ No training binaries found[/yellow]")
            console.print("  Run: cargo build --release --examples")
            checks.append(False)
    else:
        console.print(f"[yellow]⚠ Binaries directory not found[/yellow]")
        console.print("  Run: cargo build --release --examples")
        checks.append(False)

    return all(checks)


def upload_to_volume(force: bool = False, dry_run: bool = False) -> bool:
    """Upload binaries and test data to volume."""
    script = Path(__file__).parent / "upload_to_runpod_volume.py"

    cmd = [str(script), "--all"]
    if force:
        cmd.append("--force")
    if dry_run:
        cmd.append("--dry-run")

    return run_command(cmd, "Upload binaries and test data to volume", dry_run)


def deploy_pod(args: argparse.Namespace, dry_run: bool = False) -> Optional[str]:
    """
    Deploy training pod and return pod ID if successful.

    Args:
        args: Parsed command-line arguments
        dry_run: If True, show plan without deploying

    Returns:
        Pod ID if deployment succeeded, None otherwise
    """
    script = Path(__file__).parent / "runpod_deploy_production.py"

    cmd = [str(script)]

    # Add training mode
    if args.smoke_test:
        cmd.append("--smoke-test")
    elif args.full_training:
        cmd.append("--full-training")

    # Add model and training params
    if args.model:
        cmd.extend(["--model", args.model])
    if args.epochs:
        cmd.extend(["--epochs", str(args.epochs)])
    if args.batch_size:
        cmd.extend(["--batch-size", str(args.batch_size)])
    if args.dataset:
        cmd.extend(["--dataset", args.dataset])
    if args.use_int8:
        cmd.append("--use-int8")

    # Add GPU preference
    if args.gpu_type:
        cmd.extend(["--gpu-type", args.gpu_type])

    # Add execution options
    if dry_run:
        cmd.append("--dry-run")
    if args.monitor:
        cmd.append("--monitor")
    if args.yes:
        cmd.append("--yes")

    success = run_command(cmd, "Deploy training pod", dry_run)

    if not success:
        return None

    # Parse pod ID from output (if not dry-run)
    # TODO: Modify runpod_deploy_production.py to output pod ID in machine-readable format
    # For now, return placeholder
    return "pod-id-placeholder" if not dry_run else "dry-run-pod-id"


def download_models(model: str, dry_run: bool = False) -> bool:
    """Download trained models from volume via S3."""
    if dry_run:
        console.print("\n[dim]Would download models from volume[/dim]")
        return True

    console.print("\n[bold cyan]Downloading Models from Volume...[/bold cyan]")

    try:
        import boto3
        from botocore.exceptions import ClientError
        from dotenv import load_dotenv

        # Load credentials
        env_path = Path(__file__).parent.parent / ".env.runpod"
        load_dotenv(env_path)

        # Initialize S3 client
        s3_client = boto3.client(
            "s3",
            aws_access_key_id=os.getenv("RUNPOD_S3_ACCESS_KEY"),
            aws_secret_access_key=os.getenv("RUNPOD_S3_SECRET"),
            region_name=os.getenv("RUNPOD_S3_REGION"),
            endpoint_url=os.getenv("RUNPOD_S3_ENDPOINT"),
        )

        volume_id = os.getenv("RUNPOD_VOLUME_ID")

        # List models in volume
        prefix = f"models/{model}/"
        response = s3_client.list_objects_v2(Bucket=volume_id, Prefix=prefix)

        if "Contents" not in response:
            console.print(f"[yellow]⚠ No models found at {prefix}[/yellow]")
            return False

        # Download each model file
        local_dir = Path("models/runpod_trained") / model
        local_dir.mkdir(parents=True, exist_ok=True)

        for obj in response["Contents"]:
            s3_key = obj["Key"]
            filename = Path(s3_key).name
            local_path = local_dir / filename

            console.print(f"  Downloading {filename}...")
            s3_client.download_file(volume_id, s3_key, str(local_path))
            console.print(f"  [green]✓ Saved to {local_path}[/green]")

        console.print(f"\n[green]✓ Downloaded {len(response['Contents'])} model files[/green]")
        return True

    except Exception as e:
        console.print(f"[red]✗ Failed to download models: {e}[/red]")
        return False


def cleanup_pod(pod_id: str, dry_run: bool = False) -> bool:
    """Terminate pod after training completes."""
    if dry_run:
        console.print(f"\n[dim]Would terminate pod: {pod_id}[/dim]")
        return True

    console.print(f"\n[bold cyan]Terminating Pod: {pod_id}[/bold cyan]")

    # TODO: Implement GraphQL mutation to terminate pod
    # For now, just print instructions
    console.print("[yellow]⚠ Auto-cleanup not yet implemented[/yellow]")
    console.print(f"  Terminate manually: https://www.runpod.io/console/pods/{pod_id}")

    return True


def main():
    parser = argparse.ArgumentParser(
        description="Full RunPod deployment orchestrator (upload → deploy → monitor → download)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Complete smoke test workflow
  %(prog)s --smoke-test

  # Full training with cleanup
  %(prog)s --full-training --cleanup

  # Skip upload if binaries already on volume
  %(prog)s --smoke-test --skip-upload

  # Force re-upload binaries
  %(prog)s --full-training --force-upload

  # Custom configuration
  %(prog)s --model tft --epochs 25 --batch-size 16

  # Dry run (show plan)
  %(prog)s --smoke-test --dry-run
        """,
    )

    # Training mode
    parser.add_argument("--smoke-test", action="store_true", help="Quick smoke test")
    parser.add_argument("--full-training", action="store_true", help="Full training run")

    # Upload options
    parser.add_argument(
        "--skip-upload", action="store_true", help="Skip upload (binaries already on volume)"
    )
    parser.add_argument(
        "--force-upload",
        action="store_true",
        help="Force re-upload even if checksums match",
    )

    # Training parameters
    parser.add_argument("--model", default="tft", help="Model to train (default: tft)")
    parser.add_argument("--epochs", type=int, help="Number of epochs")
    parser.add_argument("--batch-size", type=int, help="Batch size")
    parser.add_argument("--dataset", help="Dataset filename")
    parser.add_argument("--use-int8", action="store_true", help="Use INT8 quantization (TFT only)")

    # GPU selection
    parser.add_argument("--gpu-type", help="Preferred GPU type")

    # Execution options
    parser.add_argument("--monitor", action="store_true", help="Monitor pod in real-time")
    parser.add_argument("--cleanup", action="store_true", help="Terminate pod after completion")
    parser.add_argument("--yes", action="store_true", help="Skip confirmation prompts")
    parser.add_argument("--dry-run", action="store_true", help="Show plan without executing")

    args = parser.parse_args()

    # Display header
    console.print(Panel.fit(
        "[bold cyan]🦊 Foxhunt Full RunPod Deployment Orchestrator[/bold cyan]",
        border_style="cyan"
    ))

    # Check prerequisites
    if not check_prerequisites():
        console.print("\n[red]✗ Prerequisites check failed[/red]")
        console.print("  Fix the issues above and try again")
        return 1

    # Step 1: Upload binaries and test data (optional)
    if not args.skip_upload:
        console.print("\n[bold]Step 1: Upload Binaries and Test Data[/bold]")
        if not upload_to_volume(force=args.force_upload, dry_run=args.dry_run):
            console.print("\n[red]✗ Upload failed[/red]")
            return 1
    else:
        console.print("\n[bold]Step 1: Upload[/bold] [dim](skipped)[/dim]")

    # Step 2: Deploy pod
    console.print("\n[bold]Step 2: Deploy Training Pod[/bold]")
    pod_id = deploy_pod(args, dry_run=args.dry_run)
    if not pod_id:
        console.print("\n[red]✗ Deployment failed[/red]")
        return 1

    # Step 3: Monitor (handled by runpod_deploy_production.py if --monitor flag passed)
    if args.monitor:
        console.print("\n[bold]Step 3: Monitor Training[/bold]")
        console.print("[dim](Monitoring handled by deployment script)[/dim]")

    # Step 4: Download models
    if not args.dry_run:
        console.print("\n[bold]Step 4: Download Trained Models[/bold]")
        console.print("[yellow]⚠ Waiting for training to complete...[/yellow]")
        console.print("  Run this manually after training: ./scripts/download_models_from_volume.py")
        # TODO: Poll pod status and download automatically when complete

    # Step 5: Cleanup (optional)
    if args.cleanup:
        console.print("\n[bold]Step 5: Cleanup Pod[/bold]")
        cleanup_pod(pod_id, dry_run=args.dry_run)

    # Final summary
    console.print("\n" + "=" * 70)
    console.print("[bold green]✓ Deployment Orchestration Complete![/bold green]")
    console.print("=" * 70)

    if not args.dry_run:
        console.print(f"\n📊 Pod ID: {pod_id}")
        console.print(f"🔗 Console: https://www.runpod.io/console/pods/{pod_id}")

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        console.print("\n\n[yellow]⚠ Interrupted by user[/yellow]")
        sys.exit(130)