#!/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 flock-based lock file to prevent concurrent invocations from racing.
Uses state file guard to prevent re-trigger loops.
Performs rollback if the target service fails to start.
Logs to ~/.hermes/logs/model-switch.log for troubleshooting.

Single source of truth for model definitions: models.py
"""
import os
import sys
import time
import fcntl
import logging
import subprocess
import tempfile
import urllib.request
import json
import yaml
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")
LOCK_FILE = os.path.expanduser("~/.hermes/.model-switch.lock")
LOG_DIR = os.path.expanduser("~/.hermes/logs")
LOG_FILE = os.path.join(LOG_DIR, "model-switch.log")

# ── Logging ────────────────────────────────────────────────────────
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")

# Also echo to journal for systemd visibility
_echo = lambda msg: print(msg, flush=True)


# ── Lock file ──────────────────────────────────────────────────────
_LOCK_FD: int | None = None


def acquire_lock() -> bool:
    """Acquire an exclusive lock on LOCK_FILE. Returns True if acquired."""
    global _LOCK_FD
    try:
        fd = os.open(LOCK_FILE, os.O_CREAT | os.O_RDWR, 0o644)
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        _LOCK_FD = fd
        # Write our PID so we can debug stale locks
        os.ftruncate(fd, 0)
        os.write(fd, f"{os.getpid()}\n".encode())
        return True
    except (IOError, OSError):
        if _LOCK_FD is None and 'fd' in locals():
            os.close(fd)
        return False


def release_lock():
    """Release the lock file."""
    global _LOCK_FD
    if _LOCK_FD is not None:
        try:
            fcntl.flock(_LOCK_FD, fcntl.LOCK_UN)
            os.close(_LOCK_FD)
        except (IOError, OSError):
            pass
        _LOCK_FD = None
        try:
            os.unlink(LOCK_FILE)
        except OSError:
            pass


# ── Helpers ────────────────────────────────────────────────────────
OPENROUTER_OVERRIDE = os.path.expanduser("~/.hermes/cache/openrouter_override.json")


def parse_yaml_simple(filepath: str) -> dict:
    """Ultra-simple YAML parser for ~/.hermes/config.yaml — returns flat-ish dict.
    Only handles basic indented key-value pairs and leaf values."""
    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, base_url: str = ""):
    """Atomically update only the top-level model section in config.yaml."""
    provider_norm = (provider or "").strip()
    model_norm = (model or "").strip()
    matched_local = next(
        (
            m for m in MODELS
            if m.get("model_name") == model_norm or m.get("name") == model_norm
        ),
        None,
    )
    if provider_norm == "local" or matched_local:
        provider_norm = "custom:local"
        if matched_local:
            model_norm = matched_local.get("name") or model_norm

    with open(filepath, encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}
    if not isinstance(cfg, dict):
        cfg = {}
    model_cfg = cfg.setdefault("model", {})
    if not isinstance(model_cfg, dict):
        model_cfg = {}
        cfg["model"] = model_cfg
    model_cfg["provider"] = provider_norm
    model_cfg["default"] = model_norm
    if base_url:
        model_cfg["base_url"] = base_url
    directory = os.path.dirname(filepath) or "."
    fd, tmp_path = tempfile.mkstemp(prefix=".config.yaml.", suffix=".tmp", dir=directory, text=True)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            yaml.safe_dump(cfg, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp_path, filepath)
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


def get_config_settings() -> dict:
    """Read model.provider and model.default from config.yaml."""
    if not os.path.exists(CONFIG):
        log.warning("Config file not found: %s", 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."""
    if not os.path.exists(STATE_FILE):
        return (None, None)
    with open(STATE_FILE) as f:
        lines = f.read().strip().split("\n")
    model = lines[0] if lines else None
    provider = lines[1] if len(lines) > 1 else None
    return (model, provider)


def save_good_state(model_name: str, provider: str):
    """Persist the last good model+provider."""
    content = f"{model_name}\n{provider}\n"
    with open(STATE_FILE, "w") as f:
        f.write(content)
    log.debug("Good state saved: provider=%s model=%s", provider, model_name)


def is_allowed_openrouter_model(model_id: str) -> bool:
    """Check if an OpenRouter model is in the curated override list."""
    if not os.path.exists(OPENROUTER_OVERRIDE):
        log.warning("OpenRouter override file not found — cannot validate")
        return True
    try:
        with open(OPENROUTER_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) as e:
        log.error("Failed to parse OpenRouter override: %s", e)
        return True


def get_active_service() -> dict | None:
    """Find which model service is currently active (systemd), skipping preserved ones."""
    for m in MODELS:
        if m.get("preserve"):
            continue
        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 is_cloud_model(model_name: str) -> bool:
    """Check if a model_name refers to a cloud model (not in MODELS list).
    Matches on both model_name (GGUF filename) and name (friendly name)."""
    return not any(
        m["model_name"] == model_name or m["name"] == model_name for m in MODELS
    )


def stop_service(model: dict) -> bool:
    """Stop a single model's systemd service. Returns True on success."""
    log.info("Stopping %s (%s)...", model["name"], model["service"])
    try:
        r = subprocess.run(
            ["sudo", "systemctl", "stop", model["service"]],
            capture_output=True, text=True, timeout=45,
        )
        if r.returncode == 0:
            log.info("  ✅ Stopped %s", model["name"])
            return True
        log.warning("  ⚠️  Stop %s returned %d: %s",
                    model["name"], r.returncode, r.stderr.strip()[:200])
    except subprocess.TimeoutExpired:
        log.warning("  ⏳ Stop %s still draining after 45s — waiting for unit to become inactive", model["name"])

    deadline = time.time() + 90
    while time.time() < deadline:
        try:
            probe = subprocess.run(
                ["systemctl", "is-active", model["service"]],
                capture_output=True, text=True, timeout=5,
            )
            state = probe.stdout.strip()
            if state in {"inactive", "failed", "deactivating"}:
                if state != "deactivating":
                    log.info("  ✅ Stopped %s (final state: %s)", model["name"], state)
                    return True
        except Exception:
            pass
        time.sleep(2)

    try:
        probe = subprocess.run(
            ["systemctl", "is-active", model["service"]],
            capture_output=True, text=True, timeout=5,
        )
        final_state = probe.stdout.strip() or "unknown"
    except Exception:
        final_state = "unknown"
    log.warning("  ⚠️  Stop %s did not settle inactive in time (state=%s)", model["name"], final_state)
    return final_state in {"inactive", "failed"}


def start_service(model: dict) -> bool:
    """Start a model's systemd service. Returns True on success."""
    log.info("Starting %s (%s)...", model["name"], model["service"])
    r = subprocess.run(
        ["sudo", "systemctl", "start", model["service"]],
        capture_output=True, text=True, timeout=300,
    )
    if r.returncode == 0:
        log.info("  ✅ Started %s", model["name"])
    else:
        log.error("  ❌ Start %s failed (%d): %s",
                   model["name"], r.returncode, r.stderr.strip()[:300])
    return r.returncode == 0


def write_notification(target: dict, action: str = "switched"):
    """Write a one-shot notification file for the watchdog to deliver."""
    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 wait_for_model(port: int, timeout_sec: int = 60) -> tuple[bool, float]:
    """Wait for the model API to respond. Returns (success, seconds_waited)."""
    url = f"http://127.0.0.1:{port}/v1/models"
    start = time.time()
    for _ in range(timeout_sec):
        try:
            resp = urllib.request.urlopen(urllib.request.Request(url), timeout=2)
            if resp.status == 200:
                return True, time.time() - start
        except Exception:
            pass
        time.sleep(1)
    return False, time.time() - start


# ── Core logic ────────────────────────────────────────────────────
def swap_to(target: dict) -> bool:
    """Stop the currently active service, start target.
    Does NOT write config.yaml (caller already did). Only rollback writes config.
    """
    active = get_active_service()

    # Already the right one — nothing to do
    if active and active["name"] == target["name"]:
        log.info("Already on %s — nothing to swap", target["name"])
        return True

    previous = active

    # 1. Stop active model if needed (skip preserved/background services)
    if active and active["provider"] == "local" and target["provider"] == "local":
        if active.get("preserve"):
            _echo(f"  🛡  Preserving {active['name']} (background service)")
        else:
            _echo(f"  ⏹  Stopping {active['name']}...")
            stop_service(active)
            _echo("  💤 Waiting 3s for VRAM drain...")
            time.sleep(3)
    elif not active:
        log.info("No active service to stop")

    # 2. Start the target
    if target["provider"] == "local":
        _echo(f"  ▶️  Starting {target['name']}...")
        ok = start_service(target)
        if ok:
            _echo(f"  ⏳ Waiting for {target['name']} to respond...")
            loaded, elapsed = wait_for_model(target['port'], timeout_sec=60)
            if loaded:
                _echo(f"  ✅ {target['name']} responded in {elapsed:.1f}s")
            else:
                _echo(f"  ⚠️  {target['name']} API not responding after {elapsed:.0f}s — continuing")
            log.info("Switch complete: %s → %s (API ready: %s in %.1fs)",
                     previous["name"] if previous else "(none)", target["name"],
                     "yes" if loaded else "no", elapsed)
            correct_base_url = f"http://127.0.0.1:{target['port']}/v1"
            write_model_yaml(CONFIG, target.get("provider", "custom:local"),
                             target["model_name"], correct_base_url)
            write_notification(target, "switched")
            return True

        # 3. Rollback on failure
        _echo(f"  ❌ Failed to start {target['name']}")
        log.error("Switch FAILED — attempting rollback")
        if previous:
            _echo(f"  ↩️  Rolling back to {previous['name']}...")
            rollback_ok = start_service(previous)
            if rollback_ok:
                log.info("Rollback successful: restarted %s", previous["name"])
                write_model_yaml(CONFIG, previous.get("provider", "custom:local"),
                                 previous["model_name"],
                                 f"http://127.0.0.1:{previous['port']}/v1")
                write_notification(previous, "rollback")
                _echo(f"  ✅ Rolled back to {previous['name']}")
            else:
                log.critical("Rollback FAILED — %s could not be restarted!", previous["name"])
                _echo(f"  ❌ Rollback failed! {previous['name']} not running!")
            return False
        else:
            log.error("No previous service to roll back to")
            _echo("  ⚠️  No previous service to roll back to")
            return False

    # Cloud model — no service to swap
    _echo(f"  ☁️  Cloud model {target['model_name']} — no service swap needed")
    return True


# ── Entry point ────────────────────────────────────────────────────
def main():
    if not acquire_lock():
        log.debug("Lock held by another process — exiting")
        sys.exit(0)
    try:
        _main_impl()
    finally:
        release_lock()


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

    if not config_model:
        log.info("No model in config — exiting")
        return

    # ── OpenRouter model 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}' is NOT in your curated OpenRouter model list!")
            log.warning("BLOCKED: OpenRouter model '%s' not in allowed list", config_model)
            if last_good_model and last_good_model != config_model:
                _echo(f"  ↩️  Reverting to {last_good_model}")
                log.info("Reverting config to last good model: %s", last_good_model)
                write_model_yaml(CONFIG, last_good_provider or "openrouter", last_good_model)
            else:
                fallback = "Qwen3.6-35B-A3B-UD-Q4_K_S.gguf"
                _echo(f"  ↩️  No last good state — reverting to local: {fallback}")
                log.info("Reverting config to local fallback: %s", fallback)
                write_model_yaml(CONFIG, "custom:local", fallback)
            settings = get_config_settings()
            save_good_state(settings["model"], settings["provider"])
            return
        save_good_state(config_model, config_provider)
        log.info("OpenRouter model validated: %s", config_model)
        return

    # ── Cloud model (no local service to manage) ──
    if is_cloud_model(config_model):
        log.info("Config set to cloud model: %s — no service to swap", config_model)
        save_good_state(config_model, config_provider)
        return

    # ── Local model swap ──
    target = next(
        (m for m in MODELS if m["model_name"] == config_model or m["name"] == config_model),
        None,
    )
    if not target:
        log.warning("Unknown model in config: %s — ignoring", config_model)
        return

    # Guard: skip if we already handled this model
    last_good_model, _ = get_last_good_state()
    if last_good_model == config_model:
        target = next((m for m in MODELS if m["model_name"] == config_model or m["name"] == config_model), None)
        if target:
            active = get_active_service()
            if active and active["name"] == target["name"]:
                log.debug("State unchanged (%s) — service healthy, skipping", config_model)
                return
            elif active:
                log.warning("State says %s but %s is active — fixing mismatch",
                            target["name"], active["name"])
            else:
                log.warning("State says %s but service is dead — restarting", target["name"])
                _echo(f" 🔄 {target['name']} is dead — restarting...")
                ok = start_service(target)
                if ok:
                    correct_url = f"http://127.0.0.1:{target['port']}/v1"
                    write_model_yaml(CONFIG, target.get("provider", "custom:local"),
                                     target["model_name"], correct_url)
                    write_notification(target, "restarted")
                    save_good_state(config_model, config_provider)
                    _echo(f" ✅ Restarted {target['name']}")
                else:
                    _echo(f"  ❌ Failed to restart {target['name']}")
                return
        else:
            log.debug("State unchanged (%s) — skipping", config_model)
            return

    # Perform the swap
    log.info("Config changed to local model: %s", target["name"])
    _echo(f"  🔄 Auto-switching to {target['name']}...")
    result = swap_to(target)

    if result:
        save_good_state(config_model, config_provider)
        _echo(f"  ✅ Switched to {target['name']}")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        log.error("Unhandled error in auto-model-switch", exc_info=True)
        _echo(f"  ❌ Auto-switch error: {e}")
        sys.exit(1)