#!/usr/bin/env python
"""
Checkpoint Monitor for VLM Evaluation

Monitors output_models directory for new checkpoints, validates completeness,
classifies model type (AR vs MDM), and submits Slurm evaluation jobs.

Usage:
    nohup python chengyue_script/dvlm/evaluation/monitor_checkpoints.py \
        --watch_dir output_models \
        --poll_interval 60 \
        > monitor.log 2>&1 &
"""

import argparse
import json
import os
import subprocess
import time
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, Tuple, Set


# ==============================================================================
# Configuration
# ==============================================================================

PROJECT_ROOT = "/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"
CONDA_ENV = "shiyi-vila-eval"
DEFAULT_TASKS = "SEEDBench2_Plus AI2D_TEST RealWorldQA DocVQA_VAL ChartQA_TEST MathVista_MINI MMMU_DEV_VAL MMBench_DEV_EN_V11 GQA_TestDev_Balanced TextVQA_VAL POPE"

SLURM_TEMPLATE = """#!/bin/bash
#SBATCH -A llmservice_fm_vision
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 4:00:00
#SBATCH -N 1
#SBATCH -J eval_{job_name}
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=512G
#SBATCH -o {output_dir}/slurm_out.log
#SBATCH -e {output_dir}/slurm_err.log

cd {project_root}
source ~/.bashrc
conda activate {conda_env} || true

# Run VLMEvalKit evaluation
MODEL_PATH="{model_path}" \\
NUM_GPUS=8 \\
TASKS="{tasks}" \\
OUTPUT_BASE="{output_dir}" \\
USE_MDM={use_mdm} \\
bash chengyue_script/dvlm/evaluation/vlmeval/run_all_benchmarks.sh

# Mark evaluation as complete
touch "{output_dir}/.eval_done"
"""


# ==============================================================================
# Logging
# ==============================================================================

def log(msg: str, level: str = "INFO"):
    """Print timestamped log message."""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{level}] {msg}", flush=True)


# ==============================================================================
# Checkpoint Completeness Validation
# ==============================================================================

def check_trainer_state(checkpoint_path: Path) -> Optional[bool]:
    """Check if trainer_state.json exists and is valid."""
    trainer_state = checkpoint_path / "trainer_state.json"
    if not trainer_state.exists():
        return None  # File doesn't exist, can't determine
    try:
        with open(trainer_state) as f:
            data = json.load(f)
        # Check for expected fields
        if "global_step" in data and "epoch" in data:
            return True
        return False
    except (json.JSONDecodeError, IOError):
        return False


def check_done_file(checkpoint_path: Path) -> Optional[bool]:
    """Check if .done file exists."""
    done_file = checkpoint_path / ".done"
    if done_file.exists():
        return True
    return None  # File doesn't exist, can't determine


def check_file_stability(checkpoint_path: Path, wait_seconds: int = 30) -> bool:
    """Check if files haven't changed in the last N seconds."""
    def get_total_size():
        total = 0
        for f in checkpoint_path.rglob("*"):
            if f.is_file():
                try:
                    total += f.stat().st_size
                except OSError:
                    pass
        return total
    
    initial_size = get_total_size()
    time.sleep(wait_seconds)
    final_size = get_total_size()
    
    return initial_size == final_size and initial_size > 0


def check_model_files(checkpoint_path: Path) -> bool:
    """Check if model files (safetensors or bin) exist."""
    safetensors = list(checkpoint_path.glob("*.safetensors"))
    bin_files = list(checkpoint_path.glob("*.bin"))
    config_file = checkpoint_path / "config.json"
    return (len(safetensors) > 0 or len(bin_files) > 0) and config_file.exists()


def validate_checkpoint_completeness(checkpoint_path: Path, max_retries: int = 5) -> bool:
    """
    Validate checkpoint completeness using multiple methods.
    Returns True if checkpoint is complete, False otherwise.
    
    For FSDP checkpoints that don't have trainer_state.json:
    - Check model files exist (safetensors/bin + config.json)
    - Check file stability (10 seconds)
    """
    log(f"Validating completeness for: {checkpoint_path}")
    
    for attempt in range(max_retries):
        # Check if model files exist
        has_model_files = check_model_files(checkpoint_path)
        
        if not has_model_files:
            log(f"  No model files found (attempt {attempt + 1}/{max_retries})")
            time.sleep(10)
            continue
        
        # Check trainer_state.json (optional for FSDP)
        trainer_state_result = check_trainer_state(checkpoint_path)
        
        # Check .done file (optional)
        done_file_result = check_done_file(checkpoint_path)
        
        # Check file stability (10 seconds)
        log("  Checking file stability (10s)...")
        stability_result = check_file_stability(checkpoint_path, wait_seconds=10)
        
        # Print status
        status_parts = []
        status_parts.append(f"model_files: {'✓' if has_model_files else '✗'}")
        status_parts.append(f"trainer_state: {'✓' if trainer_state_result else '?' if trainer_state_result is None else '✗'}")
        status_parts.append(f"done_file: {'✓' if done_file_result else '?'}")
        status_parts.append(f"file_stability: {'✓' if stability_result else '✗'}")
        log(f"  Status: {', '.join(status_parts)}")
        
        # Complete if: model files exist AND file stability passed
        if has_model_files and stability_result:
            log("  ✓ Checkpoint is COMPLETE")
            return True
        elif not stability_result:
            log(f"  Files still changing (attempt {attempt + 1}/{max_retries}), waiting...")
            time.sleep(10)
        else:
            log("  ✗ Checkpoint is INCOMPLETE")
            return False
    
    log("  ✗ Max retries reached, checkpoint incomplete")
    return False


# ==============================================================================
# Model Type Classification
# ==============================================================================

def classify_by_folder_name(checkpoint_path: Path) -> Optional[str]:
    """Classify model type by folder name."""
    path_lower = str(checkpoint_path).lower()
    if "mdm" in path_lower:
        return "MDM"
    elif "ar" in path_lower:
        return "AR"
    return None


def classify_by_config(checkpoint_path: Path) -> Optional[str]:
    """Classify model type by config.json content."""
    config_path = checkpoint_path / "config.json"
    if not config_path.exists():
        return None
    
    try:
        with open(config_path) as f:
            config = json.load(f)
        
        # Check architectures field
        architectures = config.get("architectures", [])
        for arch in architectures:
            arch_lower = arch.lower()
            if "mdm" in arch_lower:
                return "MDM"
            elif "fast_dllm" in arch_lower:
                # Fast_dLLM without MDM is AR
                return "AR"
        
        # Check model_type field
        model_type = config.get("model_type", "").lower()
        if "mdm" in model_type:
            return "MDM"
        
        # Check auto_map for custom classes
        auto_map = config.get("auto_map", {})
        for key, value in auto_map.items():
            if isinstance(value, str):
                value_lower = value.lower()
                if "mdm" in value_lower:
                    return "MDM"
        
        return "AR"  # Default to AR if Fast_dLLM model
        
    except (json.JSONDecodeError, IOError):
        return None


def classify_model_type(checkpoint_path: Path, max_retries: int = 10) -> Optional[str]:
    """
    Classify model as AR or MDM using multiple methods.
    Returns "AR", "MDM", or None if classification fails.
    
    Priority: If folder name contains "mdm", it's MDM (folder name takes priority).
    Otherwise use config.json or folder name.
    """
    return "MDM"
    log(f"Classifying model type for: {checkpoint_path}")
    
    # Method 1: Folder name
    folder_result = classify_by_folder_name(checkpoint_path)
    
    # Method 2: config.json
    config_result = classify_by_config(checkpoint_path)
    
    # Print status
    log(f"  folder_name: {folder_result or '?'}, config.json: {config_result or '?'}")
    
    # Priority: folder name MDM takes precedence (config.json may not reflect MDM training)
    if folder_result == "MDM":
        log(f"  ✓ Model type: MDM (folder name priority)")
        return "MDM"
    
    # Otherwise use config.json if available
    if config_result is not None:
        log(f"  ✓ Model type: {config_result}")
        return config_result
    
    # Fallback to folder name
    if folder_result is not None:
        log(f"  ✓ Model type: {folder_result}")
        return folder_result
    
    log("  ✗ Could not classify model type")
    return None


# ==============================================================================
# Slurm Job Submission
# ==============================================================================

def submit_eval_job(
    checkpoint_path: Path,
    model_type: str,
    tasks: str = DEFAULT_TASKS,
) -> Optional[str]:
    """
    Generate and submit Slurm evaluation job.
    Returns job ID if successful, None otherwise.
    
    Results are stored in model_root/eval_vlmeval/checkpoint-XXX/
    """
    # Store results in model root, not in checkpoint directory
    model_root = checkpoint_path.parent
    ckpt_name = checkpoint_path.name
    output_dir = model_root / "eval_vlmeval" / ckpt_name
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Generate job name from checkpoint path
    parent_name = model_root.name
    job_name = f"{model_type}_{parent_name}_{ckpt_name}"[:50]  # Truncate for Slurm
    
    # Determine if MDM model
    use_mdm = "1" if model_type == "MDM" else "0"
    
    # Generate Slurm script
    script_content = SLURM_TEMPLATE.format(
        job_name=job_name,
        project_root=PROJECT_ROOT,
        conda_env=CONDA_ENV,
        model_path=str(checkpoint_path),
        tasks=tasks,
        output_dir=str(output_dir),
        use_mdm=use_mdm,
    )
    
    # Save script
    script_path = output_dir / "run_eval.sh"
    with open(script_path, "w") as f:
        f.write(script_content)
    
    log(f"Generated Slurm script: {script_path}")
    
    # Submit job
    try:
        result = subprocess.run(
            ["sbatch", str(script_path)],
            capture_output=True,
            text=True,
            timeout=30,
        )
        
        if result.returncode == 0:
            # Parse job ID from output like "Submitted batch job 12345"
            output = result.stdout.strip()
            job_id = output.split()[-1] if output else "unknown"
            log(f"✓ Submitted Slurm job: {job_id}")
            
            # Record submission
            submission_record = output_dir / ".submitted"
            with open(submission_record, "w") as f:
                json.dump({
                    "job_id": job_id,
                    "submitted_at": datetime.now().isoformat(),
                    "model_type": model_type,
                    "tasks": tasks,
                }, f, indent=2)
            
            return job_id
        else:
            log(f"✗ sbatch failed: {result.stderr}", level="ERROR")
            return None
            
    except subprocess.TimeoutExpired:
        log("✗ sbatch timed out", level="ERROR")
        return None
    except FileNotFoundError:
        log("✗ sbatch not found (not on Slurm cluster?)", level="ERROR")
        return None


# ==============================================================================
# Checkpoint Discovery
# ==============================================================================

def _check_and_add_checkpoint(ckpt_dir: Path, processed: Set[str], new_checkpoints: list):
    """Helper to check if checkpoint should be added to new_checkpoints list."""
    if not ckpt_dir.is_dir():
        return
    if not ckpt_dir.name.startswith("checkpoint-"):
        return
    
    ckpt_key = str(ckpt_dir)
    if ckpt_key in processed:
        return
    
    # Check if already evaluated
    # New path: model_root/eval_vlmeval/checkpoint-XXX/.eval_done
    model_root = ckpt_dir.parent
    ckpt_name = ckpt_dir.name
    eval_done = model_root / "eval_vlmeval" / ckpt_name / ".eval_done"
    submitted = model_root / "eval_vlmeval" / ckpt_name / ".submitted"
    
    # Also check old paths for backwards compatibility
    eval_done_old = ckpt_dir / "eval_vlmeval" / ".eval_done"
    submitted_old = ckpt_dir / "eval_vlmeval" / ".submitted"
    eval_done_old2 = ckpt_dir / "eval_results" / ".eval_done"
    submitted_old2 = ckpt_dir / "eval_results" / ".submitted"
    
    if any([eval_done.exists(), submitted.exists(), 
            eval_done_old.exists(), submitted_old.exists(),
            eval_done_old2.exists(), submitted_old2.exists()]):
        processed.add(ckpt_key)
        return
    
    new_checkpoints.append(ckpt_dir)


def find_new_checkpoints(watch_dir: Path, processed: Set[str]) -> list:
    """
    Find checkpoint directories that haven't been processed yet.
    
    Supports two directory structures:
    1. Single model: watch_dir/checkpoint-*
    2. Multi model: watch_dir/model_name/checkpoint-*
    """
    new_checkpoints = []
    
    # Check if watch_dir directly contains checkpoint-* (single model mode)
    has_direct_checkpoints = any(
        d.name.startswith("checkpoint-") and d.is_dir() 
        for d in watch_dir.iterdir()
    )
    
    if has_direct_checkpoints:
        # Single model mode: watch_dir/checkpoint-*
        for ckpt_dir in watch_dir.iterdir():
            _check_and_add_checkpoint(ckpt_dir, processed, new_checkpoints)
    else:
        # Multi model mode: watch_dir/model_name/checkpoint-*
        for model_dir in watch_dir.iterdir():
            if not model_dir.is_dir():
                continue
            
            for ckpt_dir in model_dir.iterdir():
                _check_and_add_checkpoint(ckpt_dir, processed, new_checkpoints)
    
    return new_checkpoints


# ==============================================================================
# Main Monitor Loop
# ==============================================================================

def monitor_loop(
    watch_dir: Path,
    poll_interval: int = 60,
    tasks: str = DEFAULT_TASKS,
):
    """Main monitoring loop."""
    log(f"Starting checkpoint monitor")
    log(f"  Watch directory: {watch_dir}")
    log(f"  Poll interval: {poll_interval}s")
    log(f"  Tasks: {tasks}")
    log("=" * 60)
    
    processed: Set[str] = set()
    
    while True:
        try:
            # Find new checkpoints
            new_checkpoints = find_new_checkpoints(watch_dir, processed)
            
            if new_checkpoints:
                log(f"Found {len(new_checkpoints)} new checkpoint(s)")
            
            for ckpt_path in new_checkpoints:
                log("-" * 40)
                log(f"Processing: {ckpt_path}")
                
                # Validate completeness
                if not validate_checkpoint_completeness(ckpt_path):
                    log(f"Skipping incomplete checkpoint: {ckpt_path}")
                    continue
                
                # Classify model type
                model_type = classify_model_type(ckpt_path)
                if model_type is None:
                    log(f"Skipping unclassified checkpoint: {ckpt_path}", level="WARN")
                    processed.add(str(ckpt_path))
                    continue
                
                # Submit evaluation job
                job_id = submit_eval_job(ckpt_path, model_type, tasks)
                if job_id:
                    log(f"✓ Evaluation queued for {ckpt_path}")
                else:
                    log(f"✗ Failed to queue evaluation for {ckpt_path}", level="ERROR")
                
                processed.add(str(ckpt_path))
            
            # Wait before next poll
            time.sleep(poll_interval)
            
        except KeyboardInterrupt:
            log("Monitor stopped by user")
            break
        except Exception as e:
            log(f"Error in monitor loop: {e}", level="ERROR")
            time.sleep(poll_interval)


# ==============================================================================
# CLI
# ==============================================================================

def main():
    parser = argparse.ArgumentParser(
        description="Monitor checkpoints and auto-submit evaluation jobs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # Run in foreground
    python monitor_checkpoints.py --watch_dir output_models
    
    # Run in background
    nohup python monitor_checkpoints.py --watch_dir output_models > monitor.log 2>&1 &
        """,
    )
    
    parser.add_argument(
        "--watch_dir",
        type=str,
        default="output_models",
        help="Directory to monitor for new checkpoints",
    )
    parser.add_argument(
        "--poll_interval",
        type=int,
        default=60,
        help="Seconds between directory scans (default: 60)",
    )
    parser.add_argument(
        "--tasks",
        type=str,
        default=DEFAULT_TASKS,
        help=f"Comma-separated evaluation tasks (default: {DEFAULT_TASKS})",
    )
    
    args = parser.parse_args()
    
    watch_dir = Path(args.watch_dir)
    if not watch_dir.exists():
        log(f"Watch directory does not exist: {watch_dir}", level="ERROR")
        return 1
    
    monitor_loop(
        watch_dir=watch_dir,
        poll_interval=args.poll_interval,
        tasks=args.tasks,
    )
    
    return 0


if __name__ == "__main__":
    exit(main())

