#!/usr/bin/env python3
"""
verify-switches.py — Automated local model switch verification test.

Tests a sequence of local↔local and local↔cloud model switches, verifying
each step: path watcher, systemd service, health endpoint, config persistence,
and VRAM sanity.

Usage:
    python3 ~/.hermes/skills/software-development/local-model-switch-optimization/scripts/verify-switches.py

The script ends with one warm local model (default: the last target).
It does NOT change the Hermes model back to what it was before — callers
should restore afterward if needed.

Exit code: 0 if all switches pass, 1 if any failed.
"""

import json, subprocess, sys, time, urllib.request, yaml
from pathlib import Path

# ── Config ──────────────────────────────────────────────────────────────
CONFIG_PATH = Path.home() / ".hermes" / "config.yaml"
MODELS = [
    # (target_name, provider, base_url, port, service_name, load_timeout)
    ("Qwen3.6-35B-A3B-UD",  "custom:local", "http://127.0.0.1:8081/v1", 8081, "llama-server.service", 20),
    ("Qwen3.5-9B-UD",       "custom:local", "http://127.0.0.1:8082/v1", 8082, "llama-server-9b.service", 15),
    ("Gemma-4-26B",         "custom:local", "http://127.0.0.1:8090/v1", 8090, "llama-server-gemma4.service", 30),
]
CLOUD = ("deepseek-v4-flash", "deepseek", "https://api.deepseek.com/v1")

SEQUENCE = [
    ("cloud→35b",     *MODELS[0]),
    ("35b→9b",        *MODELS[1]),
    ("9b→gemma4",     *MODELS[2]),
    ("gemma4→cloud",  *CLOUD,   None, None, 2),
    ("cloud→35b",     *MODELS[0]),
]

# ── Helpers ─────────────────────────────────────────────────────────────

def run(cmd, timeout=30):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
    return r.stdout.strip(), r.returncode

def check_health(port, timeout=15):
    try:
        resp = urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=timeout)
        return resp.status == 200
    except Exception:
        return False

def get_vram():
    out, _ = run("nvidia-smi --query-gpu=memory.used,memory.total --format=csv,noheader,nounits")
    used, total = out.replace(" MiB", "").split(",")
    return int(used.strip()), int(total.strip())

def write_config(name, provider, url):
    cfg = yaml.safe_load(CONFIG_PATH.open())
    cfg["model"]["default"] = name
    cfg["model"]["provider"] = provider
    cfg["model"]["base_url"] = url
    tmp = Path("/tmp/config_verify_switch.yaml")
    with tmp.open("w") as f:
        yaml.dump(cfg, f, default_flow_style=False)
    subprocess.run(f"cp {tmp} {CONFIG_PATH}", shell=True, timeout=5)

# ── Run ─────────────────────────────────────────────────────────────────

errors = []
passed = 0

print(f"Verifying {len(SEQUENCE)} switches across {len(MODELS)} local models + cloud")
print()

for i, entry in enumerate(SEQUENCE):
    label = entry[0]
    name, provider, url = entry[1], entry[2], entry[3]
    port, service, wait = entry[4], entry[5], entry[6]

    print(f"[{i+1}/{len(SEQUENCE)}] {label} → {name}")

    # Write config atomically
    write_config(name, provider, url)
    print(f"  Config written: {name} / {provider}")

    # Wait for path watcher + service transition
    time.sleep(max(2, wait))

    # Path watcher check
    out, _ = run("systemctl is-active hermes-config.path")
    if out != "active":
        errors.append(f"{label}: path watcher {out}")
        print(f"  FAIL: path watcher {out}")
        continue
    print(f"  Path watcher: active")

    if port and service:
        # Service check
        out, _ = run(f"systemctl is-active {service}")
        if out != "active":
            errors.append(f"{label}: {service} is {out}")
            print(f"  FAIL: service {service} is {out}")
            continue
        print(f"  Service {service}: active")

        # Health check
        if not check_health(port, timeout=15):
            errors.append(f"{label}: health check timeout on port {port}")
            print(f"  FAIL: health check on port {port}")
            continue
        print(f"  Health on port {port}: ok")

    # VRAM sanity
    used, total = get_vram()
    print(f"  VRAM: {used} MiB / {total} MiB")

    # Config persistence
    cfg = yaml.safe_load(CONFIG_PATH.open())
    if cfg["model"]["default"] != name:
        errors.append(f"{label}: config default drifted to {cfg['model']['default']}")
        print(f"  FAIL: config default mismatch")
        continue
    print(f"  Config persisted: OK")
    passed += 1
    print()

print(f"{'='*60}")
print(f"Results: {passed}/{len(SEQUENCE)} passed, {len(errors)} failed")
for e in errors:
    print(f"  FAIL: {e}")

sys.exit(1 if errors else 0)