#!/usr/bin/env python3
"""
Quick-start a local model service by gguf filename.

Called from gateway after /model switch to a local model.
Stops any other running local service first, then starts the target.
Verifies the model actually loads, then sends a Telegram notification.

Usage: quick-start-model.py <gguf-filename>
"""

import subprocess
import sys
import os
import time
import urllib.request
import urllib.error
from datetime import datetime
from pathlib import Path

_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

CONFIG_PATH = os.path.expanduser("~/.hermes/config.yaml")
NOTIFY_FILE = os.path.expanduser("~/.hermes/.model-switch-notify")
HERMES_HOME = os.path.expanduser("~/.hermes")


def log(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_dir = os.path.join(HERMES_HOME, "logs")
    os.makedirs(log_dir, exist_ok=True)
    with open(os.path.join(log_dir, "model-switch.log"), "a") as f:
        f.write(f"{ts} [quick-start] {msg}\n")


def send_notification(message: str):
    """Deliver a notification via Hermes cron watchdog.
    Works on any platform (Telegram, Discord, etc.) because Hermes cron
    auto-resolves the delivery target to the current origin platform."""
    import subprocess as _sp
    # Write the notification file for the watchdog to pick up
    try:
        with open(NOTIFY_FILE, "w") as f:
            f.write(message + "\n")
    except Exception as e:
        log(f"Failed to write notify file: {e}")
        return

    # Trigger the watchdog immediately via cron run
    watchdog_id = "a35af7b299a1"
    try:
        hermes_bin = os.path.expanduser("~/.local/bin/hermes")
        _sp.run(
            [hermes_bin, "cron", "run", watchdog_id],
            capture_output=True, timeout=120,
        )
        log("Watchdog triggered for delivery")
    except Exception as e:
        log(f"Failed to trigger watchdog: {e}")


def get_vram() -> str:
    """Get current VRAM usage."""
    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:
                used_gb = int(parts[0]) / 1024
                total_gb = int(parts[1]) / 1024
                return f"{used_gb:.1f}G / {total_gb:.1f}G"
    except Exception:
        pass
    return "?"


def wait_for_model(port: int, timeout_sec: int = 120) -> 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 i in range(timeout_sec):
        try:
            req = urllib.request.Request(url)
            resp = urllib.request.urlopen(req, timeout=2)
            if resp.status == 200:
                elapsed = time.time() - start
                return True, elapsed
        except (urllib.error.URLError, ConnectionError, OSError):
            pass
        time.sleep(1)
    return False, time.time() - start


def update_config(port: int, model_name: str, provider: str):
    """Update model.default, model.provider, and model.base_url in config.yaml.
    Written BEFORE the path watcher fires so auto-model-switch.py sees the
    correct target model and doesn't revert the switch."""
    base_url = f"http://127.0.0.1:{port}/v1"
    with open(CONFIG_PATH) as f:
        lines = f.readlines()
    new_lines = []
    wrote_default = False
    wrote_provider = False
    wrote_base_url = False
    for line in lines:
        stripped = line.lstrip()
        # model.default
        if stripped.startswith("default:") and not wrote_default:
            new_lines.append(f"  default: {model_name}\n")
            wrote_default = True
        # model.provider
        elif stripped.startswith("provider:") and not stripped.startswith("providers:") and not wrote_provider:
            new_lines.append(f"  provider: {provider}\n")
            wrote_provider = True
        # model.base_url
        elif stripped.startswith("base_url:") and not stripped.startswith("base_urls:") and not wrote_base_url:
            new_lines.append(f"  base_url: {base_url}\n")
            wrote_base_url = True
        else:
            new_lines.append(line)
    # Append any sections that were missing
    last_model_key = None
    for i, line in enumerate(new_lines):
        if line.lstrip().startswith("model:") and line.strip() == "model:":
            last_model_key = i
    if not wrote_default:
        new_lines.append(f"  default: {model_name}\n")
    if not wrote_provider:
        new_lines.append(f"  provider: {provider}\n")
    if not wrote_base_url:
        new_lines.append(f"  base_url: {base_url}\n")
    with open(CONFIG_PATH, "w") as f:
        f.writelines(new_lines)
    log(f"Updated config: default={model_name} provider={provider} base_url={base_url}")


def main():
    if len(sys.argv) < 2:
        sys.exit(0)

    model_name = sys.argv[1]
    target = next((m for m in MODELS if m["model_name"] == model_name), None)
    if not target:
        log(f"Unknown model: {model_name}")
        send_notification(f"❌ Unknown model: {model_name}")
        sys.exit(1)

    name = target["name"]
    port = target["port"]
    service = target["service"]

    log(f"Starting switch to {name} (port {port})")

    # ── Check if already running with correct model ──
    r = subprocess.run(
        ["systemctl", "is-active", service],
        capture_output=True, text=True, timeout=10,
    )
    if r.stdout.strip() == "active":
        update_config(port, target["model_name"], target["provider"])
        log(f"{name} already running")
        # Still verify the model is actually responsive
        ready, elapsed = wait_for_model(port, timeout_sec=10)
        if ready:
            vram = get_vram()
            msg = (
                f"✅ **{name}** already running\n"
                f"━━━━━━━━━━━━━━━━━\n"
                f"**Port:** {port}\n"
                f"**VRAM:** {vram}\n"
                f"**Model:** {target['model_name']}\n"
            )
            send_notification(msg)
        sys.exit(0)

    # ── Stop any other running local service ──
    for m in MODELS:
        if m["service"] == service:
            continue
        r = subprocess.run(
            ["systemctl", "is-active", m["service"]],
            capture_output=True, text=True, timeout=5,
        )
        if r.stdout.strip() == "active":
            log(f"Stopping {m['name']}")
            subprocess.run(
                ["sudo", "systemctl", "stop", m["service"]],
                capture_output=True, timeout=30,
            )

    # ── Start the target service ──
    log(f"Starting {name}...")
    r = subprocess.run(
        ["sudo", "systemctl", "start", service],
        capture_output=True, timeout=180,
    )

    if r.returncode != 0:
        err = r.stderr.strip()[:200] if r.stderr else "unknown error"
        log(f"Failed to start {name}: {err}")
        send_notification(f"❌ **{name}** failed to start\n```\n{err}\n```")
        sys.exit(1)

    # ── Update config (model.default, model.provider, base_url) ──
    # Do this BEFORE waiting for the model so the systemd path watcher
    # sees the correct target model and doesn't revert the switch.
    update_config(port, target["model_name"], target["provider"])

    # ── Wait for model to load and respond ──
    log(f"Waiting for {name} to respond on port {port}...")
    ready, elapsed = wait_for_model(port, timeout_sec=120)

    # ── Get VRAM usage ──
    vram = get_vram()
    ctx = f"{target.get('context', 0):,}" if target.get("context") else "?"

    if ready:
        msg = (
            f"✅ **{name}** loaded\n"
            f"━━━━━━━━━━━━━━━━━\n"
            f"**Port:** {port}\n"
            f"**VRAM:** {vram}\n"
            f"**Context:** {ctx} tokens\n"
            f"**Load time:** {elapsed:.1f}s\n"
        )
        log(f"Ready in {elapsed:.1f}s — {vram}")
        send_notification(msg)
    else:
        msg = (
            f"⚠️ **{name}** started but not responding\n"
            f"━━━━━━━━━━━━━━━━━\n"
            f"**Port:** {port}\n"
            f"**Service:** active\n"
            f"**Waited:** {elapsed:.0f}s — still loading\n"
        )
        log(f"Started but not responding after {elapsed:.0f}s")
        send_notification(msg)

    sys.exit(0)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        log(f"Unhandled error: {e}")
        with open(NOTIFY_FILE, "w") as f:
            f.write(f"❌ Model switch error: {e}\n")
        sys.exit(1)