#!/usr/bin/env python3
"""
Quick Binary Upload Workflow Script

Uploads release binaries to RunPod S3 for fast development iteration.
Automatically finds binaries, validates executability, tracks upload progress,
and generates timestamped filenames for versioning.

Usage:
    python3 scripts/upload_binary.py --binary-name hyperopt_mamba2_demo
    python3 scripts/upload_binary.py --binary-path ./target/release/examples/custom_binary --force
    python3 scripts/upload_binary.py --binary-name hyperopt_tft_demo --no-timestamp
"""

import argparse
import os
import stat
import sys
from datetime import datetime
from pathlib import Path

# Check for .venv activation
if not hasattr(sys, 'real_prefix') and not (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
    print("WARNING: Not running in a virtual environment (.venv)")
    print("         Recommended: source .venv/bin/activate")
    print()

# Get project root for later use
script_dir = Path(__file__).parent
project_root = script_dir.parent

try:
    # Import from runpod module (clean architecture - no sys.path hacks)
    from runpod import S3Client
    from runpod.config import get_config
    from runpod.errors import S3Error, ConfigurationError
    from rich.console import Console
    from rich.progress import (
        Progress,
        BarColumn,
        DownloadColumn,
        TransferSpeedColumn,
        TimeRemainingColumn,
        TextColumn,
    )

    HAS_DEPENDENCIES = True
except ImportError as e:
    print(f"ERROR: Failed to import required modules: {e}")
    print("\nRequired dependencies:")
    print("  - runpod module")
    print("  - rich (pip install rich)")
    print("\nInstall with: pip install -r runpod/requirements.txt")
    sys.exit(1)

console = Console()

# Default binary search path
DEFAULT_BINARY_DIR = project_root / "target" / "release" / "examples"


def find_binary(binary_name: str) -> Path | None:
    """
    Find binary in target/release/examples/ directory.

    Args:
        binary_name: Name of binary (without path)

    Returns:
        Path to binary or None if not found
    """
    # Try exact match first
    binary_path = DEFAULT_BINARY_DIR / binary_name
    if binary_path.exists():
        return binary_path

    # Try with common suffixes
    for suffix in ["", "-*"]:  # Try exact, then with build hash
        pattern = f"{binary_name}{suffix}"
        matches = list(DEFAULT_BINARY_DIR.glob(pattern))

        # Filter out .d dependency files
        matches = [m for m in matches if not m.name.endswith(".d")]

        if matches:
            # Return most recent if multiple matches
            return max(matches, key=lambda p: p.stat().st_mtime)

    return None


def validate_binary(binary_path: Path) -> tuple[bool, str]:
    """
    Validate binary exists and is executable.

    Args:
        binary_path: Path to binary

    Returns:
        Tuple of (is_valid, error_message)
    """
    if not binary_path.exists():
        return False, f"Binary not found: {binary_path}"

    if not binary_path.is_file():
        return False, f"Not a file: {binary_path}"

    # Check if executable (Unix permissions)
    file_stat = binary_path.stat()
    if not (file_stat.st_mode & stat.S_IXUSR):
        return False, f"Binary is not executable: {binary_path}"

    # Check file size (warn if suspiciously small)
    file_size = file_stat.st_size
    if file_size < 1024 * 100:  # Less than 100KB
        return False, f"Binary suspiciously small ({file_size} bytes): {binary_path}"

    return True, ""


def generate_s3_key(binary_path: Path, use_timestamp: bool = True, cuda_variant: bool = True) -> str:
    """
    Generate S3 key with optional timestamp.

    Args:
        binary_path: Path to binary
        use_timestamp: Add timestamp to filename
        cuda_variant: Add _cuda suffix

    Returns:
        S3 key (e.g., "binaries/hyperopt_mamba2_demo_cuda_20251030_120000")
    """
    binary_name = binary_path.stem  # Remove -hash suffix if present

    # Clean up name (remove build hash like -875831286f8f5985)
    if "-" in binary_name:
        # Only remove if it looks like a hash (32 hex chars after dash)
        parts = binary_name.rsplit("-", 1)
        if len(parts) == 2 and len(parts[1]) == 16 and all(c in "0123456789abcdef" for c in parts[1]):
            binary_name = parts[0]

    # Build S3 key components
    components = [binary_name]

    if cuda_variant:
        components.append("cuda")

    if use_timestamp:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        components.append(timestamp)

    filename = "_".join(components)

    return f"binaries/{filename}"


def format_size(bytes_size: int) -> str:
    """Format byte size to human-readable string."""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if bytes_size < 1024.0:
            return f"{bytes_size:.1f} {unit}"
        bytes_size /= 1024.0
    return f"{bytes_size:.1f} TB"


def upload_binary_with_progress(
    s3_client: S3Client,
    binary_path: Path,
    s3_key: str,
    force: bool = False
) -> bool:
    """
    Upload binary with rich progress bar.

    Args:
        s3_client: S3Client instance
        binary_path: Path to binary
        s3_key: Destination S3 key
        force: Skip checksum check and force upload

    Returns:
        True if upload succeeded

    Raises:
        S3Error: On upload failure
    """
    file_size = binary_path.stat().st_size

    console.print(f"\n[cyan]Upload Plan:[/cyan]")
    console.print(f"  Source:      {binary_path}")
    console.print(f"  Destination: s3://{s3_client.bucket}/{s3_key}")
    console.print(f"  Size:        {format_size(file_size)} ({file_size:,} bytes)")
    console.print(f"  Force:       {force}")
    console.print()

    # Check if upload needed (unless forced)
    if not force:
        needs_upload, reason = s3_client.needs_upload(binary_path, s3_key)
        if not needs_upload:
            console.print(f"[green]✓ Binary up-to-date[/green]: {reason}")
            console.print(f"  S3 URI: s3://{s3_client.bucket}/{s3_key}")
            return True
        console.print(f"[yellow]↻ Upload required[/yellow]: {reason}\n")

    # Upload with progress bar
    with Progress(
        TextColumn("[bold blue]{task.description}"),
        BarColumn(),
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeRemainingColumn(),
        console=console,
    ) as progress:
        task = progress.add_task(f"Uploading {binary_path.name}", total=file_size)

        def callback(bytes_transferred):
            progress.update(task, completed=bytes_transferred)

        success = s3_client.upload_binary(
            binary_path,
            s3_key,
            force=force,
            progress_callback=callback
        )

    if success:
        console.print(f"\n[green]✅ Upload complete![/green]")
        console.print(f"  S3 URI: s3://{s3_client.bucket}/{s3_key}")
        console.print(f"\n[cyan]Next steps:[/cyan]")
        console.print(f"  1. Binary available at: /runpod-volume/{s3_key}")
        console.print(f"  2. Use in deployment: --command '/runpod-volume/{s3_key} --epochs 50'")
        console.print(f"  3. Verify: aws s3 ls s3://{s3_client.bucket}/{s3_key} --profile runpod")

    return success


def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Upload Rust binaries to RunPod S3 for fast development iteration",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Upload by name (auto-finds in target/release/examples/)
  python3 scripts/upload_binary.py --binary-name hyperopt_mamba2_demo

  # Upload custom binary path
  python3 scripts/upload_binary.py --binary-path ./custom_binary --force

  # Upload without timestamp (overwrites existing)
  python3 scripts/upload_binary.py --binary-name hyperopt_tft_demo --no-timestamp --force

  # Upload without CUDA suffix
  python3 scripts/upload_binary.py --binary-name hyperopt_dqn_demo --no-cuda

  # Dry run (validate only, don't upload)
  python3 scripts/upload_binary.py --binary-name hyperopt_ppo_demo --dry-run

Features:
  - Auto-finds binaries in target/release/examples/
  - Validates binary exists and is executable
  - Checks file size (warns if < 100KB)
  - Shows progress bar with transfer speed
  - MD5 checksum validation (skips upload if unchanged)
  - Generates timestamped names for versioning
  - Returns S3 path for deployment

S3 Organization:
  s3://se3zdnb5o4/binaries/
    ├── hyperopt_mamba2_demo_cuda_20251030_120000
    ├── hyperopt_tft_demo_cuda_20251030_143000
    └── hyperopt_dqn_demo_cuda_20251030_150000

Requirements:
  - .venv activated: source .venv/bin/activate
  - .env.runpod configured with S3 credentials
  - foxhunt_runpod module installed
        """
    )

    # Binary selection (mutually exclusive)
    binary_group = parser.add_mutually_exclusive_group(required=True)
    binary_group.add_argument(
        '--binary-name',
        help='Binary name to find in target/release/examples/ (e.g., hyperopt_mamba2_demo)'
    )
    binary_group.add_argument(
        '--binary-path',
        type=Path,
        help='Direct path to binary file'
    )

    # Upload options
    parser.add_argument(
        '--force',
        action='store_true',
        help='Force upload even if checksum matches (overwrite existing)'
    )
    parser.add_argument(
        '--no-timestamp',
        action='store_true',
        help='Do not add timestamp to filename (overwrites existing binary)'
    )
    parser.add_argument(
        '--no-cuda',
        action='store_true',
        help='Do not add _cuda suffix to filename'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Validate binary only, do not upload'
    )

    args = parser.parse_args()

    # Step 1: Find binary
    console.print("\n[bold]🔍 Step 1: Locating binary[/bold]")

    if args.binary_name:
        console.print(f"  Searching for: {args.binary_name}")
        console.print(f"  Search path:   {DEFAULT_BINARY_DIR}")

        binary_path = find_binary(args.binary_name)
        if not binary_path:
            console.print(f"\n[red]❌ ERROR: Binary not found: {args.binary_name}[/red]")
            console.print(f"\nSearched in: {DEFAULT_BINARY_DIR}")
            console.print("\nTip: Build binary first with:")
            console.print(f"  cargo build --release --example {args.binary_name}")
            sys.exit(1)
    else:
        binary_path = args.binary_path

    console.print(f"  [green]✓ Found: {binary_path}[/green]")

    # Step 2: Validate binary
    console.print("\n[bold]✅ Step 2: Validating binary[/bold]")

    is_valid, error_msg = validate_binary(binary_path)
    if not is_valid:
        console.print(f"\n[red]❌ ERROR: {error_msg}[/red]")
        sys.exit(1)

    file_size = binary_path.stat().st_size
    console.print(f"  [green]✓ Valid executable[/green]")
    console.print(f"  Size: {format_size(file_size)} ({file_size:,} bytes)")

    # Step 3: Generate S3 key
    console.print("\n[bold]📝 Step 3: Generating S3 key[/bold]")

    use_timestamp = not args.no_timestamp
    use_cuda = not args.no_cuda
    s3_key = generate_s3_key(binary_path, use_timestamp=use_timestamp, cuda_variant=use_cuda)

    console.print(f"  [green]✓ S3 key: {s3_key}[/green]")

    if args.dry_run:
        console.print("\n[yellow]🏃 DRY RUN: Skipping upload[/yellow]")
        console.print(f"\nWould upload to: s3://{{bucket}}/{s3_key}")
        console.print("\nRun without --dry-run to perform actual upload")
        return

    # Step 4: Initialize S3 client
    console.print("\n[bold]🔧 Step 4: Initializing S3 client[/bold]")

    try:
        config = get_config()
        console.print(f"  [green]✓ Loaded config from: {project_root}/.env.runpod[/green]")

        s3_client = S3Client(config)
        console.print(f"  [green]✓ Connected to bucket: {s3_client.bucket}[/green]")
        console.print(f"  Endpoint: {config.runpod_s3_endpoint}")
        console.print(f"  Region:   {config.runpod_s3_region}")

    except ConfigurationError as e:
        console.print(f"\n[red]❌ Configuration error: {e}[/red]")
        console.print("\nEnsure .env.runpod exists in project root with:")
        console.print("  RUNPOD_S3_ACCESS_KEY=...")
        console.print("  RUNPOD_S3_SECRET=...")
        console.print("  RUNPOD_VOLUME_ID=...")
        sys.exit(1)

    # Step 5: Upload binary
    console.print("\n[bold]📤 Step 5: Uploading binary[/bold]")

    try:
        success = upload_binary_with_progress(
            s3_client,
            binary_path,
            s3_key,
            force=args.force
        )

        if success:
            console.print("\n[bold green]✅ SUCCESS: Binary uploaded successfully![/bold green]")
            sys.exit(0)
        else:
            console.print("\n[bold red]❌ FAILED: Upload did not complete[/bold red]")
            sys.exit(1)

    except S3Error as e:
        console.print(f"\n[red]❌ S3 Error: {e}[/red]")
        sys.exit(1)
    except Exception as e:
        console.print(f"\n[red]❌ Unexpected error: {e}[/red]")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()