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

# Add benchmarks dir to path to import config
SCRIPT_DIR = Path(__file__).parent.resolve()
OPT_DIR = Path("/opt")

# Required environment variable pointing to a local models directory
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR")
if not LOCAL_MODEL_DIR:
    print("Error: LOCAL_MODEL_DIR environment variable is required.")
    sys.exit(1)

# Configuration file path
CONFIG_FILE = os.getenv("VLLM_CONFIG_FILE", "/etc/vllm/model_config.yaml")

# Default configuration
DEFAULT_CONFIG = {
    "default": "",
    "models": {},
    "server": {
        "host": "0.0.0.0",
        "log_level": "info"
    }
}

def detect_gpus():
    """Detects AMD GPUs via rocm-smi or /dev/dri."""
    try:
        # Try rocm-smi first
        res = subprocess.run(["rocm-smi", "--showid", "--csv"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        if res.returncode == 0:
            count = res.stdout.count("GPU")
            if count > 0: return count
    except: pass
    
    # Fallback to /dev/dri/render*
    try:
        return len(list(Path("/dev/dri").glob("renderD*")))
    except:
        return 1

def load_config():
    """Load configuration from YAML file."""
    config = DEFAULT_CONFIG.copy()
    
    if Path(CONFIG_FILE).exists():
        try:
            with open(CONFIG_FILE, "r", encoding="utf-8") as f:
                user_config = yaml.safe_load(f)
                if user_config:
                    config.update(user_config)
        except Exception as e:
            print(f"Warning: Failed to load config file: {e}")
            print("Using default configuration.")
    else:
        print(f"Warning: Config file not found at {CONFIG_FILE}")
        print("Using default configuration.")
    
    return config

def nuke_vllm_cache():
    """Removes vLLM cache directory to fix potential graph/incompatibility issues."""
    cache = Path.home() / ".cache" / "vllm"
    if cache.exists():
        try:
            print(f"Clearing vLLM cache at {cache}...", end="", flush=True)
            subprocess.run(["rm", "-rf", str(cache)], check=True)
            cache.mkdir(parents=True, exist_ok=True)
            print(" Done.")
        except Exception as e:
            print(f" Failed: {e}")

def get_local_model_path(model_name):
    """Get local model path from LOCAL_MODEL_DIR."""
    # Try multiple paths:
    # 1) LOCAL_MODEL_DIR/model_name
    # 2) case-insensitive match in LOCAL_MODEL_DIR
    
    # Exact match
    candidate_exact = os.path.join(LOCAL_MODEL_DIR, model_name)
    if os.path.isdir(candidate_exact):
        return candidate_exact
    
    # Case-insensitive match
    try:
        for entry in os.listdir(LOCAL_MODEL_DIR):
            if entry.lower() == model_name.lower():
                entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
                if os.path.isdir(entry_path):
                    return entry_path
    except Exception as e:
        print(f"Error searching for model: {e}")
    
    return None

def verify_model_path(model_path):
    """Verify that the model path contains required files."""
    required = ["config.json", "pytorch_model.bin", "model.safetensors"]
    found = any(os.path.isfile(os.path.join(model_path, f)) for f in required)
    if not found:
        print(f"Error: local model dir {model_path} missing expected files {required}")
        return False
    return True

def main():
    # Load configuration
    config = load_config()
    
    # Get model name from command line or use default
    model_name = None
    if len(sys.argv) > 1:
        model_name = sys.argv[1]
    
    if not model_name:
        model_name = config.get("default", "")
        if not model_name:
            print("Error: Either specify a model name as argument or set 'default' in config.")
            sys.exit(1)
    
    # Get model configuration
    model_config = config.get("models", {}).get(model_name)
    if not model_config:
        print(f"Error: Model '{model_name}' not found in configuration.")
        sys.exit(1)
    
    # Get model path
    model_path = model_config.get("path")
    if not model_path:
        # Fallback to LOCAL_MODEL_DIR if path not specified
        model_path = get_local_model_path(model_name)
        if not model_path:
            print(f"Error: model '{model_name}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
            sys.exit(1)
    
    # Verify model path
    if not verify_model_path(model_path):
        sys.exit(1)
    
    # Detect GPU count
    gpu_count = detect_gpus()
    
    # Get server configuration
    server_config = config.get("server", {})
    host = server_config.get("host", "0.0.0.0")
    
    # Build command
    cmd = ["vllm", "serve", model_path]
    
    # Add server parameters
    cmd.extend(["--host", host])
    if "port" in model_config:
        cmd.extend(["--port", str(model_config["port"])])
    
    # Add API key for OpenAI compatibility
    api_key = model_config.get("api_key")
    if api_key:
        cmd.extend(["--api-key", api_key])
    
    # Add model parameters
    if "tensor_parallel_size" in model_config:
        tp_size = min(model_config["tensor_parallel_size"], gpu_count)
        cmd.extend(["--tensor-parallel-size", str(tp_size)])
    
    if "max_num_seqs" in model_config:
        cmd.extend(["--max-num-seqs", str(model_config["max_num_seqs"])])
    
    if "max_model_len" in model_config:
        cmd.extend(["--max-model-len", str(model_config["max_model_len"])])
    
    if "gpu_memory_utilization" in model_config:
        cmd.extend(["--gpu-memory-utilization", str(model_config["gpu_memory_utilization"])])
    
    if "dtype" in model_config:
        cmd.extend(["--dtype", model_config["dtype"]])
    else:
        cmd.extend(["--dtype", "auto"])
    
    if "max_num_batched_tokens" in model_config:
        cmd.extend(["--max-num-batched-tokens", str(model_config["max_num_batched_tokens"])])
    
    if "block_size" in model_config:
        cmd.extend(["--block-size", str(model_config["block_size"])])
    
    if "swap_space" in model_config:
        cmd.extend(["--swap-space", str(model_config["swap_space"])])
    
    if model_config.get("enforce_eager", False):
        cmd.append("--enforce-eager")
    
    if model_config.get("trust_remote", False):
        cmd.append("--trust-remote-code")
    
    # Set environment variables
    env = os.environ.copy()
    
    if model_config.get("use_rocm_attn", False):
        env["VLLM_V1_USE_PREFILL_DECODE_ATTENTION"] = "1"
        env["VLLM_USE_TRITON_FLASH_ATTN"] = "0"
    
    # Clear cache if requested
    if model_config.get("clear_cache", False):
        nuke_vllm_cache()
    
    # Print configuration
    print("\n" + "="*60)
    print(f" Launching: {model_name}")
    print(f" Model Path: {model_path}")
    print(f" Host:      {host}")
    if "port" in model_config:
        print(f" Port:      {model_config['port']}")
    if api_key:
        print(f" API Key:   {api_key[:8]}...{api_key[-4:]}")
    if "tensor_parallel_size" in model_config:
        print(f" TP Size:   {tp_size}")
    if "max_num_seqs" in model_config:
        print(f" Max Seqs:  {model_config['max_num_seqs']}")
    if "max_model_len" in model_config:
        print(f" Max Ctx:   {model_config['max_model_len']}")
    if "gpu_memory_utilization" in model_config:
        print(f" GPU Util:  {model_config['gpu_memory_utilization']}")
    print(f" Command:   {' '.join(cmd)}")
    print("="*60 + "\n")
    
    # Launch vLLM server
    os.execvpe("vllm", cmd, env)

if __name__ == "__main__":
    main()