"""Several chats at once on the Ling-mini-2.0 megakernel: a local web app.

    python server.py [--port 8790]

Then open http://127.0.0.1:8790. Every pane is its own conversation with its own cache. One thread
owns the GPU and works in rounds: each round, every chat that is reading a prompt or writing an
answer gets its own launch of the megakernel on a share of the SMs (82 / chats working, at most 5
at once; more wait their turn), all launches run together, and their words go out to the panes as
the round ends. One chat alone gets the whole GPU (~850 words/s); three together ~460 each.

Greedy decoding, thinking off, up to 4,096 tokens per conversation. Binds to 127.0.0.1 only.
"""

from __future__ import annotations

import argparse
import json
import queue
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import torch

from ling_kernel import HERE, NQ, LingKernel, State
CHUNK = 32           # answer tokens per launch
PREFILL = 64         # prompt tokens per launch (keeps rounds short while others are writing)
MAX_AT_ONCE = 5      # chats working in the same round (16 SMs each at 5: one attention block per head)
MAX_POS = 4096
MAX_NEW = 1024
PANES = 8            # conversations the page can open


class Chat:
    """One conversation: its GPU state (KV cache and buffers), its history, its listeners."""

    def __init__(self, mk: LingKernel, cid: int):
        self.mk, self.cid = mk, cid
        self.state = State(MAX_POS)
        self.out = self.state.out
        self.stream = torch.cuda.Stream()
        self.history: list[dict] = []
        self.cached: list[int] = []          # tokens that have cache entries, in order
        self.listeners: list[queue.Queue] = []
        self.job: dict | None = None         # the answer being written

    def launch(self, ins: list[int], pos0: int, n: int, forced: bool, nblocks: int) -> None:
        self.mk.launch_async(self.state, ins, pos0, n, forced, nblocks, self.stream)

    def send(self, event: dict) -> None:
        for q in list(self.listeners):
            q.put(event)


class Server:
    def __init__(self) -> None:
        from transformers import AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(HERE)
        print("loading the Ling-mini-2.0 megakernel (4-bit) ...", flush=True)
        self.mk = LingKernel(max_pos=MAX_POS, log=lambda s: print(s, flush=True))
        self.sms = self.mk.sms
        self.eos = {i for i in (self.tok.convert_tokens_to_ids("<|role_end|>"), self.tok.eos_token_id,
                                self.mk.config.get("eos_token_id")) if isinstance(i, int) and i >= 0}
        self.chats = [Chat(self.mk, i) for i in range(PANES)]
        self.scratch = Chat(self.mk, -1)     # warms the GPU's clocks without touching a conversation
        self.lock = threading.Lock()
        self.wake = threading.Event()
        self.recent: list[tuple[float, int]] = []    # (time, tokens) of finished rounds, for the total rate
        self.last_work = 0.0
        threading.Thread(target=self.loop, daemon=True).start()

    # ---- requests (HTTP threads) ----
    def start(self, cid: int, text: str) -> str | None:
        with self.lock:
            c = self.chats[cid]
            if c.job is not None:
                return "this chat is still writing"
            c.job = {"text": text, "state": "new"}
            c.send({"type": "user", "text": text})        # every open page shows it
        self.wake.set()
        return None

    def reset(self, cid: int) -> str | None:
        with self.lock:
            c = self.chats[cid]
            if c.job is not None:
                return "this chat is still writing"
            c.history, c.cached = [], []
            c.send({"type": "cleared"})
        return None

    def transcript(self, cid: int) -> dict:
        with self.lock:
            c = self.chats[cid]
            job = c.job
            msgs = list(c.history)
            if job is not None:
                msgs.append({"role": "user", "content": job["text"]})
                msgs.append({"role": "assistant", "content": job.get("shown", "")})
            return {"messages": msgs, "writing": job is not None}

    def total_rate(self) -> float:
        now = time.perf_counter()
        with self.lock:
            self.recent = [(t, n) for t, n in self.recent if now - t < 1.0]
            n = sum(k for _, k in self.recent)
        return n / 1.0

    # ---- the GPU loop (one thread) ----
    def begin(self, c: Chat) -> None:
        """Turn a new message into prompt tokens; keep the cache's common prefix."""
        job = c.job
        messages = c.history + [{"role": "user", "content": job["text"]}]
        ids = self.tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
        if hasattr(ids, "keys"):
            ids = ids["input_ids"]
        ids = list(ids)
        if len(ids) + 2 * CHUNK + 1 > MAX_POS:
            c.job = None
            c.send({"type": "error", "text": f"this conversation is full ({MAX_POS} tokens): press Clear"})
            return
        n = 0
        while n < min(len(ids) - 1, len(c.cached)) and ids[n] == c.cached[n]:
            n += 1
        job.update(state="read", messages=messages, ids=ids, pos=n, out=[], shown="", seconds=0.0, steps=0,
                   max_new=min(MAX_NEW, MAX_POS - len(ids) - CHUNK))
        c.send({"type": "start", "prompt_tokens": len(ids) - n})

    def finish(self, c: Chat, note: str = "") -> None:
        job = c.job
        text = job["shown"].strip()
        c.history = job["messages"] + [{"role": "assistant", "content": text}]
        c.cached = job["ids"] + job["out"][:-1]
        rate = job["steps"] / job["seconds"] if job["seconds"] else 0.0
        c.send({"type": "done", "tokens": len(job["out"]), "rate": rate, "note": note})
        c.job = None

    def warm(self) -> None:
        """After the GPU idled, Windows lowers its clocks: run the scratch chat until launches are fast."""
        start = time.perf_counter()
        best = float("inf")
        while time.perf_counter() - start < 2.0:
            t = time.perf_counter()
            self.scratch.launch([1000], 0, CHUNK, False, self.sms)
            torch.cuda.synchronize()
            d = time.perf_counter() - t
            if d > 0.97 * best and d < 1.2 * best:
                break
            best = min(best, d)

    def loop(self) -> None:
        while True:
            with self.lock:
                for c in self.chats:
                    if c.job is not None and c.job["state"] == "new":
                        self.begin(c)
                working = [c for c in self.chats if c.job is not None][:MAX_AT_ONCE]
            if not working:
                self.wake.wait(0.5)
                self.wake.clear()
                continue
            if time.perf_counter() - self.last_work > 1.0:
                self.warm()
            nb = max(NQ, self.sms // len(working))
            plan = []
            try:
                for c in working:
                    job = c.job
                    if job["state"] == "read":
                        k = min(PREFILL, len(job["ids"]) - job["pos"])
                        c.launch(job["ids"][job["pos"]:job["pos"] + k], job["pos"], k, True, nb)
                    else:
                        k = min(CHUNK, job["max_new"] - len(job["out"]))
                        c.launch([job["out"][-1]], job["pos"], k, False, nb)
                    plan.append((c, k))
                t = time.perf_counter()
                torch.cuda.synchronize()
                dt = time.perf_counter() - t
            except Exception as e:                        # report it and drop the round's jobs
                torch.cuda.synchronize()
                with self.lock:
                    for c, _ in plan or [(c, 0) for c in working]:
                        if c.job is not None:
                            c.send({"type": "error", "text": f"GPU error: {e}"})
                            c.job = None
                continue
            self.last_work = time.perf_counter()
            made = 0
            with self.lock:
                for c, k in plan:
                    job = c.job
                    outs = c.out[:k].tolist()
                    if job["state"] == "read":
                        job["pos"] += k
                        if job["pos"] == len(job["ids"]):    # the prompt is read: its last step's word is the first answer word
                            job["state"] = "write"
                            job["out"] = [outs[k - 1]]
                            made += 1
                            self.emit(c)
                            if outs[k - 1] in self.eos:
                                self.finish(c)
                        continue
                    job["seconds"] += dt
                    job["steps"] += k
                    job["pos"] += k
                    for t_ in outs:
                        job["out"].append(t_)
                        made += 1
                        if t_ in self.eos:
                            break
                    self.emit(c)
                    if job["out"][-1] in self.eos:
                        self.finish(c)
                    elif len(job["out"]) >= job["max_new"]:
                        self.finish(c, "stopped at the length limit")
                self.recent.append((time.perf_counter(), made))

    def emit(self, c: Chat) -> None:
        job = c.job
        words = [t for t in job["out"] if t not in self.eos]
        text = self.tok.decode(words, skip_special_tokens=True)
        if text.endswith("�"):                         # half of a character: wait for the rest
            text = text.rstrip("�")
        if len(text) > len(job["shown"]) and text.startswith(job["shown"]):
            c.send({"type": "delta", "text": text[len(job["shown"]):]})
            job["shown"] = text
        elif text != job["shown"] and not text.startswith(job["shown"]):
            c.send({"type": "replace", "text": text})
            job["shown"] = text


def handler(server: Server):
    page = (HERE / "multichat.html").read_bytes()

    class H(BaseHTTPRequestHandler):
        def log_message(self, *a):                          # quiet
            pass

        def reply(self, code: int, body: bytes, ctype: str = "application/json") -> None:
            self.send_response(code)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(len(body)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            u = urlparse(self.path)
            if u.path == "/":
                return self.reply(200, page, "text/html; charset=utf-8")
            if u.path == "/history":
                cid = int(parse_qs(u.query).get("chat", ["0"])[0])
                if not 0 <= cid < PANES:
                    return self.reply(404, b"{}")
                return self.reply(200, json.dumps(server.transcript(cid)).encode())
            if u.path == "/status":
                with server.lock:
                    busy = [c.cid for c in server.chats if c.job is not None]
                return self.reply(200, json.dumps({"total_rate": server.total_rate(), "busy": busy,
                                                   "sms": server.sms, "panes": PANES}).encode())
            if u.path == "/events":
                cid = int(parse_qs(u.query).get("chat", ["0"])[0])
                if not 0 <= cid < PANES:
                    return self.reply(404, b"{}")
                q: queue.Queue = queue.Queue()
                c = server.chats[cid]
                with server.lock:
                    c.listeners.append(q)
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.send_header("Cache-Control", "no-store")
                self.end_headers()
                try:
                    self.wfile.write(b": open\n\n")
                    self.wfile.flush()
                    while True:
                        try:
                            ev = q.get(timeout=15)
                            self.wfile.write(f"data: {json.dumps(ev)}\n\n".encode())
                        except queue.Empty:
                            self.wfile.write(b": keep-alive\n\n")
                        self.wfile.flush()
                except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError, OSError):
                    pass
                finally:
                    with server.lock:
                        if q in c.listeners:
                            c.listeners.remove(q)
                return
            return self.reply(404, b"{}")

        def do_POST(self):
            u = urlparse(self.path)
            n = int(self.headers.get("Content-Length", "0") or 0)
            try:
                body = json.loads(self.rfile.read(n) or b"{}")
                cid = int(body.get("chat", -1))
            except (ValueError, TypeError):
                return self.reply(400, b'{"error": "bad request"}')
            if not 0 <= cid < PANES:
                return self.reply(404, b'{"error": "no such chat"}')
            if u.path == "/send":
                text = str(body.get("text", "")).strip()
                if not text:
                    return self.reply(400, b'{"error": "empty message"}')
                err = server.start(cid, text)
            elif u.path == "/reset":
                err = server.reset(cid)
            else:
                return self.reply(404, b"{}")
            if err:
                return self.reply(409, json.dumps({"error": err}).encode())
            return self.reply(200, b'{"ok": true}')

    return H


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8790)
    args = ap.parse_args()
    sys.stdout.reconfigure(encoding="utf-8")
    server = Server()
    httpd = ThreadingHTTPServer(("127.0.0.1", args.port), handler(server))
    httpd.daemon_threads = True
    print(f"open http://127.0.0.1:{args.port}  ({server.sms} SMs; up to {MAX_AT_ONCE} chats writing at once)", flush=True)
    httpd.serve_forever()


if __name__ == "__main__":
    main()