"""OpenAI-RVS-Wrapper — der ARIA-LLM-Cluster als OpenAI-Endpunkt auf 127.0.0.1.

Laeuft als kleiner Container NEBEN einem Programm, das nur "OpenAI-kompatibel"
kann (OpenHands, Continue, Open WebUI …). Fuer das Programm ist der Cluster
dann einfach http://127.0.0.1:8099/v1. Der Wrapper selbst braucht nur eine
ausgehende Verbindung zum RVS — wie jede KI-Box auch.

  Programm ──HTTP──▶ Wrapper ──RVS──▶ ARIA-Bridge (Router) ──RVS──▶ KI-Box / Karte

  GET  /                     Hilfsseite: welche Modelle gibt es, was eintragen
  GET  /v1/models            jedes geladene Modell EINMAL (egal auf wie vielen Karten)
  POST /v1/chat/completions  auch mit stream=true und tools
  GET  /health               RVS verbunden? Antwortet die ARIA-Bridge?

Welche Box/Karte antwortet, entscheidet die ARIA-Bridge (Last, Ausfall-Ersatz).
Fehler kommen als HTTP-Fehler mit OpenAI-Fehlerobjekt zurueck, nie als normale
Antwort — ein Programm wuerde einen Fehlertext sonst als Aussage der KI in
seinen Verlauf schreiben.

Schnittstelle dahinter: docs/LLM-SCHNITTSTELLE.md
"""

from __future__ import annotations

import asyncio
import hmac
import json
import logging
import os
import ssl
import time
import uuid

import websockets

logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"),
                    format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("openai-rvs-wrapper")

RVS_HOST = os.getenv("RVS_HOST", "").strip()
RVS_PORT = os.getenv("RVS_PORT", "443").strip()
RVS_TLS = os.getenv("RVS_TLS", "true").strip().lower() == "true"
RVS_TOKEN = os.getenv("RVS_TOKEN", "").strip()
PORT = int(os.getenv("WRAPPER_PORT", "8099"))
API_KEY = os.getenv("WRAPPER_API_KEY", "").strip()
REQUEST_TIMEOUT_S = float(os.getenv("REQUEST_TIMEOUT_S", "660"))   # > CLUSTER_LLM_TIMEOUT_S der Bridge
MODELS_CACHE_S = float(os.getenv("MODELS_CACHE_S", "10"))
MAX_BODY = 32 * 1024 * 1024


# ── RVS ─────────────────────────────────────────────────────────────

class Rvs:
    """Eine Verbindung, ein Leser. Antworten werden per (Typ, requestId) an
    die wartende Anfrage gegeben — der RVS-Raum ist ein Verteiler, wir sehen
    auch Nachrichten fuer andere Teilnehmer und ignorieren sie."""

    def __init__(self):
        self.ws = None
        self.wait: dict = {}
        self.connected_since = 0.0

    @property
    def connected(self) -> bool:
        return self.ws is not None

    async def run(self) -> None:
        scheme = "wss" if RVS_TLS else "ws"
        url = f"{scheme}://{RVS_HOST}:{RVS_PORT}/?token={RVS_TOKEN}"
        backoff = 1
        while True:
            try:
                ctx = ssl.create_default_context() if RVS_TLS else None
                async with websockets.connect(url, ssl=ctx, max_size=64 * 1024 * 1024,
                                              ping_interval=20, ping_timeout=20) as ws:
                    self.ws, self.connected_since, backoff = ws, time.time(), 1
                    logger.info("RVS verbunden: %s://%s:%s", scheme, RVS_HOST, RVS_PORT)
                    async for raw in ws:
                        try:
                            m = json.loads(raw)
                        except Exception:
                            continue
                        p = m.get("payload") or {}
                        f = self.wait.get((m.get("type"), p.get("requestId")))
                        if f and not f.done():
                            f.set_result(p)
            except Exception as e:
                logger.warning("RVS-Verbindung weg (%s) — neuer Versuch in %ds", e, backoff)
            finally:
                self.ws = None
                for f in list(self.wait.values()):
                    if not f.done():
                        f.set_exception(ConnectionError("RVS-Verbindung unterbrochen"))
            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, 30)

    async def call(self, mtype: str, answer: str, payload: dict, timeout: float) -> dict:
        if not self.ws:
            raise ConnectionError("keine Verbindung zum RVS")
        rid = str(uuid.uuid4())
        fut = asyncio.get_running_loop().create_future()
        self.wait[(answer, rid)] = fut
        try:
            await self.ws.send(json.dumps({"type": mtype, "payload": dict(payload, requestId=rid),
                                           "timestamp": int(time.time() * 1000)}))
            return await asyncio.wait_for(fut, timeout)
        finally:
            self.wait.pop((answer, rid), None)


rvs = Rvs()
_models_cache = {"at": 0.0, "models": None}


async def cluster_models() -> list:
    """[{model, ctx, available, cards}] — jedes Modell einmal. Kurz gecacht,
    damit ein Programm, das /v1/models dauernd fragt, den Cluster nicht flutet."""
    if _models_cache["models"] is not None and time.time() - _models_cache["at"] < MODELS_CACHE_S:
        return _models_cache["models"]
    res = await rvs.call("llm_cluster_models_request", "llm_cluster_models", {}, 6)
    _models_cache.update(at=time.time(), models=res.get("models") or [])
    return _models_cache["models"]


# ── HTTP ────────────────────────────────────────────────────────────

def oai_error(message: str, code: str, etype: str = "invalid_request_error") -> dict:
    return {"error": {"message": message, "type": etype, "param": None, "code": code}}


async def reply(writer, status: int, payload, ctype="application/json; charset=utf-8", headers=None):
    body = payload if isinstance(payload, bytes) else json.dumps(payload, ensure_ascii=False).encode("utf-8")
    reason = {200: "OK", 400: "Bad Request", 401: "Unauthorized", 404: "Not Found", 413: "Payload Too Large",
              502: "Bad Gateway", 503: "Service Unavailable", 504: "Gateway Timeout"}.get(status, "Error")
    head = (f"HTTP/1.1 {status} {reason}\r\nContent-Type: {ctype}\r\nContent-Length: {len(body)}\r\n"
            "Connection: close\r\nAccess-Control-Allow-Origin: *\r\n")
    for k, v in (headers or {}).items():
        head += f"{k}: {v}\r\n"
    writer.write((head + "\r\n").encode("utf-8") + body)
    await writer.drain()


def authorized(headers: dict) -> bool:
    if not API_KEY:
        return True
    auth = headers.get("authorization", "")
    token = auth[7:].strip() if auth.lower().startswith("bearer ") else ""
    return bool(token) and hmac.compare_digest(token, API_KEY)


async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
    try:
        line = await asyncio.wait_for(reader.readline(), 15)
        if not line:
            return
        try:
            method, path, _ = line.decode("utf-8", "ignore").strip().split(" ", 2)
        except ValueError:
            return await reply(writer, 400, oai_error("bad request line", "bad_request"))
        headers = {}
        while True:
            h = await asyncio.wait_for(reader.readline(), 10)
            if not h or h in (b"\r\n", b"\n"):
                break
            k, _, v = h.decode("utf-8", "ignore").partition(":")
            headers[k.strip().lower()] = v.strip()
        n = int(headers.get("content-length", "0") or "0")
        if n > MAX_BODY:
            return await reply(writer, 413, oai_error("Anfrage zu gross", "too_large"))
        body = await reader.readexactly(n) if n else b""
        path = path.split("?", 1)[0].rstrip("/") or "/"

        if method == "OPTIONS":
            return await reply(writer, 200, b"", "text/plain", {
                "Access-Control-Allow-Headers": "Authorization, Content-Type",
                "Access-Control-Allow-Methods": "GET, POST, OPTIONS"})
        if method == "GET" and path == "/":
            return await reply(writer, 200, HELP_PAGE.encode("utf-8"), "text/html; charset=utf-8")
        if method == "GET" and path == "/health":
            bridge_ok = False
            if rvs.connected:
                try:
                    await rvs.call("llm_cluster_models_request", "llm_cluster_models", {}, 4)
                    bridge_ok = True
                except Exception:
                    pass
            return await reply(writer, 200 if bridge_ok else 503, {
                "rvs": rvs.connected, "bridge": bridge_ok,
                "rvsUrl": f"{'wss' if RVS_TLS else 'ws'}://{RVS_HOST}:{RVS_PORT}", "apiKey": bool(API_KEY)})
        if method == "GET" and path == "/api/models":
            # Fuer die Hilfsseite. Nur Modellnamen — kein Schutz noetig.
            try:
                return await reply(writer, 200, {"ok": True, "models": await cluster_models()})
            except Exception as e:
                return await reply(writer, 503, {"ok": False, "error": _unreachable(e)})

        if not path.startswith("/v1/"):
            return await reply(writer, 404, oai_error(f"{method} {path} gibt es nicht", "not_found"))
        if not authorized(headers):
            return await reply(writer, 401, oai_error("API-Key fehlt oder ist falsch (Authorization: Bearer <KEY>).",
                                                      "invalid_api_key", "authentication_error"))

        if method == "GET" and path == "/v1/models":
            try:
                models = await cluster_models()
            except Exception as e:
                return await reply(writer, 503, oai_error(_unreachable(e), "cluster_unreachable", "server_error"))
            now = int(time.time())
            return await reply(writer, 200, {"object": "list", "data": [
                {"id": m["model"], "object": "model", "created": now, "owned_by": "aria-cluster",
                 "context_length": m.get("ctx") or None} for m in models]})

        if method == "POST" and path == "/v1/chat/completions":
            try:
                req = json.loads(body.decode("utf-8", "ignore") or "{}")
            except Exception as e:
                return await reply(writer, 400, oai_error(f"kein gueltiges JSON: {e}", "bad_json"))
            payload = {"model": req.get("model") or "", "messages": req.get("messages"),
                       "max_tokens": req.get("max_completion_tokens") or req.get("max_tokens"),
                       "temperature": req.get("temperature"), "stop": req.get("stop"), "tools": req.get("tools")}
            try:
                res = await rvs.call("llm_cluster_request", "llm_cluster_response",
                                     {k: v for k, v in payload.items() if v is not None}, REQUEST_TIMEOUT_S)
            except asyncio.TimeoutError:
                return await reply(writer, 504, oai_error(
                    "Der Cluster hat nicht rechtzeitig geantwortet.", "timeout", "server_error"))
            except Exception as e:
                return await reply(writer, 503, oai_error(_unreachable(e), "cluster_unreachable", "server_error"))
            if not res.get("ok"):
                code = res.get("errorCode") or "model_error"
                status = {"bad_request": 400, "model_unavailable": 404, "timeout": 504}.get(code, 502)
                return await reply(writer, status, oai_error(
                    res.get("error") or "Fehler", "model_not_found" if code == "model_unavailable" else code,
                    "invalid_request_error" if status < 500 else "server_error"))
            return await completion(writer, req, res)

        return await reply(writer, 404, oai_error(f"{method} {path} gibt es nicht", "not_found"))
    except (asyncio.IncompleteReadError, asyncio.TimeoutError, ConnectionError):
        pass
    except Exception:
        logger.exception("HTTP-Fehler")
        try:
            await reply(writer, 502, oai_error("interner Fehler im Wrapper", "internal", "server_error"))
        except Exception:
            pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


def _unreachable(e: Exception) -> str:
    if not rvs.connected:
        return f"Keine Verbindung zum RVS ({RVS_HOST}:{RVS_PORT}) — RVS_HOST/RVS_TOKEN in der .env pruefen."
    if isinstance(e, asyncio.TimeoutError):
        return "RVS verbunden, aber die ARIA-Bridge antwortet nicht (laeuft ARIA? stimmt das RVS_TOKEN?)."
    return str(e)


async def completion(writer, req: dict, res: dict) -> None:
    cid = "chatcmpl-" + uuid.uuid4().hex[:24]
    created = int(time.time())
    model = res.get("model") or req.get("model") or ""
    content = res.get("content") or ""
    tcs = res.get("tool_calls") or None
    finish = "tool_calls" if tcs else "stop"
    usage = res.get("usage") or {}
    extra = {"X-Aria-Node": str(res.get("node", "")), "X-Aria-Gpu": str(res.get("gpu", "")),
             "X-Aria-Attempts": str(res.get("attempts", 1))}
    if not req.get("stream"):
        msg = {"role": "assistant", "content": content if (content or not tcs) else None}
        if tcs:
            msg["tool_calls"] = tcs
        return await reply(writer, 200, {
            "id": cid, "object": "chat.completion", "created": created, "model": model,
            "choices": [{"index": 0, "message": msg, "finish_reason": finish}], "usage": usage}, headers=extra)
    # stream=true: die Boxen liefern (noch) am Stueck — die fertige Antwort geht
    # im Stream-Format in einem Rutsch raus. Funktioniert mit jedem Client.
    head = ("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream; charset=utf-8\r\nCache-Control: no-cache\r\n"
            "Connection: close\r\nAccess-Control-Allow-Origin: *\r\n"
            + "".join(f"{k}: {v}\r\n" for k, v in extra.items()) + "\r\n")
    writer.write(head.encode("utf-8"))

    def chunk(delta: dict, fin=None) -> bytes:
        return ("data: " + json.dumps({"id": cid, "object": "chat.completion.chunk", "created": created,
                                       "model": model, "choices": [{"index": 0, "delta": delta,
                                                                    "finish_reason": fin}]},
                                      ensure_ascii=False) + "\n\n").encode("utf-8")
    delta = {"role": "assistant", "content": content}
    if tcs:
        delta["tool_calls"] = [dict(tc, index=i) for i, tc in enumerate(tcs)]
    writer.write(chunk(delta))
    writer.write(chunk({}, finish))
    if usage and (req.get("stream_options") or {}).get("include_usage"):
        writer.write(("data: " + json.dumps({"id": cid, "object": "chat.completion.chunk", "created": created,
                                             "model": model, "choices": [], "usage": usage}) + "\n\n").encode())
    writer.write(b"data: [DONE]\n\n")
    await writer.drain()


HELP_PAGE = """<!doctype html>
<html lang="de"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>ARIA LLM-Cluster</title>
<style>
 body{margin:0;background:#0A0A14;color:#E0E0F0;font:14px/1.5 -apple-system,"Segoe UI",Roboto,sans-serif}
 main{max-width:760px;margin:0 auto;padding:24px 16px}
 h1{font-size:20px;margin:0 0 4px} .sub{color:#8888AA;margin:0 0 20px}
 .card{background:#12121E;border:1px solid #2A2A3E;border-radius:10px;padding:14px 16px;margin-bottom:14px}
 .k{color:#8888AA;font-size:12px} code{background:#0A0A14;border:1px solid #2A2A3E;border-radius:5px;padding:2px 6px}
 table{width:100%;border-collapse:collapse} td,th{padding:7px 6px;text-align:left;border-bottom:1px solid #1E1E2E}
 th{color:#8888AA;font-weight:500;font-size:12px} button{background:#1E1E2E;color:#E0E0F0;border:1px solid #333;
 border-radius:5px;padding:3px 9px;cursor:pointer;font:inherit;font-size:12px} button:hover{border-color:#0096FF}
 .ok{color:#34C759}.err{color:#FF6B6B}.mute{color:#555570}
</style></head><body><main>
<h1>ARIA LLM-Cluster</h1>
<p class="sub">OpenAI-Endpunkt für Programme auf diesem Rechner — dahinter alle Modelle, die auf den KI-Boxen geladen sind.</p>
<div class="card"><div class="k">Status</div><div id="status" class="mute">prüfe …</div></div>
<div class="card">
 <div class="k">Im Programm eintragen (OpenAI-kompatibel)</div>
 <p style="margin:6px 0">Base URL: <code id="base"></code> <button onclick="copy(document.getElementById('base').textContent)">kopieren</button></p>
 <p style="margin:6px 0">API-Key: <span id="keyinfo" class="mute"></span></p>
 <p style="margin:6px 0" class="k">Läuft das Programm selbst in Docker, statt 127.0.0.1 den Namen
 <code>openai-rvs-wrapper</code> nehmen und das Netz <code>aria-llm</code> anhängen (siehe README).</p>
</div>
<div class="card"><div class="k">Modelle im Cluster — Name genau so eintragen</div>
 <table><thead><tr><th>Modell</th><th>Kontext</th><th>Karten</th><th></th></tr></thead><tbody id="models">
 <tr><td colspan="4" class="mute">lade …</td></tr></tbody></table></div>
<p class="mute" style="font-size:12px">Auf welcher Box und Karte ein Auftrag läuft, entscheidet der Cluster — bei
Last verteilt, bei Ausfall einer Box übernimmt eine andere mit demselben Modell.</p>
</main><script>
const esc=s=>String(s??'').replace(/[&<>"']/g,c=>({'&':'&amp;','<':'&lt;','>':'&gt;','"':'&quot;',"'":'&#39;'}[c]));
function copy(t){navigator.clipboard&&navigator.clipboard.writeText(t)}
document.getElementById('base').textContent=location.origin+'/v1';
async function load(){
 try{const h=await (await fetch('/health')).json();
  document.getElementById('status').innerHTML=h.bridge?'<span class="ok">✓ verbunden mit dem Cluster</span> <span class="mute">('+esc(h.rvsUrl)+')</span>'
   :(h.rvs?'<span class="err">RVS verbunden, aber die ARIA-Bridge antwortet nicht</span>':'<span class="err">keine Verbindung zum RVS — .env prüfen</span>');
  document.getElementById('keyinfo').innerHTML=h.apiKey?'der WRAPPER_API_KEY aus der .env':'beliebig (kein Key gesetzt — der Wrapper ist nur lokal erreichbar)';
 }catch(e){document.getElementById('status').innerHTML='<span class="err">Wrapper antwortet nicht</span>'}
 try{const d=await (await fetch('/api/models')).json();const tb=document.getElementById('models');
  if(!d.ok){tb.innerHTML='<tr><td colspan="4" class="err">'+esc(d.error)+'</td></tr>';return}
  if(!d.models.length){tb.innerHTML='<tr><td colspan="4" class="mute">Kein Modell geladen — in der ARIA-Diagnostic unter Einstellungen → LLM-Karten.</td></tr>';return}
  tb.innerHTML=d.models.map(m=>'<tr><td><code>'+esc(m.model)+'</code></td><td>'+(m.ctx?esc(m.ctx)+' Tokens':'<span class="mute">?</span>')
   +'</td><td>'+esc(m.available)+'</td><td><button onclick="copy(\\''+esc(m.model)+'\\')">kopieren</button></td></tr>').join('');
 }catch(e){}
}
load();setInterval(load,10000);
</script></body></html>
"""


async def main() -> None:
    if not RVS_HOST or not RVS_TOKEN:
        logger.error("RVS_HOST und RVS_TOKEN muessen in der .env stehen — Abbruch")
        return
    asyncio.create_task(rvs.run())
    server = await asyncio.start_server(handle, "0.0.0.0", PORT)
    logger.info("OpenAI-Endpunkt auf :%d (%s) — Hilfsseite: http://127.0.0.1:%d/", PORT,
                "mit API-Key" if API_KEY else "ohne API-Key", PORT)
    async with server:
        await server.serve_forever()


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