#!/usr/bin/env python3
"""
RunPod Network Volume S3 Upload Script

Uploads binaries and test data to RunPod network volume via S3 API with smart
checksum comparison to avoid unnecessary uploads.

Usage:
    ./scripts/upload_to_runpod_volume.py --binaries       # Upload binaries only
    ./scripts/upload_to_runpod_volume.py --test-data      # Upload test data only
    ./scripts/upload_to_runpod_volume.py --all            # Upload everything
    ./scripts/upload_to_runpod_volume.py --all --force    # Force re-upload (ignore checksums)
    ./scripts/upload_to_runpod_volume.py --all --dry-run  # Show what would be uploaded
"""

import argparse
import hashlib
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
from dotenv import load_dotenv
from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TaskID,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)
from rich.table import Table

console = Console()


class RunPodVolumeUploader:
    """Uploads files to RunPod network volume via S3 API."""

    def __init__(self, env_file: str = ".env.runpod"):
        """Initialize uploader with credentials from env file."""
        # Load credentials
        env_path = Path(__file__).parent.parent / env_file
        if not env_path.exists():
            console.print(f"[red]✗ Error: {env_file} not found at {env_path}[/red]")
            console.print(f"[yellow]Run this script from the project root directory[/yellow]")
            sys.exit(1)

        load_dotenv(env_path)

        # Validate required env vars
        required_vars = [
            "RUNPOD_S3_ACCESS_KEY",
            "RUNPOD_S3_SECRET",
            "RUNPOD_S3_ENDPOINT",
            "RUNPOD_S3_REGION",
            "RUNPOD_VOLUME_ID",
        ]
        missing = [var for var in required_vars if not os.getenv(var)]
        if missing:
            console.print(f"[red]✗ Missing required env vars: {', '.join(missing)}[/red]")
            sys.exit(1)

        # Initialize S3 client
        self.bucket = os.getenv("RUNPOD_VOLUME_ID")
        self.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"),
            config=Config(
                signature_version="s3v4",
                s3={"addressing_style": "path"},
                retries={"max_attempts": 3, "mode": "standard"},
            ),
        )

        console.print(f"[green]✓ Initialized S3 client[/green]")
        console.print(f"  Endpoint: {os.getenv('RUNPOD_S3_ENDPOINT')}")
        console.print(f"  Volume ID: {self.bucket}")

    def compute_md5(self, file_path: Path) -> str:
        """Compute MD5 checksum of local file."""
        md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                md5.update(chunk)
        return md5.hexdigest()

    def needs_upload(self, local_file: Path, s3_key: str) -> Tuple[bool, Optional[str]]:
        """
        Check if file needs upload based on checksum.

        Returns:
            (needs_upload, reason)
        """
        local_md5 = self.compute_md5(local_file)

        try:
            response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            remote_etag = response["ETag"].strip('"')

            if local_md5 == remote_etag:
                return False, "up-to-date (checksum match)"
            else:
                return True, "changed (checksum mismatch)"
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return True, "new file (doesn't exist on S3)"
            else:
                return True, f"error checking: {e}"

    def upload_file_with_progress(
        self, local_file: Path, s3_key: str, progress: Progress, task: TaskID
    ) -> bool:
        """Upload file with progress bar."""
        file_size = local_file.stat().st_size

        try:
            # Upload with progress callback
            def callback(bytes_transferred):
                progress.update(task, completed=bytes_transferred)

            self.s3_client.upload_file(
                str(local_file),
                self.bucket,
                s3_key,
                Callback=callback,
                ExtraArgs={"ContentType": "application/octet-stream"},
            )

            # Verify upload
            response = self.s3_client.head_object(Bucket=self.bucket, Key=s3_key)
            remote_size = response["ContentLength"]

            if remote_size == file_size:
                return True
            else:
                console.print(
                    f"[red]✗ Size mismatch: local={file_size}, remote={remote_size}[/red]"
                )
                self.s3_client.delete_object(Bucket=self.bucket, Key=s3_key)
                return False

        except ClientError as e:
            console.print(f"[red]✗ Upload failed: {e}[/red]")
            return False

    def list_volume_contents(self, prefix: str = "") -> List[Dict]:
        """List contents of volume directory."""
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
            return response.get("Contents", [])
        except ClientError as e:
            console.print(f"[red]✗ Failed to list volume: {e}[/red]")
            return []

    def create_directory(self, s3_key: str, dry_run: bool = False) -> bool:
        """Create directory on volume (by uploading empty marker object)."""
        if dry_run:
            console.print(f"[dim]Would create directory: {s3_key}[/dim]")
            return True

        try:
            # Check if directory already exists
            objects = self.list_volume_contents(prefix=s3_key)
            if objects:
                console.print(f"[dim]  Directory exists: {s3_key}[/dim]")
                return True

            # Create empty directory marker (trailing slash)
            self.s3_client.put_object(Bucket=self.bucket, Key=s3_key)
            console.print(f"[green]✓ Created directory: {s3_key}[/green]")
            return True
        except ClientError as e:
            console.print(f"[red]✗ Failed to create directory {s3_key}: {e}[/red]")
            return False

    def upload_binaries(self, force: bool = False, dry_run: bool = False) -> bool:
        """Upload training binaries to volume."""
        console.print("\n[bold cyan]Step 1: Upload Binaries[/bold cyan]")

        binaries = [
            "train_tft_parquet",
            "train_mamba2_parquet",
            "train_dqn",
            "train_ppo",
        ]

        binaries_dir = Path("target/release/examples")
        if not binaries_dir.exists():
            console.print(f"[red]✗ Binaries directory not found: {binaries_dir}[/red]")
            console.print("[yellow]Run: cargo build --release --examples[/yellow]")
            return False

        # Check which binaries exist locally
        existing_binaries = []
        for binary in binaries:
            binary_path = binaries_dir / binary
            if binary_path.exists():
                existing_binaries.append(binary)
            else:
                console.print(f"[yellow]⚠ Binary not found: {binary}[/yellow]")

        if not existing_binaries:
            console.print("[red]✗ No binaries found to upload[/red]")
            return False

        console.print(f"Found {len(existing_binaries)}/{len(binaries)} binaries")

        # Upload each binary
        stats = {"uploaded": 0, "skipped": 0, "failed": 0, "total_bytes": 0}

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            DownloadColumn(),
            TransferSpeedColumn(),
            TimeRemainingColumn(),
            console=console,
        ) as progress:
            for binary in existing_binaries:
                local_file = binaries_dir / binary
                s3_key = f"binaries/{binary}"
                file_size = local_file.stat().st_size

                # Check if upload needed
                if not force:
                    needs_upload, reason = self.needs_upload(local_file, s3_key)
                    if not needs_upload:
                        console.print(f"[green]✓ {binary}[/green]: {reason}")
                        stats["skipped"] += 1
                        continue
                    else:
                        console.print(f"[cyan]↻ {binary}[/cyan]: {reason}")

                if dry_run:
                    console.print(f"[dim]Would upload: {binary} ({file_size:,} bytes)[/dim]")
                    stats["uploaded"] += 1
                    stats["total_bytes"] += file_size
                    continue

                # Upload with progress bar
                task = progress.add_task(f"Uploading {binary}", total=file_size)
                success = self.upload_file_with_progress(local_file, s3_key, progress, task)

                if success:
                    console.print(f"[green]✓ Uploaded: {binary} ({file_size:,} bytes)[/green]")
                    stats["uploaded"] += 1
                    stats["total_bytes"] += file_size
                else:
                    console.print(f"[red]✗ Failed: {binary}[/red]")
                    stats["failed"] += 1

        # Summary
        console.print(f"\n[bold]Summary:[/bold]")
        console.print(f"  Uploaded: {stats['uploaded']}")
        console.print(f"  Skipped: {stats['skipped']}")
        console.print(f"  Failed: {stats['failed']}")
        console.print(f"  Total bytes: {stats['total_bytes']:,}")

        return stats["failed"] == 0

    def upload_test_data(self, force: bool = False, dry_run: bool = False) -> bool:
        """Upload test data files to volume."""
        console.print("\n[bold cyan]Step 2: Upload Test Data[/bold cyan]")

        test_data_dir = Path("test_data")
        if not test_data_dir.exists():
            console.print(f"[red]✗ Test data directory not found: {test_data_dir}[/red]")
            return False

        # Find all parquet files
        parquet_files = list(test_data_dir.rglob("*.parquet"))
        if not parquet_files:
            console.print("[yellow]⚠ No parquet files found[/yellow]")
            return True

        console.print(f"Found {len(parquet_files)} parquet files")

        # Upload each file
        stats = {"uploaded": 0, "skipped": 0, "failed": 0, "total_bytes": 0}

        with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            DownloadColumn(),
            TransferSpeedColumn(),
            TimeRemainingColumn(),
            console=console,
        ) as progress:
            for local_file in parquet_files:
                # Preserve directory structure
                rel_path = local_file.relative_to(test_data_dir)
                s3_key = f"test_data/{rel_path}"
                file_size = local_file.stat().st_size

                # Check if upload needed
                if not force:
                    needs_upload, reason = self.needs_upload(local_file, s3_key)
                    if not needs_upload:
                        console.print(f"[green]✓ {rel_path}[/green]: {reason}")
                        stats["skipped"] += 1
                        continue
                    else:
                        console.print(f"[cyan]↻ {rel_path}[/cyan]: {reason}")

                if dry_run:
                    console.print(f"[dim]Would upload: {rel_path} ({file_size:,} bytes)[/dim]")
                    stats["uploaded"] += 1
                    stats["total_bytes"] += file_size
                    continue

                # Upload with progress bar
                task = progress.add_task(f"Uploading {rel_path}", total=file_size)
                success = self.upload_file_with_progress(local_file, s3_key, progress, task)

                if success:
                    console.print(f"[green]✓ Uploaded: {rel_path} ({file_size:,} bytes)[/green]")
                    stats["uploaded"] += 1
                    stats["total_bytes"] += file_size
                else:
                    console.print(f"[red]✗ Failed: {rel_path}[/red]")
                    stats["failed"] += 1

        # Summary
        console.print(f"\n[bold]Summary:[/bold]")
        console.print(f"  Uploaded: {stats['uploaded']}")
        console.print(f"  Skipped: {stats['skipped']}")
        console.print(f"  Failed: {stats['failed']}")
        console.print(f"  Total bytes: {stats['total_bytes']:,}")

        return stats["failed"] == 0

    def create_output_directories(self, dry_run: bool = False) -> bool:
        """Create output directories on volume."""
        console.print("\n[bold cyan]Step 3: Create Output Directories[/bold cyan]")

        directories = [
            "models/",
            "models/tft/",
            "models/mamba2/",
            "models/dqn/",
            "models/ppo/",
            "logs/",
            "checkpoints/",
        ]

        # Directory creation is non-critical - models will create subdirs as needed
        for directory in directories:
            self.create_directory(directory, dry_run)

        # Always return success - directory creation is optional
        return True

    def verify_upload(self) -> None:
        """Verify all uploads by listing volume contents."""
        console.print("\n[bold cyan]Step 4: Verify Upload[/bold cyan]")

        # List binaries
        binaries = self.list_volume_contents(prefix="binaries/")
        console.print(f"\n[bold]Binaries:[/bold] ({len(binaries)} files)")
        for obj in binaries[:10]:  # Show first 10
            size_mb = obj["Size"] / (1024 * 1024)
            console.print(f"  {obj['Key']} ({size_mb:.1f} MB)")

        # List test data
        test_data = self.list_volume_contents(prefix="test_data/")
        console.print(f"\n[bold]Test Data:[/bold] ({len(test_data)} files)")
        for obj in test_data[:10]:  # Show first 10
            size_mb = obj["Size"] / (1024 * 1024)
            console.print(f"  {obj['Key']} ({size_mb:.1f} MB)")

        # List models directory
        models = self.list_volume_contents(prefix="models/")
        console.print(f"\n[bold]Models:[/bold] ({len(models)} files)")
        if models:
            for obj in models[:5]:
                size_mb = obj["Size"] / (1024 * 1024)
                console.print(f"  {obj['Key']} ({size_mb:.1f} MB)")


def main():
    parser = argparse.ArgumentParser(
        description="Upload binaries and test data to RunPod network volume via S3 API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Upload binaries only
  ./scripts/upload_to_runpod_volume.py --binaries

  # Upload test data only
  ./scripts/upload_to_runpod_volume.py --test-data

  # Upload everything
  ./scripts/upload_to_runpod_volume.py --all

  # Force re-upload (ignore checksums)
  ./scripts/upload_to_runpod_volume.py --all --force

  # Dry run (show what would be uploaded)
  ./scripts/upload_to_runpod_volume.py --all --dry-run
        """,
    )

    parser.add_argument(
        "--binaries", action="store_true", help="Upload training binaries"
    )
    parser.add_argument(
        "--test-data", action="store_true", help="Upload test data files"
    )
    parser.add_argument(
        "--all", action="store_true", help="Upload binaries and test data"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force re-upload (ignore checksums)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be uploaded without uploading",
    )
    parser.add_argument(
        "--env-file",
        default=".env.runpod",
        help="Path to env file (default: .env.runpod)",
    )

    args = parser.parse_args()

    # Validate arguments
    if not (args.binaries or args.test_data or args.all):
        parser.print_help()
        console.print("\n[red]Error: Must specify --binaries, --test-data, or --all[/red]")
        sys.exit(1)

    # Initialize uploader
    uploader = RunPodVolumeUploader(env_file=args.env_file)

    # Execute uploads
    success = True

    if args.all or args.binaries:
        if not uploader.upload_binaries(force=args.force, dry_run=args.dry_run):
            success = False

    if args.all or args.test_data:
        if not uploader.upload_test_data(force=args.force, dry_run=args.dry_run):
            success = False

    if args.all:
        if not uploader.create_output_directories(dry_run=args.dry_run):
            success = False

    # Verify upload
    if not args.dry_run:
        uploader.verify_upload()

    # Exit with appropriate status
    if success:
        console.print("\n[bold green]✓ Upload complete![/bold green]")
        sys.exit(0)
    else:
        console.print("\n[bold red]✗ Upload failed![/bold red]")
        sys.exit(1)


if __name__ == "__main__":
    main()