"""
ARIA Local-LLM-Adapter (AI-Box) — Plan B, Phase B0.

Bruecke zwischen RVS und dem lokalen llama.cpp-Server. Spiegelt das Muster der
whisper-bridge: verbindet sich per WebSocket mit dem RVS (Token-Room, TLS mit
ws-Fallback, Reconnect-Backoff), lauscht auf `llm_request` und ruft den lokalen
llama.cpp-`/v1/chat/completions`-Endpoint (OpenAI-kompatibel), antwortet mit
`llm_response` (korreliert per requestId).

Topologie: AI-Box steht zuhause, ARIA im RZ — die Kommunikation laeuft ueber
den RVS (wie TTS/STT), keine IPs zu pflegen. Nur URL + Token.

Env:
  RVS_HOST, RVS_PORT, RVS_TLS, RVS_TLS_FALLBACK, RVS_TOKEN  (wie f5tts/whisper)
  LLAMA_URL         Default http://llama:8081  (llama.cpp im selben Compose-Netz)
  LLM_MODEL         optionaler Modell-Name fuer llama (llama.cpp ignoriert ihn
                    meist, dient nur der Transparenz im Log)
  LLM_TIMEOUT_SEC   Default 60

Bewusst NICHT-streamend in B0 (volle llm_response). Token-Streaming (llm_partial)
kommt in B2 zusammen mit TTS-on-first-sentence.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import time

import httpx
import websockets

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("llm-adapter")

RVS_HOST = os.getenv("RVS_HOST", "").strip()
RVS_PORT = os.getenv("RVS_PORT", "443").strip()
RVS_TLS = os.getenv("RVS_TLS", "true").lower() == "true"
RVS_TLS_FALLBACK = os.getenv("RVS_TLS_FALLBACK", "true").lower() == "true"
RVS_TOKEN = os.getenv("RVS_TOKEN", "").strip()

LLAMA_URL = os.getenv("LLAMA_URL", "http://llama:8081").rstrip("/")
LLM_MODEL = os.getenv("LLM_MODEL", "qwen3-8b")
LLM_TIMEOUT_SEC = float(os.getenv("LLM_TIMEOUT_SEC", "60"))

# ── Compute-Fleet: Worker-Identitaet & Registrierung ──────────────
# Meldet sich bei der aria-bridge (worker_hello) + periodischer worker_ping.
NODE_NAME = os.getenv("NODE_NAME", "node").strip() or "node"
GPU_IDS = os.getenv("NVIDIA_VISIBLE_DEVICES", "").strip()
WORKER_SERVICE = "llm"
INSTANCE_ID = f"{WORKER_SERVICE}@{NODE_NAME}"
WORKER_PING_INTERVAL_S = int(os.getenv("WORKER_PING_INTERVAL_S", "10"))
# Empfangs-Watchdog: kommt in RX_STALE_S kein Broadcast rein (ein echter Raum hat
# staendig Traffic, z.B. sat_hello alle 25s / Brain-Polling), gilt die Verbindung
# als halb-tot (Caddy pongt die WS-Pings selbst) -> Zwangs-Reconnect.
RX_STALE_S = int(os.getenv("RX_STALE_S", "60"))
_inflight = 0   # laufende llm_requests (busy-Report im ping)
# Qwen3 hat Thinking-Mode default AN — dann verbraet es Tokens in einem
# <think>-Block und liefert (bei kleinem max_tokens) leeren/abgeschnittenen
# content, ausserdem 3x langsamer. ARIAs schnelles Tier will KEIN Grübeln
# (grübeln = harter Turn = Claude). Wir schalten Thinking daher per
# chat_template_kwargs ab (Qwen3-Template versteht enable_thinking=false;
# andere Templates ignorieren das kwarg). Bei einem Modell, das darauf
# empfindlich reagiert: LLM_DISABLE_THINKING=false setzen.
LLM_DISABLE_THINKING = os.getenv("LLM_DISABLE_THINKING", "true").lower() == "true"

# ── Modell-Verwaltung (Stage D): Adapter besitzt llama-swaps Config ──
# llama-swap liest die GENERIERTE Config (beschreibbar, im /models-Bind). Wir
# erzeugen sie aus dem Basis-Template (kuratierte Defaults) + der persistenten
# Box-Registry (per Diagnostic hinzugefuegte Modelle). So werden neue Modelle
# ohne Image-Rebuild waehlbar.
import yaml  # pyyaml
BASE_CONFIG_PATH = os.getenv("LLAMA_BASE_CONFIG", "/llamaswap/config.yaml")
GEN_CONFIG_PATH = os.getenv("LLAMA_GEN_CONFIG", "/models/llama-swap.config.yaml")
REGISTRY_PATH = os.getenv("LLM_REGISTRY", "/models/aria_models.json")

# ── Auslastungs-Monitor (Stage E) ──────────────────────────
import node_stats
STATS_PATH = os.getenv("STATS_PATH", "/models/aria_stats.json")
_total_tokens = 0   # kumulativ, fuer den Token-Graph
_stats = node_stats.NodeStats(INSTANCE_ID, NODE_NAME, STATS_PATH,
                              token_getter=lambda: _total_tokens, logger=logger)


def _load_registry() -> list:
    try:
        with open(REGISTRY_PATH) as f:
            data = json.load(f)
        return data if isinstance(data, list) else []
    except Exception:
        return []


def _save_registry(reg: list) -> None:
    try:
        tmp = REGISTRY_PATH + ".tmp"
        with open(tmp, "w") as f:
            json.dump(reg, f, indent=2)
        os.replace(tmp, REGISTRY_PATH)
    except Exception as e:
        logger.warning("Registry speichern fehlgeschlagen: %s", e)


def _generate_config() -> int:
    """Schreibt die llama-swap-Config aus Basis-Template + Registry. Gibt die
    Anzahl Modelle zurueck. Idempotent, bei jeder Aenderung + beim Start."""
    base = {}
    try:
        with open(BASE_CONFIG_PATH) as f:
            base = yaml.safe_load(f) or {}
    except Exception as e:
        logger.warning("Basis-Template %s nicht lesbar (%s)", BASE_CONFIG_PATH, e)
    models = dict(base.get("models") or {})
    for e in _load_registry():
        key = (e.get("key") or "").strip()
        repo = (e.get("hfRepo") or "").strip()
        if not key or not repo:
            continue
        quant = (e.get("quant") or "Q4_K_M").strip()
        ctx = int(e.get("ctx") or 8192)
        ngl = int(e.get("ngl") or 99)
        models[key] = {
            "cmd": (f"llama-server --port ${{PORT}} --host 127.0.0.1\n"
                    f"-hf {repo}:{quant}\n-ngl {ngl} -c {ctx} --jinja"),
            "ttl": 3600,
        }
    out = dict(base)
    out["models"] = models
    try:
        os.makedirs(os.path.dirname(GEN_CONFIG_PATH), exist_ok=True)
        tmp = GEN_CONFIG_PATH + ".tmp"
        with open(tmp, "w") as f:
            yaml.safe_dump(out, f, sort_keys=False, default_flow_style=False)
        os.replace(tmp, GEN_CONFIG_PATH)
        logger.info("llama-swap-Config generiert: %d Modelle → %s", len(models), GEN_CONFIG_PATH)
    except Exception as e:
        logger.error("Config schreiben fehlgeschlagen: %s", e)
    return len(models)


async def _reload_llama() -> None:
    """Stoesst llama-swap-Reload an. Viele Builds watchen die Config-Datei ohnehin;
    zusaetzlich versuchen wir bekannte Reload-Endpunkte (Fehler ignoriert)."""
    for path in ("/api/config/reload", "/reload"):
        try:
            async with httpx.AsyncClient(timeout=10) as c:
                r = await c.post(f"{LLAMA_URL}{path}")
                if r.status_code < 400:
                    logger.info("llama-swap reload via %s", path)
                    return
        except Exception:
            pass
    logger.info("llama-swap reload: kein Endpoint — verlasse mich auf File-Watch")


async def _send(ws, mtype: str, payload: dict) -> None:
    try:
        await ws.send(json.dumps({
            "type": mtype,
            "payload": payload,
            "timestamp": int(time.time() * 1000),
        }))
    except Exception as e:
        logger.warning("Send fehlgeschlagen (%s): %s", mtype, e)


async def _call_llama(messages: list, *, max_tokens: int, temperature: float,
                      stop, tools=None, model=None) -> dict:
    """Ruft llama.cpp/llama-swap /v1/chat/completions (OpenAI-Format). Gibt
    {ok, content, tool_calls, error} zurueck — wirft nie.

    model: welches Modell llama-swap laden soll (B0.5). Kommt aus dem Request
    (Brain -> local_llm.json). Faellt auf LLM_MODEL (env) zurueck.
    tools: optionale OpenAI-Tool-Definitionen (B1b). Qwen3 (--jinja) kann
    natives Tool-Calling und liefert dann message.tool_calls."""
    body = {
        "model": model or LLM_MODEL,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False,
    }
    if stop:
        body["stop"] = stop
    if tools:
        body["tools"] = tools
        body["tool_choice"] = "auto"
    if LLM_DISABLE_THINKING:
        # llama.cpp (--jinja) reicht chat_template_kwargs an die Chat-Vorlage
        # weiter. Qwen3 unterdrueckt damit den <think>-Block.
        body["chat_template_kwargs"] = {"enable_thinking": False}
    try:
        async with httpx.AsyncClient(timeout=LLM_TIMEOUT_SEC) as client:
            r = await client.post(f"{LLAMA_URL}/v1/chat/completions", json=body)
            r.raise_for_status()
            data = r.json()
        msg = (data.get("choices") or [{}])[0].get("message", {}) or {}
        usage = data.get("usage") or {}
        try:
            global _total_tokens
            _total_tokens += int(usage.get("total_tokens") or 0)
        except Exception:
            pass
        return {
            "ok": True,
            "content": msg.get("content") or "",
            "tool_calls": msg.get("tool_calls") or None,
            "usage": usage,
        }
    except Exception as e:
        logger.warning("llama.cpp-Call fehlgeschlagen: %s", e)
        return {"ok": False, "content": "", "error": str(e)[:300]}


# B0.5-2: Lade-Status ans Diagnostic (service_status, service="llm"). Wir kennen
# den Download-Fortschritt nicht (llama-swap gibt ihn nicht her), aber wir melden
# den Zustand bei Modellwechsel: loading -> ready/error. _last_model = aktuell
# geladenes; _ready_models = in dieser Session schon einmal bereit gewesene
# (fuer den "frisch geladen"-Hinweis 🎉 bei langem Erst-Load).
_last_model = None
_ready_models: set = set()


async def _emit_llm_status(ws, state: str, model: str, **extra) -> None:
    await _send(ws, "service_status",
                {"service": "llm", "state": state, "model": model, **extra})


async def _fetch_available_models() -> list:
    """Fragt llama-swap ab, welche Modelle diese Box fahren kann (GET /v1/models,
    OpenAI-kompatibel → {data:[{id},...]}). Das sind die config.yaml-Keys.
    Defensiv: bei Fehler Fallback auf [LLM_MODEL]."""
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            r = await client.get(f"{LLAMA_URL}/v1/models")
            r.raise_for_status()
            data = r.json()
        ids = [m.get("id") for m in (data.get("data") or []) if m.get("id")]
        return ids or [LLM_MODEL]
    except Exception as e:
        logger.warning("llama-swap /v1/models nicht abfragbar (%s) — Fallback [%s]", e, LLM_MODEL)
        return [LLM_MODEL]


async def _announce(ws) -> None:
    """Sendet ein frisches worker_hello mit der aktuellen Modell-Liste (nach
    Provision/Remove aufrufen, damit Bridge+Diagnostic das neue Modell lernen)."""
    models = await _fetch_available_models()
    await _send(ws, "worker_hello", {
        "instanceId": INSTANCE_ID, "service": WORKER_SERVICE,
        "node": NODE_NAME, "gpus": GPU_IDS, "model": LLM_MODEL,
        "models": models,   # welche Modelle diese Box fahren kann (llama-swap-Keys)
    })
    logger.info("worker_hello: models=%s", models)


async def _worker_register(ws) -> None:
    """Meldet diesen Worker bei der aria-bridge an (worker_hello) und haelt die
    Flotten-Registry per periodischem worker_ping (mit busy-Status) frisch."""
    try:
        await _announce(ws)
        n = 0
        while True:
            await asyncio.sleep(WORKER_PING_INTERVAL_S)
            n += 1
            await _send(ws, "worker_ping",
                        {"instanceId": INSTANCE_ID, "busy": _inflight > 0})
            if n % 3 == 0:   # ~30s worker_hello wiederholen (wie der Satellit) →
                await _announce(ws)   # auch neu gestartetes Diagnostic/Bridge lernt uns
    except asyncio.CancelledError:
        raise
    except Exception:
        return   # Socket tot → still beenden; _run reconnectet + startet neu


async def _handle_llm_request(ws, payload: dict) -> None:
    global _last_model, _inflight
    _inflight += 1
    try:
        await _do_llm_request(ws, payload)
    finally:
        _inflight -= 1


async def _do_llm_request(ws, payload: dict) -> None:
    global _last_model
    req_id = payload.get("requestId", "")
    messages = payload.get("messages") or []
    if not isinstance(messages, list) or not messages:
        await _send(ws, "llm_response", {
            "requestId": req_id, "ok": False, "error": "leere/ungueltige messages",
        })
        return
    max_tokens = int(payload.get("max_tokens", 512) or 512)
    temperature = float(payload.get("temperature", 0.7) or 0.7)
    stop = payload.get("stop")
    tools = payload.get("tools") or None
    model = (payload.get("model") or "").strip() or None
    eff_model = model or LLM_MODEL

    # Modellwechsel (oder erster Request) → llama-swap laedt/swappt: Status melden.
    switching = eff_model != _last_model
    if switching:
        await _emit_llm_status(ws, "loading", eff_model)

    t0 = time.time()
    res = await _call_llama(messages, max_tokens=max_tokens,
                            temperature=temperature, stop=stop, tools=tools,
                            model=model)
    dt = time.time() - t0

    if switching:
        if res.get("ok"):
            fresh = (eff_model not in _ready_models) and dt > 25
            _ready_models.add(eff_model)
            _last_model = eff_model
            await _emit_llm_status(ws, "ready", eff_model,
                                   loadSeconds=round(dt, 1), freshlyDownloaded=fresh)
        else:
            # bei Fehler _last_model NICHT setzen → naechster Versuch meldet erneut loading
            await _emit_llm_status(ws, "error", eff_model,
                                   error=(res.get("error") or "")[:120])
    tc = res.get("tool_calls")
    logger.info("llm_request id=%s model=%s -> ok=%s %.2fs content_len=%d tool_calls=%d",
                (req_id[:8] if req_id else "?"), model or LLM_MODEL, res.get("ok"), dt,
                len(res.get("content") or ""), len(tc) if tc else 0)
    await _send(ws, "llm_response", {
        "requestId": req_id,
        "ok": res.get("ok", False),
        "content": res.get("content", ""),
        "tool_calls": tc,
        "error": res.get("error"),
        "model": model or LLM_MODEL,
        "elapsedMs": int(dt * 1000),
    })


async def _handle_provision(ws, payload: dict) -> None:
    """Fuegt ein Modell hinzu: Registry+Config schreiben, reload, dann Warmup
    (zieht das GGUF via -hf beim ersten Load). Meldet die neue Modell-Liste."""
    key = (payload.get("key") or "").strip()
    repo = (payload.get("hfRepo") or "").strip()
    if not key or not repo:
        await _send(ws, "llm_provision_result",
                    {"instanceId": INSTANCE_ID, "key": key, "ok": False, "error": "key/hfRepo fehlt"})
        return
    entry = {
        "key": key, "hfRepo": repo,
        "quant": (payload.get("quant") or "Q4_K_M").strip(),
        "ctx": int(payload.get("ctx") or 8192),
        "ngl": int(payload.get("ngl") or 99),
    }
    reg = [e for e in _load_registry() if e.get("key") != key]
    reg.append(entry)
    _save_registry(reg)
    _generate_config()
    await _reload_llama()
    await _announce(ws)   # Bridge/Diagnostic lernen das neue Modell
    # Warmup: Mini-Request → llama-swap laedt/zieht das Modell (Fortschritt via
    # service_status loading→ready, freshlyDownloaded).
    await _emit_llm_status(ws, "loading", key)
    t0 = time.time()
    res = await _call_llama([{"role": "user", "content": "hi"}],
                            max_tokens=1, temperature=0.0, stop=None, model=key)
    dt = time.time() - t0
    if res.get("ok"):
        _ready_models.add(key)
        await _emit_llm_status(ws, "ready", key, loadSeconds=round(dt, 1),
                               freshlyDownloaded=dt > 25)
    else:
        await _emit_llm_status(ws, "error", key, error=(res.get("error") or "")[:160])
    await _send(ws, "llm_provision_result",
                {"instanceId": INSTANCE_ID, "key": key, "ok": res.get("ok", False),
                 "error": res.get("error"), "elapsedMs": int(dt * 1000)})
    logger.info("provision %s (%s) → ok=%s %.1fs", key, repo, res.get("ok"), dt)


async def _handle_remove(ws, payload: dict) -> None:
    """Entfernt ein Modell aus Registry+Config (GGUF bleibt im Cache)."""
    key = (payload.get("key") or "").strip()
    if not key:
        return
    reg = [e for e in _load_registry() if e.get("key") != key]
    _save_registry(reg)
    _generate_config()
    await _reload_llama()
    await _announce(ws)
    await _send(ws, "llm_provision_result",
                {"instanceId": INSTANCE_ID, "key": key, "ok": True, "removed": True})
    logger.info("removed model %s", key)


async def _run() -> None:
    if not RVS_HOST:
        logger.error("RVS_HOST nicht gesetzt — Abbruch")
        return
    if not RVS_TOKEN:
        logger.error("RVS_TOKEN nicht gesetzt — Abbruch")
        return

    # llama-swap-Config aus Basis-Template + Registry erzeugen, BEVOR llama-swap
    # sie braucht (llama-swap restart: unless-stopped faengt die Erst-Boot-
    # Reihenfolge ab, falls es kurz vor uns startet).
    _generate_config()

    # Auslastungs-Sampler (GPU + Tokens) laeuft unabhaengig vom RVS.
    asyncio.create_task(_stats.run_sampler())

    use_tls = RVS_TLS
    retry_s = 2
    tls_fallback_tried = False

    while True:
        scheme = "wss" if use_tls else "ws"
        url = f"{scheme}://{RVS_HOST}:{RVS_PORT}?token={RVS_TOKEN}"
        masked = url.replace(RVS_TOKEN, "***") if RVS_TOKEN else url
        try:
            logger.info("Verbinde zu RVS: %s (llama=%s)", masked, LLAMA_URL)
            async with websockets.connect(
                url, ping_interval=20, ping_timeout=10, max_size=16 * 1024 * 1024
            ) as ws:
                logger.info("RVS verbunden — llm-adapter online")
                retry_s = 2
                tls_fallback_tried = False
                ping_task = asyncio.create_task(_worker_register(ws))
                while True:
                    try:
                        raw = await asyncio.wait_for(ws.recv(), timeout=RX_STALE_S)
                    except asyncio.TimeoutError:
                        logger.warning("Kein RVS-Traffic seit %ds — Verbindung halb-tot, reconnect", RX_STALE_S)
                        raise ConnectionError("rvs-stale")
                    try:
                        msg = json.loads(raw)
                    except Exception:
                        continue
                    mtype = msg.get("type")
                    payload = msg.get("payload", {}) or {}
                    # Redundanz-Routing: gezielt an eine andere Instanz adressiert
                    # → ignorieren. Ohne targetInstance → wie bisher (jeder nimmt).
                    tgt = payload.get("targetInstance")
                    if tgt and tgt != INSTANCE_ID:
                        continue
                    # Auslastungs-Monitor (node_stats_*) abfangen.
                    if await _stats.handle(ws, mtype, payload, _send):
                        continue
                    if mtype not in ("llm_request", "llm_provision_model", "llm_remove_model"):
                        continue
                    if mtype == "llm_provision_model":
                        asyncio.create_task(_handle_provision(ws, payload))
                    elif mtype == "llm_remove_model":
                        asyncio.create_task(_handle_remove(ws, payload))
                    else:
                        # Jede Anfrage nebenlaeufig — llama.cpp serialisiert intern,
                        # aber wir blockieren so nicht den Empfang weiterer Messages.
                        asyncio.create_task(_handle_llm_request(ws, payload))
        except Exception as e:
            logger.warning("RVS-Verbindung verloren/fehlgeschlagen: %s", e)
            try:
                ping_task.cancel()
            except NameError:
                pass
            if use_tls and RVS_TLS_FALLBACK and not tls_fallback_tried:
                tls_fallback_tried = True
                use_tls = False
                logger.info("TLS fehlgeschlagen — Fallback auf ws://")
                continue
            await asyncio.sleep(min(retry_s, 30))
            retry_s = min(retry_s * 2, 30)
            use_tls = RVS_TLS


if __name__ == "__main__":
    asyncio.run(_run())