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

# Add script dir to path
SCRIPT_DIR = Path("/opt/script")
OPT_DIR = Path("/opt")

# Config file path (check container path first, then local path)
CONFIG_PATH = Path("/config/config.json")
if not CONFIG_PATH.exists():
    CONFIG_PATH = Path("/config.json")

# Local model directory (container path)
LOCAL_MODEL_DIR = os.getenv("LOCAL_MODEL_DIR", "/opt/model")

# Load configuration from config.json
print(f"DEBUG: Looking for config at {CONFIG_PATH}")
print(f"DEBUG: Config file exists: {CONFIG_PATH.exists()}")
if not CONFIG_PATH.exists():
    print(f"ERROR: Config file not found at {CONFIG_PATH}")
    print(f"ERROR: Please mount config.json to /config/config.json")
    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())
    print(f"DEBUG: Loaded {len(MODELS_TO_RUN)} models from config")
    print(f"DEBUG: Default model: {DEFAULT_MODEL}")
except Exception as e:
    print(f"Error: Could not load config.json: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# Results file path
if (OPT_DIR / "max_context_results.json").exists():
    RESULTS_FILE = OPT_DIR / "max_context_results.json"
else:
    RESULTS_FILE = SCRIPT_DIR / "max_context_results.json"
HOST = os.getenv("HOST", "0.0.0.0")
PORT = os.getenv("PORT", "8000")

def check_dependencies():
    if not shutil.which("dialog"):
        print("Error: 'dialog' is required. Please install it (apt-get install dialog).")
        sys.exit(1)

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 get_verified_config(model_id, tp_size, max_seqs):
    """
    Reads max_context_results.json to find the best verified configuration.
    Returns dict: {'ctx': int, 'util': float}
    """
    default_config = {
        "ctx": int(MODEL_TABLE.get(model_id, {}).get("ctx", 8192)),
        "util": 0.90 # Safe default
    }
    
    if not RESULTS_FILE.exists():
        return default_config

    try:
        with open(RESULTS_FILE, "r") as f:
            data = json.load(f)
            
        # Filter for Model + TP + Sequences
        matches = [r for r in data 
                  if r["model"] == model_id 
                  and r["tp"] == tp_size 
                  and r["max_seqs"] == max_seqs 
                  and r["status"] == "success"]
        
        if not matches:
            # Fallback 1: Try finding match with SAME TP but ANY Sequences (e.g. 1) to get base context?
            # Actually, safer to fallback to default or try finding nearest sequence?
            # Let's try finding exact match first. If fail, return default.
            return default_config
            
        # Sort by Util desc, then Context desc
        # We prefer higher utilization if available (performance), as long as it is verified success
        matches.sort(key=lambda x: (float(x["util"]), x["max_context_1_user"]), reverse=True)
        
        best = matches[0]
        return {
            "ctx": best["max_context_1_user"],
            "util": float(best["util"])
        }
        
    except Exception as e:
        return default_config

def run_dialog(args):
    """Runs dialog and returns stderr (selection)."""
    with tempfile.NamedTemporaryFile(mode="w+") as tf:
        cmd = ["dialog"] + args
        try:
            subprocess.run(cmd, stderr=tf, check=True)
            tf.seek(0)
            return tf.read().strip()
        except subprocess.CalledProcessError:
            return None # User cancelled

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.")
            time.sleep(1)
        except Exception as e:
            print(f" Failed: {e}")

def configure_and_launch(model_idx, gpu_count):
    print(f"DEBUG: configure_and_launch called with model_idx={model_idx}, gpu_count={gpu_count}")
    model_id = MODELS_TO_RUN[model_idx]
    config = MODEL_TABLE[model_id]
    print(f"DEBUG: model_id={model_id}")
    print(f"DEBUG: LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
    
    # Determine whether we have a local copy to serve. Try multiple fallbacks:
    # 1) LOCAL_MODEL_DIR/<owner>/<repo>
    # 2) LOCAL_MODEL_DIR/<repo>
    # 3) case-insensitive match of <repo> in LOCAL_MODEL_DIR
    model_path = model_id
    print(f"DEBUG: Starting model path lookup...")
    if LOCAL_MODEL_DIR:
        print(f"DEBUG: LOCAL_MODEL_DIR is set to: {LOCAL_MODEL_DIR}")
        print(f"DEBUG: LOCAL_MODEL_DIR exists: {os.path.exists(LOCAL_MODEL_DIR)}")
        if os.path.exists(LOCAL_MODEL_DIR):
            print(f"DEBUG: LOCAL_MODEL_DIR contents: {os.listdir(LOCAL_MODEL_DIR)}")
        
        # Full repo path (owner/repo)
        candidate_full = os.path.join(LOCAL_MODEL_DIR, model_id)
        print(f"DEBUG: Checking candidate_full: {candidate_full}")
        print(f"DEBUG: candidate_full exists: {os.path.isdir(candidate_full)}")
        if os.path.isdir(candidate_full):
            model_path = candidate_full
            print(f"DEBUG: Found model at: {model_path}")
        else:
            # Repo-name only (last segment)
            repo_name = model_id.split('/')[-1]
            candidate_repo = os.path.join(LOCAL_MODEL_DIR, repo_name)
            print(f"DEBUG: Checking candidate_repo: {candidate_repo}")
            print(f"DEBUG: candidate_repo exists: {os.path.isdir(candidate_repo)}")
            if os.path.isdir(candidate_repo):
                model_path = candidate_repo
                print(f"DEBUG: Found model at: {model_path}")
            else:
                # Fallback: try to find a directory in LOCAL_MODEL_DIR that matches repo_name case-insensitively
                print(f"DEBUG: Trying case-insensitive match for: {repo_name}")
                try:
                    for entry in os.listdir(LOCAL_MODEL_DIR):
                        print(f"DEBUG: Checking entry: {entry}")
                        if entry.lower() == repo_name.lower():
                            entry_path = os.path.join(LOCAL_MODEL_DIR, entry)
                            if os.path.isdir(entry_path):
                                model_path = entry_path
                                print(f"DEBUG: Found model at: {model_path}")
                                break
                except Exception as e:
                    print(f"DEBUG: Exception during case-insensitive lookup: {e}")
    
    print(f"DEBUG: Final model_path: {model_path}")
    print(f"DEBUG: model_path == model_id: {model_path == model_id}")
    
    # if LOCAL_MODEL_DIR is specified, refuse to fall back to remote
    if LOCAL_MODEL_DIR and model_path == model_id:
        print(f"Error: model '{model_id}' not found under LOCAL_MODEL_DIR={LOCAL_MODEL_DIR}")
        print("Off‑line mode active; network downloads are disabled.")
        sys.exit(1)
    
    # Static Config
    valid_tps = config.get("valid_tp", [1])
    max_tp = max(valid_tps) if valid_tps else 1
    
    # Defaults
    current_tp = min(gpu_count, max_tp)
    current_seqs = 1 # Default to 1 concurrent user/request for stability
    
    # Initial Lookup
    verified = get_verified_config(model_id, current_tp, current_seqs)
    current_ctx = verified["ctx"]
    current_util = verified["util"]
    
    clear_cache = False
    use_eager = config.get("enforce_eager", False) # Default to model config, usually False
    use_rocm_attn = False # Default to Triton
    
    name = model_id.split("/")[-1]
    
    while True:
        cache_status = "YES" if clear_cache else "NO"
        eager_status = "YES" if use_eager else "NO"
        attn_backend = "ROCm" if use_rocm_attn else "Triton"
        
        menu_args = [
            "--clear", "--backtitle", f"AMD R9700 vLLM Launcher (GPUs: {gpu_count})",
            "--title", f"Configuration: {name}",
            "--menu", "Customize Launch Parameters:", "22", "65", "9",
            "1", f"Tensor Parallelism:   {current_tp}",
            "2", f"Concurrent Requests:  {current_seqs}",
            "3", f"Context Length:       {current_ctx} (Verified)",
            "4", f"GPU Utilization:      {current_util} (Verified)",
            "5", f"Attention Backend:    {attn_backend}",
            "6", f"Erase vLLM Cache:     {cache_status}",
            "7", f"Force Eager Mode:     {eager_status}",
            "8", "LAUNCH SERVER"
        ]
        
        choice = run_dialog(menu_args)
        if not choice: return False # Back/Cancel
        
        if choice == "1":
            # TP Selection
            new_tp = run_dialog([
                "--title", "Tensor Parallelism",
                "--rangebox", f"Set TP Size (1-{max_tp})", "10", "40", "1", str(max_tp), str(current_tp)
            ])
            if new_tp: 
                new_tp_int = int(new_tp)
                if new_tp_int != current_tp:
                    current_tp = new_tp_int
                    # RE-CALCULATE Config
                    verified = get_verified_config(model_id, current_tp, current_seqs)
                    current_ctx = verified["ctx"]
                    current_util = verified["util"]
            
        elif choice == "2":
            # Max Seqs Selection
            new_seqs = run_dialog([
                "--title", "Concurrent Requests",
                "--menu", "Select Max Concurrent Requests:", "12", "40", "4",
                "1", "1 (Latency Focus)",
                "4", "4 (Balanced)",
                "8", "8 (Throughput)",
                "16", "16 (Max Load)"
            ])
            if new_seqs:
                current_seqs = int(new_seqs)
                # RE-CALCULATE Config based on new concurrency
                verified = get_verified_config(model_id, current_tp, current_seqs)
                current_ctx = verified["ctx"]
                current_util = verified["util"]

        elif choice == "3":
            # Configured Length Override
            new_ctx = run_dialog([
                "--title", "Context Length",
                "--inputbox", f"Override verified limit ({current_ctx}):", "10", "40", str(current_ctx)
            ])
            if new_ctx: current_ctx = int(new_ctx)

        elif choice == "4":
             # Util Override
             pass 

        elif choice == "5":
            # Toggle Attention Backend
            use_rocm_attn = not use_rocm_attn

        elif choice == "6":
            # Toggle Cache
            if not clear_cache:
                # Enabling it -> Show Warning
                warn_msg = (
                    "WARNING: Erasing the vLLM cache will remove the compiled compute graphs.\n\n"
                    "This is useful if you are experiencing crashes, 'invalid graph' errors,\n"
                    "or have switched vLLM versions recently.\n\n"
                    "However, the next startup will take longer as graphs are re-compiled.\n\n"
                    "Are you sure you want to enable this?"
                )
                confirm = run_dialog([
                    "--title", "Erase Cache Warning", 
                    "--yesno", warn_msg, "12", "60"
                ])
                
                # If confirm is not None (exit 0), it is YES.
                if confirm is not None:
                     clear_cache = True
            else:
                # Disabling it -> No warning needed
                clear_cache = False
             
        elif choice == "7":
            # Toggle Eager Mode
            use_eager = not use_eager
             
        elif choice == "8":
            # Launch
            break
            
    # Build Command
    subprocess.run(["clear"])
    
    if clear_cache:
        nuke_vllm_cache()
    
    cmd = [
        "vllm", "serve", model_path,
        "--host", HOST,
        "--port", PORT,
        "--tensor-parallel-size", str(current_tp),
        "--max-num-seqs", str(current_seqs),
        "--max-model-len", str(current_ctx),
        "--gpu-memory-utilization", str(current_util),
        "--dtype", "auto"
    ]
    
    if config.get("trust_remote"): cmd.append("--trust-remote-code")
    if use_eager: cmd.append("--enforce-eager")
    
    # Env Vars
    env = os.environ.copy()
    env.update(config.get("env", {}))
    
    if use_rocm_attn:
        env["VLLM_V1_USE_PREFILL_DECODE_ATTENTION"] = "1"
        env["VLLM_USE_TRITON_FLASH_ATTN"] = "0"
        # Optional: Explicitly mention these in print
        
    
    print("\n" + "="*60)
    print(f" Launching: {name}")
    if model_path != model_id:
        print(f" (using local model at {model_path})")
    print(f" Config:    TP={current_tp} | Seqs={current_seqs} | Ctx={current_ctx} | Util={current_util}")
    print(f" Backend:   {'ROCm' if use_rocm_attn else 'Triton'}")
    if clear_cache:
        print(f" Action:    Clearing vLLM Cache (~/.cache/vllm)")
    print(f" Command:   {' '.join(cmd)}")
    print("="*60 + "\n")
    
    # Check if model path exists
    if not os.path.exists(model_path):
        print(f"ERROR: Model path does not exist: {model_path}")
        print(f"Please ensure the model is mounted at {model_path}")
        sys.exit(1)
    
    # Run vllm serve
    try:
        result = subprocess.run(cmd, env=env, check=False)
        if result.returncode != 0:
            print(f"\nERROR: vllm serve exited with code {result.returncode}")
            sys.exit(result.returncode)
    except Exception as e:
        print(f"\nERROR: Failed to start vllm serve: {e}")
        sys.exit(1)

def main():
    check_dependencies()
    gpu_count = detect_gpus()
    
    # Check if we should use default model (for docker startup)
    use_default = os.getenv("USE_DEFAULT_MODEL", "false").lower() == "true"
    
    if use_default:
        # Find the index of default model
        try:
            default_idx = MODELS_TO_RUN.index(DEFAULT_MODEL)
            print(f"Using default model: {DEFAULT_MODEL}")
            configure_and_launch(default_idx, gpu_count)
        except ValueError:
            print(f"Error: Default model {DEFAULT_MODEL} not found in configuration")
            sys.exit(1)
    else:
        while True:
            # Build Model Menu
            menu_items = []
            for i, m_id in enumerate(MODELS_TO_RUN):
                name = m_id.split("/")[-1]
                # Mark default model
                if m_id == DEFAULT_MODEL:
                    name += " (Default)"
                menu_items.extend([str(i), name])
                
            choice = run_dialog([
                "--clear", "--backtitle", f"AMD R9700 vLLM Launcher (GPUs: {gpu_count})",
                "--title", "Select Model",
                "--menu", "Choose a model to serve:", "20", "60", "10"
            ] + menu_items)
            
            if not choice:
                subprocess.run(["clear"])
                print("Selection cancelled.")
                sys.exit(0)
                
            configure_and_launch(int(choice), gpu_count)

if __name__ == "__main__":
    main()