#!/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

try:
    # Import from runpod module (clean architecture - no sys.path hacks)
    from runpod import PodMonitor, S3Client
    from runpod.config import get_config
    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
        from runpod.config import RunPodConfig
        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()