#!/usr/bin/env python3
"""
Remote llama-server Switcher (Legacy)
======================================
SSH into a server and switch between llama-server systemd services.
Generic template — customize SERVER, MODELS dict, and SSH_USER.

⚠️  Consider llama-swap instead (references/llama-swap.md):
    A single binary that auto-stops/starts models on demand,
    no systemd services, no swap scripts, no manual intervention needed.

Usage:
  ./remote-switch.py status          # Show running model and reachability
  ./remote-switch.py daily           # Switch to the "daily" model
  ./remote-switch.py coder30         # Switch to "coder30"
"""

import argparse
import subprocess
import sys
import time

# ─── CUSTOMIZE THESE ─────────────────────────────────────────────
SERVER = "192.168.1.50"       # Target machine IP or hostname
SSH_USER = "rurouni"          # SSH user on the target machine

MODELS = {
    # key:     {"unit": "<systemd-unit>",    "port": PORT, "desc": "..."}
    "daily":    {"unit": "llama-server",        "port": 8081, "desc": "Daily driver"},
    "model2":   {"unit": "llama-server-9b",     "port": 8082, "desc": "Second model"},
    "model3":   {"unit": "llama-server-coder",  "port": 9021, "desc": "Coder model"},
}
# ─────────────────────────────────────────────────────────────────


def ssh(cmd: str) -> subprocess.CompletedProcess:
    return subprocess.run(
        ["ssh", f"{SSH_USER}@{SERVER}", cmd],
        capture_output=True, text=True, timeout=120,
    )


def list_running():
    r = ssh("systemctl list-units --type=service --state=running | grep llama-server")
    return [line.strip() for line in r.stdout.splitlines() if line.strip()]


def find_running_model(running_units):
    for line in running_units:
        name = line.split()[0] if line.split() else line
        if name.endswith('.service'):
            name = name[:-8]
        for key, info in MODELS.items():
            if info["unit"] == name:
                return key
    return None


def stop_current(running_units):
    for line in running_units:
        name = line.split()[0] if line.split() else line
        svc = name if name.endswith('.service') else f"{name}.service"
        print(f"  Stopping {svc}...")
        r = ssh(f"sudo systemctl stop {svc}")
        if r.returncode != 0:
            print(f"  ⚠ Warning: {r.stderr.strip()}")
    time.sleep(2)


def wait_for_port(port: int, timeout=90):
    start = time.time()
    while time.time() - start < timeout:
        r = subprocess.run(
            ["curl", "-s", "--connect-timeout", "3",
             f"http://{SERVER}:{port}/health"],
            capture_output=True, text=True, timeout=10,
        )
        if r.returncode == 0 and '"status":"ok"' in r.stdout:
            return True
        time.sleep(3)
    return False


def switch(target: str):
    if target not in MODELS:
        print(f"Unknown model '{target}'. Options: {', '.join(sorted(MODELS.keys()))}")
        sys.exit(1)

    info = MODELS[target]
    print(f"🔄 Switching to: {info['desc']} (port {info['port']})")

    running = list_running()
    if running:
        current = find_running_model(running)
        if current == target:
            print(f"  ✓ {info['desc']} is already running")
            return
        stop_current(running)

    print(f"  Starting {info['unit']}...")
    r = ssh(f"sudo systemctl start {info['unit']}")
    if r.returncode != 0:
        print(f"  ✗ Failed: {r.stderr.strip()}")
        sys.exit(1)

    print(f"  Waiting for port {info['port']}...")
    if wait_for_port(info['port']):
        print(f"  ✓ {info['desc']} ready on port {info['port']}")
    else:
        print(f"  ✗ Timed out waiting for port {info['port']}")


def status():
    running = list_running()
    if not running:
        print("No llama-server services running.")
    else:
        print("Running:")
        for line in running:
            name = line.split()[0] if line.split() else line
            key = find_running_model([line])
            if key:
                info = MODELS[key]
                print(f"  ✓ {info['desc']}  →  port {info['port']}")
            else:
                print(f"  {line}")

    print("\nReachability:")
    for key, info in sorted(MODELS.items()):
        r = subprocess.run(
            ["curl", "-s", "--connect-timeout", "2",
             f"http://{SERVER}:{info['port']}/health"],
            capture_output=True, text=True, timeout=5,
        )
        ok = r.returncode == 0 and '"status":"ok"' in r.stdout
        print(f"  {'✓' if ok else '✗'}  port {info['port']:>5}  {info['desc']}")


def main():
    parser = argparse.ArgumentParser(description=f"Remote llama-server switcher ({SERVER})")
    parser.add_argument("action", nargs="?", default="status",
                        choices=list(MODELS.keys()) + ["status"])
    args = parser.parse_args()
    if args.action == "status":
        status()
    else:
        switch(args.action)


if __name__ == "__main__":
    main()