"""stream.py: answers questions with a language model on your own server and streams each answer,
as it is written, to the member who asked: to every tab and device of that user, in Python. A
member can stop an answer, and a page that lost its connection gets what it missed when it is
back. Any server with an OpenAI-compatible chat API works: Ollama, llama.cpp, vLLM. We use Gemma 4
with Ollama. Python 3.10+. pip install -r requirements.txt, then python stream.py (reads .env).

A connection may send 10 client events a second, so the text goes out in pieces: every 125 ms the
answer that has waited longest gets one event with all it has written since.

Its key (dashboard, New API key): rights subscribe and client_events, kind presence,
channels: prefix presence-ask-.
"""

import asyncio
import json
import threading
import time
import urllib.request

from pikoman_bot import Bot, env

ROOMS = env("ROOMS", "presence-ask-lobby").split(",")
AI_URL = env("AI_URL", "http://127.0.0.1:11434/v1")  # Ollama's OpenAI-compatible API
AI_MODEL = env("AI_MODEL", "gemma4:12b")
AT_ONCE = 2  # answers the model writes at the same time
SEND_EVERY = 0.125  # seconds between two events from the bot: 8 a second, under the limit of 10
MAX_PIECE = 6000  # characters in one event, well under the 12 KiB an event may hold
KEEP_FOR = 600  # seconds a finished answer is kept, for pages that come back

bot = Bot("Assistant", access=3)
writing = asyncio.Semaphore(AT_ONCE)
answers = {}  # "<room>/<id>" -> Answer, the one that has waited longest first


class Answer:
    """Who asked, the text so far, how much of it has gone out, and where each piece ended."""

    def __init__(self, answer_id, room, to, owner, question):
        self.id, self.room, self.to, self.owner, self.question = answer_id, room, to, owner, question
        self.text = ""
        self.sent = 0
        self.pieces = []  # where each event ended, so a page that comes back gets what follows
        self.done = False
        self.finished = False
        self.finished_at = 0.0
        self.stopped = False
        self.error = None
        self.stop = threading.Event()


def stream_model(answer, add):
    """Streams one answer from the model (in a thread: urllib waits), calling add(text) as it comes."""
    body = json.dumps({
        "model": AI_MODEL,
        "stream": True,
        "messages": [
            {"role": "system", "content": "Answer the user's question clearly and briefly."},
            {"role": "user", "content": answer.question},
        ],
    }).encode()
    headers = {"Content-Type": "application/json"}
    if env("AI_KEY"):
        headers["Authorization"] = "Bearer " + env("AI_KEY")
    request = urllib.request.Request(f"{AI_URL}/chat/completions", data=body, headers=headers, method="POST")
    with urllib.request.urlopen(request, timeout=60) as response:
        # Server-sent events: lines of "data: {json}", and "data: [DONE]" at the end.
        for raw in response:
            if answer.stop.is_set():
                return
            line = raw.decode().strip()
            if not line.startswith("data:"):
                continue
            data = line[5:].strip()
            if data == "[DONE]":
                return
            choices = json.loads(data).get("choices") or [{}]
            text = (choices[0].get("delta") or {}).get("content") or ""
            if text:
                add(text)


def append(answer, text):
    answer.text += text


async def job(answer):
    async with writing:
        if answer.done:
            return  # stopped before the model got to it
        loop = asyncio.get_running_loop()
        try:
            await asyncio.to_thread(stream_model, answer, lambda text: loop.call_soon_threadsafe(append, answer, text))
        except Exception as e:  # noqa: BLE001: the next question may go better
            if not answer.stop.is_set():
                print("model:", e)
                answer.error = "The model did not answer."
    if not answer.done:
        answer.done = True
        answer.finished_at = time.time()


async def send_piece(answer):
    """Sends the next piece of one answer to the user who asked: every tab and device of theirs."""
    text = answer.text[answer.sent:answer.sent + MAX_PIECE]
    last = answer.done and answer.sent + len(text) == len(answer.text)
    piece = {"id": answer.id, "seq": len(answer.pieces) + 1, "text": text}
    if last:
        piece["done"] = True
        if answer.stopped:
            piece["stopped"] = True
        if answer.error:
            piece["error"] = answer.error
    try:
        await answer.room.send_to(answer.to, "answer", piece)
    except Exception as e:  # noqa: BLE001: the connection is down; the text waits for it
        return print("send:", e)
    answer.sent += len(text)
    answer.pieces.append(answer.sent)
    if last:
        answer.finished = True


async def sender():
    """One event every SEND_EVERY seconds, for the answer that has waited longest and has something new."""
    while True:
        await asyncio.sleep(SEND_EVERY)
        if bot.id is None:
            continue  # not connected: the text waits until the bot is back
        now = time.time()
        for key in list(answers):
            answer = answers[key]
            if answer.finished:
                if now - answer.finished_at > KEEP_FOR:
                    del answers[key]
                continue
            if answer.sent < len(answer.text) or answer.done:
                answers[key] = answers.pop(key)  # to the back of the line
                await send_piece(answer)
                break


def assistant(room):
    pending = set()  # tasks still with the model, kept so they are not collected

    def owner_of(sender, info):
        return info.get("user_id") or sender

    def find(data):
        answer_id = str((data if isinstance(data, dict) else {}).get("id", ""))
        return answers.get(f"{room.name}/{answer_id}")

    @room.on("ask")
    async def ask(sender, data, info):
        """client-ask {id, text}: a question. The answer goes to the user who asked, on every tab."""
        data = data if isinstance(data, dict) else {}
        answer_id = str(data.get("id", ""))[:64]
        question = str(data.get("text", "")).strip()[:2000]
        if not answer_id or not question or room.is_bot(sender):
            return
        owner = owner_of(sender, info)
        to = {"user_id": info["user_id"]} if info.get("user_id") else sender
        key = f"{room.name}/{answer_id}"
        if key in answers:
            return  # asked twice: the first answer stands
        busy = any(not a.done and a.owner == owner for a in answers.values())
        if busy or len(pending) > 100:
            return await room.send_to(to, "answer", {"id": answer_id, "seq": 1, "text": "", "done": True, "error": "Busy: ask again in a moment."})
        answer = answers[key] = Answer(answer_id, room, to, owner, question)
        task = asyncio.create_task(job(answer))
        pending.add(task)
        task.add_done_callback(pending.discard)

    @room.on("stop")
    def stop(sender, data, info):
        """client-stop {id}: the model stops writing; what it wrote so far goes out, marked stopped."""
        answer = find(data)
        if not answer or answer.owner != owner_of(sender, info) or answer.done:
            return
        answer.stopped = True
        answer.done = True
        answer.finished_at = time.time()
        answer.stop.set()

    @room.on("resume")
    async def resume(sender, data, info):
        """client-resume {id, after}: the pieces after `after`, to that page alone, as many as fit."""
        answer = find(data)
        if not answer or answer.owner != owner_of(sender, info):
            return
        try:
            after = int(float(data.get("after") or 0))
        except (TypeError, ValueError):
            after = 0
        after = max(0, min(after, len(answer.pieces)))
        start = 0 if after == 0 else answer.pieces[after - 1]
        seq = after
        while seq < len(answer.pieces) and (seq == after or answer.pieces[seq] - start <= MAX_PIECE):
            seq += 1
        caught = {"id": answer.id, "seq": seq, "text": answer.text[start:answer.pieces[seq - 1] if seq else 0], "from": after}
        if seq < len(answer.pieces):
            caught["more"] = True
        elif answer.finished:
            caught["done"] = True
            if answer.stopped:
                caught["stopped"] = True
            if answer.error:
                caught["error"] = answer.error
        await room.send_to(sender, "answer", caught)


for name in ROOMS:
    assistant(bot.join(name))

print(f"Answering in {', '.join(ROOMS)} with {AI_MODEL}")
bot.start(sender())