#!/usr/bin/env python3
"""
S3 Log Monitoring Script - Real-time training log streaming

Monitors ML training runs on RunPod with S3 log tailing.
Supports both pod-based and run-based monitoring with rich output.

Features:
- List recent pods and runs
- Stream training logs in real-time
- Monitor hyperopt trials.json updates
- Follow mode for continuous streaming
- Configurable timeout
- Color-coded output (errors, warnings, success)

Usage:
    # List recent activity
    python3 scripts/monitor_logs.py

    # Monitor specific run (by run ID)
    python3 scripts/monitor_logs.py --run-id run_20251029_145151_hyperopt --follow

    # Monitor specific pod (by pod ID)
    python3 scripts/monitor_logs.py --pod-id w4srx0tgm5hfgu --follow

    # With timeout
    python3 scripts/monitor_logs.py --run-id run_20251029_145151_hyperopt --follow --timeout 30m

Requirements:
    - .venv activation (source .venv/bin/activate)
    - Dependencies: rich, boto3, pydantic-settings
    - .env.runpod with S3 credentials
"""

import os
import sys
import re
import time
import argparse
from pathlib import Path
from typing import Optional
from datetime import datetime, timedelta

# Check for .venv activation
is_venv = hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)
if not is_venv:
    print("=" * 70)
    print("ERROR: Not running in a virtual environment (.venv)")
    print("=" * 70)
    print("This script requires .venv activation to ensure correct dependencies.")
    print()
    print("To fix this:")
    print("  1. Activate .venv:   source .venv/bin/activate")
    print("  2. Install deps:     pip install -r runpod/requirements.txt")
    print("  3. Re-run script:    python3 scripts/monitor_logs.py --help")
    print("=" * 70)
    sys.exit(1)

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

# Add project root to sys.path to ensure local runpod module is used
# This must be done BEFORE importing from runpod to avoid pip package conflict
sys.path.insert(0, str(project_root))

try:
    # Import from local runpod module (not pip-installed package)
    from runpod import PodMonitor, S3Client, RunPodConfig
    from runpod.errors import S3ObjectNotFoundError
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
    from rich.text import Text
    from rich import box
except ImportError as e:
    print(f"ERROR: Could not import required modules: {e}")
    print()
    print("To fix this:")
    print("  1. Ensure .venv is activated: source .venv/bin/activate")
    print("  2. Install dependencies: pip install -r runpod/requirements.txt")
    print()
    sys.exit(1)

# Change to project root for config loading
os.chdir(project_root)

# Set environment variable to help config find .env.runpod
os.environ.setdefault('RUNPOD_CONFIG_PATH', str(project_root / '.env.runpod'))

console = Console()


def parse_timeout(timeout_str: str) -> Optional[int]:
    """
    Parse timeout string to seconds.

    Args:
        timeout_str: Timeout string (e.g., '30m', '2h', '45s')

    Returns:
        Timeout in seconds, or None if invalid
    """
    match = re.match(r'(\d+)([smh]?)$', timeout_str.lower())
    if not match:
        return None

    value = int(match.group(1))
    unit = match.group(2) or 's'

    if unit == 'h':
        return value * 3600
    elif unit == 'm':
        return value * 60
    else:  # 's'
        return value


def list_recent_runs(s3_client: S3Client, limit: int = 20) -> list[dict]:
    """
    List recent training runs from S3.

    Args:
        s3_client: S3Client instance
        limit: Maximum number of runs to return

    Returns:
        List of run metadata dicts
    """
    try:
        # List all training run directories
        response = s3_client.s3_client.list_objects_v2(
            Bucket=s3_client.bucket,
            Prefix="ml_training/training_runs/",
            Delimiter="/"
        )

        runs = []

        # Iterate through model types (mamba2, dqn, ppo, tft)
        for prefix_obj in response.get('CommonPrefixes', []):
            model_prefix = prefix_obj['Prefix']

            # List runs for this model type
            model_response = s3_client.s3_client.list_objects_v2(
                Bucket=s3_client.bucket,
                Prefix=model_prefix,
                Delimiter="/"
            )

            for run_prefix_obj in model_response.get('CommonPrefixes', []):
                run_path = run_prefix_obj['Prefix']
                run_id = run_path.rstrip('/').split('/')[-1]

                # Extract model type
                model_type = model_prefix.rstrip('/').split('/')[-1]

                # Try to get log file metadata for last modified time
                log_key = f"{run_path}logs/training.log"
                try:
                    log_obj = s3_client.s3_client.head_object(
                        Bucket=s3_client.bucket,
                        Key=log_key
                    )
                    last_modified = log_obj['LastModified']
                    log_size = log_obj['ContentLength']
                except:
                    last_modified = None
                    log_size = 0

                # Check for trials.json (hyperopt indicator)
                trials_key = f"{run_path}hyperopt/trials.json"
                has_trials = s3_client.object_exists(trials_key)

                runs.append({
                    'run_id': run_id,
                    'model_type': model_type,
                    'last_modified': last_modified,
                    'log_size': log_size,
                    'has_trials': has_trials,
                    'log_key': log_key,
                    'trials_key': trials_key if has_trials else None
                })

        # Sort by last modified (newest first)
        # Use min datetime with UTC timezone for consistency
        from datetime import timezone
        min_dt = datetime.min.replace(tzinfo=timezone.utc)
        runs.sort(key=lambda x: x['last_modified'] or min_dt, reverse=True)

        return runs[:limit]

    except Exception as e:
        console.print(f"[red]Error listing runs: {e}[/red]")
        return []


def display_recent_runs(runs: list[dict]) -> None:
    """Display recent runs in a formatted table."""
    if not runs:
        console.print("[yellow]No recent runs found[/yellow]")
        return

    table = Table(title="Recent Training Runs", box=box.ROUNDED)
    table.add_column("Run ID", style="cyan", no_wrap=True)
    table.add_column("Model", style="magenta")
    table.add_column("Last Modified", style="green")
    table.add_column("Log Size", justify="right", style="blue")
    table.add_column("Trials", justify="center")

    for run in runs:
        run_id = run['run_id']
        model_type = run['model_type'].upper()

        if run['last_modified']:
            # Format as relative time
            now = datetime.now(run['last_modified'].tzinfo)
            delta = now - run['last_modified']

            if delta < timedelta(minutes=1):
                time_str = "just now"
            elif delta < timedelta(hours=1):
                time_str = f"{int(delta.total_seconds() / 60)}m ago"
            elif delta < timedelta(days=1):
                time_str = f"{int(delta.total_seconds() / 3600)}h ago"
            else:
                time_str = f"{delta.days}d ago"
        else:
            time_str = "unknown"

        # Format log size
        if run['log_size'] > 0:
            if run['log_size'] < 1024:
                size_str = f"{run['log_size']}B"
            elif run['log_size'] < 1024 * 1024:
                size_str = f"{run['log_size'] / 1024:.1f}KB"
            else:
                size_str = f"{run['log_size'] / (1024 * 1024):.2f}MB"
        else:
            size_str = "-"

        trials_str = "✓" if run['has_trials'] else "-"

        table.add_row(run_id, model_type, time_str, size_str, trials_str)

    console.print()
    console.print(table)
    console.print()
    console.print("[dim]Use --run-id <RUN_ID> to monitor a specific run[/dim]")
    console.print()


def stream_run_logs(
    s3_client: S3Client,
    run_id: str,
    follow: bool = True,
    timeout: Optional[int] = None,
    poll_interval: int = 5
) -> None:
    """
    Stream logs for a specific run.

    Args:
        s3_client: S3Client instance
        run_id: Run ID to monitor
        follow: Continue streaming until completion
        timeout: Maximum monitoring time in seconds
        poll_interval: Polling interval in seconds
    """
    # Find the run's log path
    try:
        # Search for run across all model types
        for model_type in ['mamba2', 'dqn', 'ppo', 'tft']:
            log_key = f"ml_training/training_runs/{model_type}/{run_id}/logs/training.log"
            if s3_client.object_exists(log_key):
                trials_key = f"ml_training/training_runs/{model_type}/{run_id}/hyperopt/trials.json"
                has_trials = s3_client.object_exists(trials_key)
                break
        else:
            console.print(f"[red]Run not found: {run_id}[/red]")
            console.print("[dim]Use 'python3 scripts/monitor_logs.py' to list available runs[/dim]")
            return
    except Exception as e:
        console.print(f"[red]Error finding run: {e}[/red]")
        return

    # Display header
    console.print()
    console.print(Panel.fit(
        f"[bold cyan]Monitoring Run: {run_id}[/bold cyan]\n"
        f"Model: [magenta]{model_type.upper()}[/magenta]\n"
        f"Log: [dim]{log_key}[/dim]\n"
        f"Hyperopt: [green]Yes[/green]" if has_trials else "[dim]No[/dim]",
        title="Log Monitor",
        border_style="cyan"
    ))
    console.print()

    # Stream logs
    log_position = 0
    start_time = time.time()
    last_trials_check = 0

    # Pattern detection
    completion_patterns = [
        "Training complete",
        "Model saved to",
        "✓ Training finished",
        "SUCCESS:",
        "Hyperparameter optimization complete"
    ]

    error_patterns = [
        "CUDA out of memory",
        "RuntimeError:",
        "AssertionError:",
        "FAILED:",
        "ERROR:",
        "panic!"
    ]

    training_complete = False
    error_detected = False

    try:
        console.print("[dim]Streaming logs... (Press Ctrl+C to stop)[/dim]")
        console.print("─" * 70)

        while True:
            # Check timeout
            if timeout and (time.time() - start_time) > timeout:
                console.print()
                console.print("[yellow]⏱️  Timeout reached[/yellow]")
                break

            # Check for new log content
            try:
                content, log_position = s3_client.tail_log_file(log_key, start_byte=log_position)

                if content:
                    text = content.decode('utf-8', errors='ignore')
                    lines = text.splitlines()

                    for line in lines:
                        # Color-code output based on content
                        if any(pattern in line for pattern in error_patterns):
                            console.print(f"[red]{line}[/red]")
                            error_detected = True
                        elif "WARN" in line.upper():
                            console.print(f"[yellow]{line}[/yellow]")
                        elif "SUCCESS" in line.upper() or "✓" in line:
                            console.print(f"[green]{line}[/green]")
                        elif any(pattern in line for pattern in completion_patterns):
                            console.print(f"[bold green]{line}[/bold green]")
                            training_complete = True
                        else:
                            console.print(line)

            except S3ObjectNotFoundError:
                if not follow:
                    console.print("[yellow]⚠ Log file not found[/yellow]")
                    return
                # Still waiting for log file to appear
                pass

            # Check trials.json updates (every 30 seconds)
            if has_trials and (time.time() - last_trials_check) > 30:
                last_trials_check = time.time()
                try:
                    trials_obj = s3_client.s3_client.get_object(
                        Bucket=s3_client.bucket,
                        Key=trials_key
                    )
                    trials_content = trials_obj['Body'].read().decode('utf-8')

                    # Parse trial count
                    import json
                    trials_data = json.loads(trials_content)
                    trial_count = len(trials_data)

                    console.print(f"[dim]📊 Hyperopt trials: {trial_count}[/dim]")

                except:
                    pass

            # Check completion
            if training_complete or error_detected:
                console.print()
                console.print("─" * 70)
                if training_complete:
                    console.print("[bold green]✅ Training completed![/bold green]")
                if error_detected:
                    console.print("[bold red]❌ Error detected in logs[/bold red]")
                break

            if not follow:
                break

            time.sleep(poll_interval)

    except KeyboardInterrupt:
        console.print()
        console.print()
        console.print("[yellow]Streaming stopped by user[/yellow]")


def stream_pod_logs(
    pod_id: str,
    config,
    follow: bool = True,
    timeout: Optional[int] = None,
    poll_interval: int = 5
) -> None:
    """
    Stream logs for a specific pod using PodMonitor.

    Args:
        pod_id: Pod ID to monitor
        config: RunPodConfig instance
        follow: Continue streaming until completion
        timeout: Maximum monitoring time in seconds
        poll_interval: Polling interval in seconds
    """
    try:
        monitor = PodMonitor(pod_id=pod_id, config=config)

        # Display pod info
        monitor.display_pod_info()

        # Stream logs
        max_lines = None
        if timeout:
            # Estimate max lines based on timeout
            max_lines = timeout * 10  # Rough estimate

        monitor.stream_s3_logs(
            follow=follow,
            poll_interval=poll_interval,
            max_lines=max_lines
        )

    except Exception as e:
        console.print(f"[red]Error monitoring pod: {e}[/red]")


def main():
    """Main execution function."""
    parser = argparse.ArgumentParser(
        description="Monitor ML training logs on RunPod S3",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List recent runs
  python3 scripts/monitor_logs.py

  # Monitor specific run (continuous streaming)
  python3 scripts/monitor_logs.py --run-id run_20251029_145151_hyperopt --follow

  # Monitor specific pod
  python3 scripts/monitor_logs.py --pod-id w4srx0tgm5hfgu --follow

  # With timeout (30 minutes)
  python3 scripts/monitor_logs.py --run-id run_20251029_145151_hyperopt --follow --timeout 30m

  # One-time log snapshot (no streaming)
  python3 scripts/monitor_logs.py --run-id run_20251029_145151_hyperopt

Requirements:
  - .venv activation (source .venv/bin/activate)
  - .env.runpod with S3 credentials
  - Dependencies: rich, boto3, pydantic-settings
        """
    )

    parser.add_argument(
        '--run-id',
        help='Run ID to monitor (e.g., run_20251029_145151_hyperopt)'
    )
    parser.add_argument(
        '--pod-id',
        help='Pod ID to monitor (e.g., w4srx0tgm5hfgu)'
    )
    parser.add_argument(
        '--follow',
        action='store_true',
        help='Continuously stream logs until completion'
    )
    parser.add_argument(
        '--timeout',
        default=None,
        help='Maximum monitoring time (e.g., 30m, 2h). Default: no timeout'
    )
    parser.add_argument(
        '--interval',
        type=int,
        default=5,
        help='Log polling interval in seconds (default: 5)'
    )
    parser.add_argument(
        '--limit',
        type=int,
        default=20,
        help='Maximum number of runs to list (default: 20)'
    )

    args = parser.parse_args()

    # Parse timeout
    timeout_seconds = None
    if args.timeout:
        timeout_seconds = parse_timeout(args.timeout)
        if timeout_seconds is None:
            console.print(f"[red]Invalid timeout format: {args.timeout}[/red]")
            console.print("[dim]Use format like '30m', '2h', or '45s'[/dim]")
            sys.exit(1)

    # Validate arguments
    if args.run_id and args.pod_id:
        console.print("[red]Error: Cannot specify both --run-id and --pod-id[/red]")
        sys.exit(1)

    try:
        # Load config from explicit path
        config = RunPodConfig.load_from_file(project_root / '.env.runpod')
        s3_client = S3Client(config)

        # Monitor specific run
        if args.run_id:
            stream_run_logs(
                s3_client=s3_client,
                run_id=args.run_id,
                follow=args.follow,
                timeout=timeout_seconds,
                poll_interval=args.interval
            )

        # Monitor specific pod
        elif args.pod_id:
            stream_pod_logs(
                pod_id=args.pod_id,
                config=config,
                follow=args.follow,
                timeout=timeout_seconds,
                poll_interval=args.interval
            )

        # List recent runs
        else:
            runs = list_recent_runs(s3_client, limit=args.limit)
            display_recent_runs(runs)

    except Exception as e:
        console.print(f"[red]Error: {e}[/red]")
        sys.exit(1)


if __name__ == "__main__":
    main()