#!/usr/bin/env python3
import sys
import os
import json
import shutil
import subprocess
from pathlib import Path

# Configuration
SCRIPT_DIR = Path("/opt/script")
CONFIG_PATH = Path("/config/config.json")
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/opt/model")
HOST = os.getenv("HOST", "0.0.0.0")
PORT = os.getenv("PORT", "8000")

def log(msg):
    """Print log message with timestamp"""
    print(f"[START-VLLM] {msg}", flush=True)

def load_config():
    """Load configuration from config.json"""
    log(f"Loading config from {CONFIG_PATH}")
    
    if not CONFIG_PATH.exists():
        log(f"ERROR: Config file not found at {CONFIG_PATH}")
        sys.exit(1)
    
    try:
        with open(CONFIG_PATH, "r") as f:
            config_data = json.load(f)
        
        model_table = config_data["models"]
        default_model = config_data["default_model"]
        models_to_run = list(model_table.keys())
        
        log(f"Loaded {len(models_to_run)} models from config")
        log(f"Default model: {default_model}")
        
        return model_table, default_model, models_to_run
    except Exception as e:
        log(f"ERROR: Failed to load config: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def detect_gpus():
    """Detect AMD GPUs"""
    try:
        result = subprocess.run(
            ["rocm-smi", "--showid", "--csv"],
            capture_output=True,
            text=True,
            timeout=10
        )
        if result.returncode == 0:
            count = result.stdout.count("GPU")
            if count > 0:
                return count
    except Exception as e:
        log(f"Warning: rocm-smi failed: {e}")
    
    # Fallback to /dev/dri
    try:
        render_devices = list(Path("/dev/dri").glob("renderD*"))
        if render_devices:
            return len(render_devices)
    except Exception:
        pass
    
    log("Warning: Could not detect GPUs, assuming 1 GPU")
    return 1

def find_model_path(model_id):
    """Find local model path"""
    log(f"Looking for model: {model_id}")
    log(f"LOCAL_MODEL_DIR: {LOCAL_MODEL_DIR}")
    
    if not os.path.exists(LOCAL_MODEL_DIR):
        log(f"ERROR: LOCAL_MODEL_DIR does not exist: {LOCAL_MODEL_DIR}")
        return None
    
    # Try exact match
    candidate = os.path.join(LOCAL_MODEL_DIR, model_id)
    if os.path.isdir(candidate):
        log(f"Found model at: {candidate}")
        return candidate
    
    # Try without prefix
    repo_name = model_id.split('/')[-1] if '/' in model_id else model_id
    candidate = os.path.join(LOCAL_MODEL_DIR, repo_name)
    if os.path.isdir(candidate):
        log(f"Found model at: {candidate}")
        return candidate
    
    # Try case-insensitive match
    try:
        for entry in os.listdir(LOCAL_MODEL_DIR):
            if entry.lower() == repo_name.lower():
                entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
                if os.path.isdir(entry_path):
                    log(f"Found model at: {entry_path}")
                    return entry_path
    except Exception as e:
        log(f"ERROR: Failed to list directory: {e}")
    
    log(f"ERROR: Model not found: {model_id}")
    log(f"Available models: {os.listdir(LOCAL_MODEL_DIR)}")
    return None

def launch_model(model_id, config, model_path, gpu_count):
    """Launch vLLM server"""
    log(f"Launching model: {model_id}")
    
    # Get configuration
    valid_tp = config.get("valid_tp", [1])
    max_tp = max(valid_tp) if valid_tp else 1
    tp_size = min(gpu_count, max_tp)
    
    ctx = int(config.get("ctx", 8192))
    max_seqs = int(config.get("max_num_seqs", 64))
    gpu_util = float(config.get("gpu_util", 0.98))
    
    log(f"Config: TP={tp_size}, Ctx={ctx}, Seqs={max_seqs}, Util={gpu_util}")
    
    # Build command
    cmd = [
        "vllm", "serve", model_path,
        "--host", HOST,
        "--port", PORT,
        "--tensor-parallel-size", str(tp_size),
        "--max-num-seqs", str(max_seqs),
        "--max-model-len", str(ctx),
        "--gpu-memory-utilization", str(gpu_util),
        "--dtype", "auto"
    ]
    
    if config.get("trust_remote"):
        cmd.append("--trust-remote-code")
    
    if config.get("enforce_eager"):
        cmd.append("--enforce-eager")
    
    log(f"Command: {' '.join(cmd)}")
    
    # Set environment
    env = os.environ.copy()
    env.update(config.get("env", {}))
    
    # Launch vLLM
    log("Starting vLLM server...")
    try:
        result = subprocess.run(cmd, env=env)
        if result.returncode != 0:
            log(f"ERROR: vLLM exited with code {result.returncode}")
            sys.exit(result.returncode)
    except Exception as e:
        log(f"ERROR: Failed to start vLLM: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def main():
    """Main entry point"""
    log("Starting vLLM launcher...")
    
    # Load configuration
    model_table, default_model, models_to_run = load_config()
    
    # Detect GPUs
    gpu_count = detect_gpus()
    log(f"Detected {gpu_count} GPU(s)")
    
    # Check if we should use default model
    use_default = os.getenv("USE_DEFAULT_MODEL", "false").lower() == "true"
    
    if use_default:
        log("Using default model mode")
        model_id = default_model
    else:
        # Interactive mode - for now just use default
        log("Interactive mode not supported in container, using default model")
        model_id = default_model
    
    # Check if model is in config
    if model_id not in model_table:
        log(f"ERROR: Model {model_id} not found in config")
        sys.exit(1)
    
    config = model_table[model_id]
    
    # Find model path
    model_path = find_model_path(model_id)
    if not model_path:
        log("ERROR: Could not find local model. Offline mode is active.")
        sys.exit(1)
    
    # Launch model
    launch_model(model_id, config, model_path, gpu_count)

if __name__ == "__main__":
    main()