#!/usr/bin/env python3
"""
Auto-switch llama-server service when Hermes config model changes.

Triggered by systemd path unit (hermes-config.path) on config.yaml changes.
Uses state file guard to prevent re-trigger loops.
Performs rollback if the target service fails to start.
Writes notification file for the cron watchdog to deliver.

Single source of truth for model definitions: models.py

INSTALLED PATH: ~/.hermes/scripts/auto-model-switch.py
This copy is canonical for reference.

KEY DESIGN:
  - PathModified trigger on config.yaml → instant response
  - State file (.last-model-state) prevents re-trigger loops
  - Health-check guard: even if state matches, verify service is actually running
  - Rollback: if target fails to start, restart the previous service
  - Notification file: writes .model-switch-notify for watchdog to deliver
  - OpenRouter validation: blocks models not in curated override list
"""

import subprocess
import os
import sys
import json
import logging
from datetime import datetime

# Ensure shared models.py is importable
_script_dir = os.path.dirname(os.path.abspath(__file__))
if _script_dir not in sys.path:
    sys.path.insert(0, _script_dir)
from models import MODELS

# ── Paths ──
CONFIG = os.path.expanduser("~/.hermes/config.yaml")
STATE_FILE = os.path.expanduser("~/.hermes/.last-model-state")
NOTIFY_FILE = os.path.expanduser("~/.hermes/.model-switch-notify")
LOG_DIR = os.path.expanduser("~/.hermes/logs")
LOG_FILE = os.path.join(LOG_DIR, "model-switch.log")

os.makedirs(LOG_DIR, exist_ok=True)
logging.basicConfig(
    filename=LOG_FILE,
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("model-switch")
_echo = lambda msg: print(msg, flush=True)


def parse_yaml_simple(filepath: str) -> dict:
    """Simple YAML parser for ~/.hermes/config.yaml — flat key/value return."""
    result = {}
    current_key = None
    current_indent = 0
    with open(filepath) as f:
        for line in f:
            stripped = line.rstrip()
            if not stripped or stripped.startswith("#"):
                continue
            indent = len(line) - len(line.lstrip())
            if ":" in stripped:
                key_part, _, val_part = stripped.partition(":")
                key = key_part.strip()
                val = val_part.strip().strip("'\"")
                if indent == 0 and val == "":
                    current_key = key; current_indent = 0
                elif indent > 0 and val == "":
                    current_key = key; current_indent = indent
                elif current_key and indent > current_indent:
                    result[f"{current_key}.{key}"] = val if val else ""
                elif val:
                    result[key] = val if val else ""
                else:
                    current_key = key
            else:
                result[current_key] = stripped.strip()
    return result


def write_model_yaml(filepath: str, provider: str, model: str):
    """Write model.provider and model.default into config.yaml in-place."""
    with open(filepath) as f:
        lines = f.readlines()
    new_lines = []
    wrote_provider = False
    wrote_default = False
    for line in lines:
        if line.lstrip().startswith("provider:") and not line.lstrip().startswith("providers:"):
            if not wrote_provider:
                new_lines.append(f"  provider: {provider}\n")
                wrote_provider = True
            else:
                continue
        elif line.lstrip().startswith("default:"):
            if not wrote_default:
                new_lines.append(f"  default: {model}\n")
                wrote_default = True
            else:
                continue
        else:
            new_lines.append(line)
    if not wrote_provider:
        new_lines.append(f"  provider: {provider}\n")
    if not wrote_default:
        new_lines.append(f"  default: {model}\n")
    with open(filepath, "w") as f:
        f.writelines(new_lines)


def get_config_settings() -> dict:
    """Read model.provider and model.default from config.yaml."""
    if not os.path.exists(CONFIG):
        return {"provider": None, "model": None}
    parsed = parse_yaml_simple(CONFIG)
    return {"provider": parsed.get("model.provider", ""), "model": parsed.get("model.default", "")}


def get_last_good_state() -> tuple:
    """Read last-known-GOOD state from state file. Returns (model_name, provider)."""
    if not os.path.exists(STATE_FILE):
        return (None, None)
    with open(STATE_FILE) as f:
        lines = f.read().strip().split("\n")
    return (lines[0] if lines else None, lines[1] if len(lines) > 1 else None)


def save_good_state(model_name: str, provider: str):
    """Persist the last good model+provider."""
    with open(STATE_FILE, "w") as f:
        f.write(f"{model_name}\n{provider}\n")


def is_allowed_openrouter_model(model_id: str) -> bool:
    """Check if an OpenRouter model is in the curated override list."""
    override = os.path.expanduser("~/.hermes/cache/openrouter_override.json")
    if not os.path.exists(override):
        return True  # fail open
    try:
        with open(override) as f:
            data = json.load(f)
        allowed = [m["id"] for m in data["providers"]["openrouter"]["models"]]
        return model_id in allowed
    except (json.JSONDecodeError, KeyError):
        return True


def get_active_service() -> dict | None:
    """Find which model service is currently active (systemd)."""
    for m in MODELS:
        r = subprocess.run(["systemctl", "is-active", m["service"]],
                           capture_output=True, text=True, timeout=5)
        if r.stdout.strip() == "active":
            return m
    return None


def stop_service(model: dict) -> bool:
    log.info("Stopping %s (%s)...", model["name"], model["service"])
    r = subprocess.run(["sudo", "systemctl", "stop", model["service"]],
                       capture_output=True, text=True, timeout=30)
    ok = r.returncode == 0
    log.info("  %s Stopped %s", "✅" if ok else "⚠️", model["name"])
    return ok


def start_service(model: dict) -> bool:
    log.info("Starting %s (%s)...", model["name"], model["service"])
    r = subprocess.run(["sudo", "systemctl", "start", model["service"]],
                       capture_output=True, text=True, timeout=30)
    ok = r.returncode == 0
    log.info("  %s Started %s", "✅" if ok else "❌", model["name"])
    return ok


def write_notification(target: dict, action: str = "switched"):
    """Write a one-shot notification file for the cron watchdog."""
    vram = "?"
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.used,memory.total",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5)
        if r.returncode == 0:
            parts = r.stdout.strip().split(", ")
            if len(parts) == 2:
                vram = f"{parts[0]} MB / {parts[1]} MB"
    except Exception:
        pass

    ts = datetime.now().strftime("%I:%M:%S %p")
    msg = (
        f"🔀 **Model Switch**\n"
        f"━━━━━━━━━━━━━━━\n"
        f"**Event:** {action}\n"
        f"**Model:** {target['name']}\n"
        f"**Port:** {target['port']}\n"
        f"**VRAM:** {vram}\n"
        f"**Context:** {target.get('context', '?')}K\n"
        f"**Time:** {ts} PDT"
    )
    os.makedirs(os.path.dirname(NOTIFY_FILE), exist_ok=True)
    with open(NOTIFY_FILE, "w") as f:
        f.write(msg + "\n")


def swap_to(target: dict) -> bool:
    """Stop active service, start target. Roll back on failure."""
    active = get_active_service()
    if active and active["name"] == target["name"]:
        log.info("Already on %s — nothing to swap", target["name"])
        return True

    previous = active
    if active:
        _echo(f"  ⏹  Stopping {active['name']}...")
        stop_service(active)

    _echo(f"  ▶️  Starting {target['name']}...")
    ok = start_service(target)

    if ok:
        write_notification(target, "switched")
        return True

    # Rollback
    _echo(f"  ❌ Failed to start {target['name']}")
    if previous:
        _echo(f"  ↩️  Rolling back to {previous['name']}...")
        if start_service(previous):
            write_notification(previous, "rollback")
            _echo(f"  ✅ Rolled back to {previous['name']}")
        else:
            _echo(f"  ❌ Rollback failed!")
    return False


def main():
    settings = get_config_settings()
    config_model = settings["model"]
    config_provider = settings["provider"]

    if not config_model:
        sys.exit(0)

    # ── OpenRouter validation ──
    if config_provider == "openrouter":
        last_good_model, last_good_provider = get_last_good_state()
        if not is_allowed_openrouter_model(config_model):
            _echo(f"  ❌ '{config_model}' NOT in curated OpenRouter list!")
            if last_good_model and last_good_model != config_model:
                write_model_yaml(CONFIG, last_good_provider or "openrouter", last_good_model)
            else:
                write_model_yaml(CONFIG, "custom:local", "qwen-35b-fallback.gguf")
            settings = get_config_settings()
            save_good_state(settings["model"], settings["provider"])
            sys.exit(1)
        save_good_state(config_model, config_provider)
        sys.exit(0)

    # ── Local model swap ──
    target = next((m for m in MODELS if m["model_name"] == config_model), None)
    if not target:
        # Cloud model — no service to swap
        save_good_state(config_model, config_provider)
        sys.exit(0)

    # State guard + health check
    last_good_model, _ = get_last_good_state()
    if last_good_model == config_model:
        active = get_active_service()
        if active and active["name"] == target["name"]:
            sys.exit(0)  # healthy, skip
        # Dead or mismatched — fall through to restart
        if active:
            log.warning("State says %s but %s is active — fixing", target["name"], active["name"])
        else:
            _echo(f"  🔄 {target['name']} is dead — restarting...")
            if start_service(target):
                save_good_state(config_model, config_provider)
                write_notification(target, "restarted")
                _echo(f"  ✅ Restarted {target['name']}")
            sys.exit(0 if ok else 1)

    # Perform the swap
    _echo(f"  🔄 Auto-switching to {target['name']}...")
    result = swap_to(target)
    if result:
        save_good_state(config_model, config_provider)
        _echo(f"  ✅ Now using {target['name']}")
    else:
        log.warning("Swap failed — state NOT saved")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        log.exception("Unhandled error")
        _echo(f"  ❌ Error: {e}")
        sys.exit(1)