"""
ARIA Local-LLM-Adapter (Gamebox) — 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: Gamebox 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"))
# 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"


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) -> dict:
    """Ruft llama.cpp /v1/chat/completions (OpenAI-Format). Gibt
    {ok, content, tool_calls, error} zurueck — wirft nie.

    tools: optionale OpenAI-Tool-Definitionen (B1b). llama.cpp (--jinja) mit
    Qwen3 kann natives Tool-Calling und liefert dann message.tool_calls."""
    body = {
        "model": 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 {}
        return {
            "ok": True,
            "content": msg.get("content") or "",
            "tool_calls": msg.get("tool_calls") or None,
            "usage": data.get("usage"),
        }
    except Exception as e:
        logger.warning("llama.cpp-Call fehlgeschlagen: %s", e)
        return {"ok": False, "content": "", "error": str(e)[:300]}


async def _handle_llm_request(ws, payload: dict) -> None:
    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
    t0 = time.time()
    res = await _call_llama(messages, max_tokens=max_tokens,
                            temperature=temperature, stop=stop, tools=tools)
    dt = time.time() - t0
    tc = res.get("tool_calls")
    logger.info("llm_request id=%s -> ok=%s %.2fs content_len=%d tool_calls=%d",
                (req_id[:8] if req_id else "?"), 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": LLM_MODEL,
        "elapsedMs": int(dt * 1000),
    })


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

    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}/ws?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
                async for raw in ws:
                    try:
                        msg = json.loads(raw)
                    except Exception:
                        continue
                    if msg.get("type") != "llm_request":
                        continue
                    payload = msg.get("payload", {}) or {}
                    # 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)
            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())