"""translator.py: translates each chat message with a language model on your own server and posts
the translation in the room, in Python. Any server with an OpenAI-compatible chat API works:
Ollama, llama.cpp, vLLM. We use Gemma 4 with Ollama. A connection may send 10 client events a
second, so when the room is busy the translations go out together, a few in each event.
Python 3.10+. pip install -r requirements.txt, then python translator.py (reads .env).

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

import asyncio
import json
import time
import urllib.request

from pikoman_bot import Bot, env

ROOMS = env("ROOMS", "presence-intl-lobby").split(",")
LANGUAGE = env("LANGUAGE", "English")
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  # questions to the model at the same time
SEND_EVERY = 0.125  # seconds between two events from the bot: 8 a second, under the limit of 10
USER_CHARS_PER_MINUTE = 5000  # what one user may send to the model; a model's time costs

bot = Bot("Translator", access=3)
asking = asyncio.Semaphore(AT_ONCE)
to_send = {}  # room name -> [{from, name, text, original}], oldest room first
used = {}  # what each user has sent to the model this minute, by user id, or by connection without one
used_minute = 0


def within_user_limit(user, text):
    global used_minute
    minute = int(time.time() // 60)
    if minute != used_minute:
        used_minute = minute
        used.clear()
    total = used.get(user, 0) + len(text)
    if total > USER_CHARS_PER_MINUTE:
        return False
    used[user] = total
    return True


def translate(text):
    """One message through the model (in a thread: urllib waits): the translation, or None."""
    body = json.dumps({
        "model": AI_MODEL,
        "temperature": 0,
        "messages": [
            {"role": "system", "content": f"Translate the user's chat message into {LANGUAGE}. Answer with the translation only. If it is already in {LANGUAGE}, answer with it unchanged."},
            {"role": "user", "content": text},
        ],
    }).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=30) as response:
        answer = json.load(response)["choices"][0]["message"]["content"].strip()
    return answer or None


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

    @room.on("message")
    def message(sender, data, info):
        text = str((data or {}).get("text", "")).strip()[:1000]
        if not text or room.is_bot(sender):
            return
        if len(pending) > 200:
            return print(f"{room.name}: too much to translate, a message skipped")
        member = room.member(sender)
        if not within_user_limit(str(member.get("user_id", sender)), text):
            return print(f"{room.name}: a user is over their limit, a message skipped")
        name = member.get("name", "Someone")
        task = asyncio.create_task(job(sender, name, text))
        pending.add(task)
        task.add_done_callback(pending.discard)

    async def job(sender, name, text):
        async with asking:
            try:
                translated = await asyncio.to_thread(translate, text)
            except Exception as e:  # noqa: BLE001: the next message may go better
                return print("translate:", e)
        if translated and translated != text:
            to_send.setdefault(room.name, []).append({"from": sender, "name": name, "text": translated, "original": text})


async def sender():
    """One event every SEND_EVERY seconds, to the room that has waited longest, with all it has."""
    while True:
        await asyncio.sleep(SEND_EVERY)
        if not to_send:
            continue
        name = next(iter(to_send))
        items = to_send.pop(name)
        await bot.rooms[name].send("translations", {"language": LANGUAGE, "items": items[:20]})
        if len(items) > 20:
            to_send[name] = items[20:]  # the rest next time (12 KiB per event)


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

print(f"Translating {', '.join(ROOMS)} into {LANGUAGE} with {AI_MODEL}")
bot.start(sender())