#!/usr/bin/env python3
"""
Hermes Local Model Switcher — Template
---------------------------------------
Usage: python3 ~/.hermes/scripts/model-switcher      (interactive menu)
       python3 ~/.hermes/scripts/model-switcher 9b    (direct switch)

Interactive TUI menu for hot-swapping between local llama.cpp models.
Stops the current model, starts the selected one, updates Hermes config,
waits for the model to respond. Edit the MODELS list below for your setup.

Each model entry can include:
  - name, desc, detail   (display string + tooltip)
  - service              (systemd unit name)
  - port                 (llama-server port)
  - provider / model_name (Hermes custom_provider config keys)
  - context              (context length, informational)
  - short                (quick-switch alias)
"""

import subprocess
import sys
import os
import time
import urllib.request

# ── EDIT THIS LIST for your models ────────────────────────────────────
MODELS = [
    {
        "name": "Qwen3.6-35B",
        "desc": "BEST All-Purpose: biggest, smartest, handles everything",
        "detail": "MoE 35B (3B active) - daily driver for complex reasoning, coding, creative tasks",
        "service": "llama-server.service",
        "port": 8081,
        "provider": "custom:local",
        "model_name": "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
        "context": 128000,
        "short": "35b",
    },
    {
        "name": "Qwen3.5-9B",
        "desc": "FAST All-Purpose: lighter daily driver, good reasoning + speed",
        "detail": "9B dense - balanced pick for quick tasks, general chat, simpler queries",
        "service": "llama-server-9b.service",
        "port": 8082,
        "provider": "custom:local",
        "model_name": "Qwen3.5-9B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf",
        "context": 64000,
        "short": "9b",
    },
    {
        "name": "Llama-3.2-3B",
        "desc": "TINY FAST: ultra-light Llama for simple tasks, maximum speed",
        "detail": "3.2B Llama Q5_K_M - super fast response, perfect for routing/classification/simple QA",
        "service": "llama-server-llama3.service",
        "port": 8086,
        "provider": "custom:local",
        "model_name": "Llama-3.2-3B-Instruct-Q5_K_M.gguf",
        "context": 32768,
        "short": "llama3",
    },
    {
        "name": "SmolLM3-3B",
        "desc": "SMALLEST FASTEST: lightweight champ, instant responses",
        "detail": "3B SmolLM Q4_K_M - the speed king for trivial tasks, memory/agent routing",
        "service": "llama-server-smol.service",
        "port": 8087,
        "provider": "custom:local",
        "model_name": "HuggingFaceTB_SmolLM3-3B-Q4_K_M.gguf",
        "context": 16384,
        "short": "smol",
    },
    {
        "name": "Phi-4-14B",
        "desc": "SOLID ALL-ROUNDER: Microsoft dense 14B, fast on GPU",
        "detail": "14B Phi-4 Q4_K - solid reasoning, fits GPU, great middle ground",
        "service": "llama-server-phi4.service",
        "port": 8088,
        "provider": "custom:local",
        "model_name": "phi-4-Q4_K.gguf",
        "context": 65536,
        "short": "phi4",
    },
    {
        "name": "Gemma-4-26B-A4B",
        "desc": "GOOGLE MOE: latest MoE w/ vision, 256K ctx, big brain",
        "detail": "MoE 26B (8 active) Google Gemma 4 - native vision, 256K context",
        "service": "llama-server-gemma4.service",
        "port": 8090,
        "provider": "custom:local",
        "model_name": "gemma-4-26B-A4B-it-Q4_K_M.gguf",
        "context": 131072,
        "short": "gemma4",
    },
    {
        "name": "Gemma-3-12B",
        "desc": "GOOGLE DENSE: community #1 pick for 11GB VRAM",
        "detail": "12B Gemma 3 Q4_K_M - native 128K ctx, fits GPU, top-rated",
        "service": "llama-server-gemma3.service",
        "port": 8089,
        "provider": "custom:local",
        "model_name": "gemma-3-12b-it-Q4_K_M.gguf",
        "context": 131072,
        "short": "gemma3",
    },
]
# ──────────────────────────────────────────────────────────────────────

HERMES_CLI = os.path.expanduser("~/.local/bin/hermes")
HERMES_CONFIG = os.path.expanduser("~/.hermes/config.yaml")


def get_active_service():
    for model in MODELS:
        r = subprocess.run(["systemctl", "is-active", model["service"]],
                           capture_output=True, text=True)
        if r.stdout.strip() == "active":
            return model
    return None


def get_current_config():
    """Find which model Hermes config currently points to by matching model.default."""
    try:
        with open(HERMES_CONFIG) as f:
            cfg_name = None
            for line in f:
                line = line.strip()
                if line.startswith("default:") and ".gguf" in line:
                    cfg_name = line.split("default:")[-1].strip().strip("'\"")
                    break
            if cfg_name:
                for m in MODELS:
                    if m["model_name"] == cfg_name:
                        return m
    except:
        pass
    return None


def swap(target):
    active = get_active_service()
    current_cfg = get_current_config()

    if active and active["name"] == target["name"] \
       and current_cfg and current_cfg["name"] == target["name"]:
        print(f"\n  ✅ {target['name']} is already running and configured")
        return True

    print(f"\n  🔄 Switching to {target['name']}...", flush=True)

    # Stop current (or all, to be safe)
    if active:
        subprocess.run(["sudo", "systemctl", "stop", active["service"]],
                       capture_output=True, text=True)
        print(f"  ⏹  Stopped {active['name']:20s} ✅", flush=True)
    else:
        for m in MODELS:
            subprocess.run(["sudo", "systemctl", "stop", m["service"]],
                           capture_output=True, text=True)

    r = subprocess.run(["sudo", "systemctl", "start", target["service"]],
                       capture_output=True, text=True)
    if r.returncode != 0:
        print(f"  ❌ Failed to start {target['name']}: {r.stderr.strip()}", flush=True)
        return False
    print(f"  ▶️  Starting {target['name']:20s} ✅", flush=True)

    for key, value in [
        ("model.provider", target["provider"]),
        ("model.default", target["model_name"]),
        ("model.base_url", f"http://127.0.0.1:{target['port']}/v1"),
    ]:
        subprocess.run([HERMES_CLI, "config", "set", key, value],
                       capture_output=True, text=True)

    print(f"  ⏳ Waiting for {target['name']} to load...", end="", flush=True)
    for _ in range(30):
        try:
            req = urllib.request.Request(f"http://127.0.0.1:{target['port']}/v1/models")
            resp = urllib.request.urlopen(req, timeout=2)
            if resp.status == 200:
                print(" ✅", flush=True)
                break
        except:
            pass
        print(".", end="", flush=True)
        time.sleep(2)
    else:
        print(" ❌ (timeout)", flush=True)
        return False

    print(f"\n  ✅ Now using {target['name']} on port {target['port']}", flush=True)
    print(f"     Context: {target['context']:,} tokens | {target['desc']}", flush=True)
    print(f"     Run /reset in Hermes to start a fresh session with this model.", flush=True)
    return True


def show_all_models():
    print()
    print(f"  {'#':<3} {'Model':22s} {'Status':12s} {'Port':6s} {'Context':>8s}  Use Case")
    print(f"  {'─'*3} {'─'*22} {'─'*12} {'─'*6} {'─'*8}  {'─'*50}")
    for i, model in enumerate(MODELS):
        active = subprocess.run(
            ["systemctl", "is-active", model["service"]],
            capture_output=True, text=True
        ).stdout.strip()
        status = "🟢 ACTIVE" if active == "active" else "⚪"
        ctx = f"{model['context']:,}"
        print(f"  [{i+1}] {model['name']:22s} {status:12s} {model['port']:6s} {ctx:>8s}  {model['desc']}")
    print()


def main():
    if len(sys.argv) > 1:
        arg = sys.argv[1].lower().replace("-", "").replace(" ", "")
        for model in MODELS:
            if arg == model["short"].lower() or arg in model["name"].lower().replace(".", "").replace("-", ""):
                return 0 if swap(model) else 1
        print(f"Unknown model. Try: {', '.join(m['short'] for m in MODELS)}")
        return 1

    while True:
        os.system('clear' if os.name == 'posix' else 'cls')

        active = get_active_service()
        print("╔══════════════════════════════════════════════════════╗")
        print("║      🤖  Hermes Local Model Switcher               ║")
        print("╚══════════════════════════════════════════════════════╝")

        if active:
            print(f"\n  🟢 Currently running: {active['name']}")
            print(f"     {active['detail']}")
        else:
            print(f"\n  ⚪ No local model currently running")

        show_all_models()

        print("  Quick shortcuts:", ", ".join(m['short'] for m in MODELS))
        print("  [q] Quit\n")

        choice = input("  Select model (# or name): ").strip().lower()

        if choice in ('q', 'quit', 'exit'):
            print("\n  👋 Bye!")
            break

        matched = [m for m in MODELS if choice == m["short"].lower()]
        if not matched:
            try:
                idx = int(choice) - 1
                if 0 <= idx < len(MODELS):
                    matched = [MODELS[idx]]
            except ValueError:
                for m in MODELS:
                    if choice in m["name"].lower().replace("-", "").replace(".", ""):
                        matched = [m]
                        break

        if matched:
            swap(matched[0])
            input("\n  Press Enter to continue...")
        else:
            print(f"\n  ❌ Invalid choice. Try a number (1-{len(MODELS)}) or shortcut name.")
            input("  Press Enter...")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n  👋 Bye!")
