"""Consume real GGUF tokens and native P3Q/Jupiter receipts from JetBase."""
import json
import os
from getpass import getpass
from urllib.request import Request, urlopen
from urllib.error import HTTPError

BASE_URL = os.environ.get("JETBASE_URL", "https://jetbase-production.up.railway.app").rstrip("/")


def sse_events(response):
    event, data, event_id = "message", [], None
    for raw in response:
        line = raw.decode("utf-8").rstrip("\r\n")
        if not line:
            if data:
                yield event, json.loads("\n".join(data)), event_id
            event, data = "message", []
        elif line.startswith("event:"):
            event = line[6:].strip()
        elif line.startswith("data:"):
            data.append(line[5:].lstrip())
        elif line.startswith("id:"):
            event_id = line[3:].strip()


def run_twin_rail(prompt, key=None, targets=("phi2", "qwen"), run_id=None,
                  include_lattice=True, on_event=None, max_tokens=128):
    """One concurrent token stream; native lattice results arrive as typed events.

    Token generation and the native Rex P3Q computation are distinct operations.
    Native receipts carry the documented models, source hashes and material IDs.
    """
    key = key or os.environ.get("AURA_OPERATOR_KEY") or getpass("Existing AURA operator key: ")
    payload = {"prompt": prompt, "targets": list(targets), "max_tokens": max_tokens,
               "include_lattice": include_lattice}
    if run_id:
        payload["run_id"] = run_id
    request = Request(BASE_URL + "/api/clinical/quantum/stream", method="POST",
                      data=json.dumps(payload).encode(), headers={
                          "Content-Type": "application/json", "Accept": "text/event-stream",
                          "x-aura-operator-key": key})
    state = {"phi2": "", "qwen": "", "receipts": [], "errors": [], "completed": False}
    try:
        with urlopen(request, timeout=930) as response:
            if "text/event-stream" not in response.headers.get("Content-Type", ""):
                raise RuntimeError("Expected a live SSE response")
            for event, value, event_id in sse_events(response):
                if event == "token":
                    seat = value["seat"]
                    content = value.get("content") or value.get("reasoning_content") or ""
                    state[seat] += content
                    if on_event is None and content:
                        print(f"[{seat}] {content}", end="", flush=True)
                elif event in ("lattice", "jupiter"):
                    state["receipts"].append({"event": event, "data": value})
                elif event == "error":
                    state["errors"].append(value.get("error", "stream error"))
                elif event == "done":
                    state["completed"] = value.get("ok") is True
                if on_event:
                    on_event(event, value, event_id)
    except HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"JetBase HTTP {error.code}: {detail}") from error
    if state["errors"] or not state["completed"]:
        raise RuntimeError("; ".join(state["errors"]) or "Stream ended before verified completion")
    return state


def notebook_twin_rail(prompt, **options):
    """Optional notebook widgets; no background worker touches notebook UI."""
    import ipywidgets as widgets
    from IPython.display import display
    rails = {seat: widgets.Textarea(description=seat, layout={"width": "48%", "height": "220px"})
             for seat in ("phi2", "qwen")}
    status = widgets.Label(value="Connecting live twin rail…")
    receipt = widgets.Textarea(description="Receipts", layout={"width": "100%", "height": "220px"})
    display(status, widgets.HBox(list(rails.values())), receipt)

    def update(event, value, event_id):
        if event == "token":
            rails[value["seat"]].value += value.get("content") or value.get("reasoning_content") or ""
        elif event in ("lattice", "jupiter"):
            receipt.value += json.dumps({"event": event, "data": value}, indent=2) + "\n"
        elif event == "done":
            status.value = "Verified completion" if value.get("ok") else "Execution requires attention"
        elif event == "error":
            status.value = value.get("error", "Stream error")
    return run_twin_rail(prompt, on_event=update, **options)