"""LFM3-1.7B with its inference modules behind one OpenAI-compatible endpoint: for a client it is just a model.

Every request goes through the modules measured in REPORT.md, chosen by the kind of request:
  * code (asks for a program / function): deliberation + test-and-repair on the examples in the statement;
  * visual (draw / build an SVG, a web page, a canvas animation, a three.js scene): deliberation + render in headless
    Chrome + repair of what the browser reports (needs playwright and Chrome);
  * multiple choice (>= 3 lettered options): deliberation with temporal consensus;
  * math / short-answer question: deliberation (answer read-out at the budget);
  * anything else (writing, explanations): deliberation with the free-form budget fallback.
Model "lfm3-deliberate-x4" (or "parallel": k in the request) adds parallel thinking: k independent reasonings and the
weighted vote (math / choice only; AIME test 52.5 -> 66.7 at 4x tokens).

The backend is a llama-server with the model; each request takes 2 of its slots (2k with parallel thinking), so
-np sets how many requests run at once.
usage:
  python lfm3_server.py --start-backend gguf/lfm3-orig-q5km.gguf --slots 16      # starts llama-server too
  python lfm3_server.py --backend http://localhost:8080                           # uses a running one
  then POST http://localhost:8090/v1/chat/completions {"model": "lfm3-deliberate", "messages": [...]}
"""
import argparse
import json
import os
import re
import subprocess
import threading
import time
import urllib.request
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from lfm3_deliberate import Deliberator

HERE = os.path.dirname(os.path.abspath(__file__))

CODE_STRONG = re.compile(r"```python|\bdef \w+\(|standard input|\bstdin\b|>>>|\bassert \w", re.I)
CODE_ASK = re.compile(r"\b(write|implement|complete|fix|create)\b[^\n]{0,80}\b(python|program|function)\b|"
                      r"(напиши|напишите|реализуй|реализуйте|исправь|создай)[^\n]{0,80}(python|программ|функци|скрипт|код)",
                      re.I)
CODE_HINT = re.compile(r"\b(python|input|output|returns?|integers?|array|list|string)\b|"
                       r"(python|возвращ|список|массив|строк|ввод|вывод|чисел|число)", re.I)
MATH_TEX = re.compile(r"\\boxed|\\frac|\$[^$\n]+\$|\d\s*[+*/^×÷]\s*\(?\d")
MATH_WORDS = re.compile(r"\b(compute|calculate|evaluate|simplify|solve|how many|how much|find the (value|number|sum|"
                        r"remainder|product|smallest|largest|least|greatest))\b|(вычисли|найди|найдите|сколько|реши|"
                        r"решите|докажи)", re.I)


VISUAL = re.compile(r"\b(svg|html|css|canvas|three\.js|web ?page|landing page|website)\b|(сайт|страниц|верстк)", re.I)
VISUAL_ASK = re.compile(r"\b(draw|build|make|create|design|render|animate)\b|(нарисуй|сделай|создай|сверстай|построй)",
                        re.I)


def route(text):
    """Picks the module configuration that was measured for this kind of request (REPORT.md)."""
    if len(re.findall(r"^\s*\(?[A-J][.)]\s+\S", text, re.M)) >= 3:
        return "choice"
    if VISUAL.search(text) and VISUAL_ASK.search(text) and not re.search(r"\bpython\b", text, re.I):
        return "visual"
    if CODE_STRONG.search(text) or (CODE_ASK.search(text) and CODE_HINT.search(text)):
        return "code"
    if re.search(r"answer is:?\s*<", text):  # "end your response with 'The answer is: <answer>'"
        return "phrase"
    if len(text) < 4000 and (MATH_TEX.search(text) or (MATH_WORDS.search(text) and re.search(r"\d", text))):
        return "math"
    return "free"


class SlotPool:
    def __init__(self, n):
        self.free, self.cv = list(range(n)), threading.Condition()

    def take(self, k):
        with self.cv:
            self.cv.wait_for(lambda: len(self.free) >= k)
            got, self.free = self.free[:k], self.free[k:]
            return got

    def give(self, slots):
        with self.cv:
            self.free += slots
            self.cv.notify_all()


def start_backend(gguf, slots, ctx, bin_dir, port):
    # llama-server scans its working directory for backends; the project path is not ASCII, so run it from its own
    cmd = [os.path.join(bin_dir, "llama-server.exe" if os.name == "nt" else "llama-server"), "-m", os.path.abspath(gguf),
           "-ngl", "99", "-np", str(slots), "-c", str(slots * ctx), "-fa", "on", "--port", str(port), "--no-webui"]
    os.makedirs(os.path.join(HERE, "results"), exist_ok=True)
    log = open(os.path.join(HERE, "results", "server_lfm3.log"), "w")
    proc = subprocess.Popen(cmd, cwd=bin_dir, stdout=log, stderr=subprocess.STDOUT)
    for _ in range(300):
        try:
            if b'"ok"' in urllib.request.urlopen(f"http://localhost:{port}/health", timeout=2).read():
                return proc
        except Exception:
            pass
        if proc.poll() is not None:
            raise SystemExit(f"llama-server exited, see results/server_lfm3.log")
        time.sleep(1)
    raise SystemExit("llama-server did not become ready")


def make_handler(delib, pool, args):
    class H(BaseHTTPRequestHandler):
        def _send(self, code, obj):
            data = json.dumps(obj, ensure_ascii=False).encode()
            self.send_response(code)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def log_message(self, fmt, *a):
            pass

        def do_GET(self):
            if self.path.rstrip("/") in ("/v1/models", "/models"):
                return self._send(200, {"object": "list", "data": [
                    {"id": m, "object": "model", "owned_by": "local"} for m in ("lfm3-deliberate", "lfm3-deliberate-x4")]})
            if self.path.rstrip("/") == "/health":
                return self._send(200, {"status": "ok"})
            self._send(404, {"error": "not found"})

        def do_POST(self):
            if self.path.rstrip("/") not in ("/v1/chat/completions", "/chat/completions"):
                return self._send(404, {"error": "not found"})
            try:
                body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))))
                msgs = [{"role": m["role"], "content": m["content"] if isinstance(m["content"], str) else
                         "".join(p.get("text", "") for p in m["content"])} for m in body["messages"]]
            except Exception as e:
                return self._send(400, {"error": f"bad request: {e}"})
            last = next((m["content"] for m in reversed(msgs) if m["role"] == "user"), "")
            kind = body.get("lfm3_kind") or route(last)
            k = int(body.get("parallel") or (4 if str(body.get("model", "")).endswith("-x4") else 1))
            k = k if kind in ("math", "choice") else 1
            budget = int(body.get("max_tokens") or body.get("max_completion_tokens") or args.budget)
            slots = pool.take(2 * k)
            t0 = time.time()
            try:
                import copy
                d = copy.copy(delib)
                d.budget = max(2048, min(budget, args.budget))
                d.reserve = min(d.reserve, d.budget // 4)
                d.main_slot, d.probe_slot = slots[0], slots[1]
                seed = int(body.get("seed", 0))
                if kind == "code":
                    r = d.solve_code(msgs, seed=seed)
                    info = {"mode": r["mode"], "public_tests": r["public"], "repair_rounds": r["rounds"]}
                elif kind == "visual":
                    r = d.solve_visual(msgs, seed=seed)
                    info = {"mode": r["mode"], "render_problems": r["problems"], "repair_rounds": r["rounds"]}
                elif k > 1:
                    r = d.ask_parallel(msgs, kind, k=k, seed=seed, slots=slots)
                    info = {"votes": r["votes"], "modes": [o["mode"] for o in r["outputs"]],
                            "thinking_tokens": sum(o["thinking_tokens"] for o in r["outputs"])}
                else:
                    r = d.ask(msgs, kind, seed=seed)
                    info = {"mode": r["mode"], "thinking_tokens": r["thinking_tokens"]}
            except Exception as e:
                return self._send(500, {"error": f"{type(e).__name__}: {e}"})
            finally:
                pool.give(slots)
            info.update(kind=kind, parallel=k, seconds=round(time.time() - t0, 1))
            answer = r["answer"]
            rid, created = "chatcmpl-" + uuid.uuid4().hex[:24], int(time.time())
            model = body.get("model", "lfm3-deliberate")
            if body.get("stream"):
                self.send_response(200)
                self.send_header("Content-Type", "text/event-stream")
                self.end_headers()
                for delta, fin in (({"role": "assistant", "content": answer}, None), ({}, "stop")):
                    chunk = {"id": rid, "object": "chat.completion.chunk", "created": created, "model": model,
                             "choices": [{"index": 0, "delta": delta, "finish_reason": fin}]}
                    self.wfile.write(b"data: " + json.dumps(chunk, ensure_ascii=False).encode() + b"\n\n")
                self.wfile.write(b"data: [DONE]\n\n")
                return
            self._send(200, {"id": rid, "object": "chat.completion", "created": created, "model": model,
                             "choices": [{"index": 0, "message": {"role": "assistant", "content": answer},
                                          "finish_reason": "stop"}],
                             "usage": {"completion_tokens": info.get("thinking_tokens", 0)}, "lfm3": info})
    return H


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--backend", default="http://localhost:8080")
    ap.add_argument("--start-backend", metavar="GGUF", help="start llama-server with this model first")
    import shutil
    found = shutil.which("llama-server")
    ap.add_argument("--llama-bin", default=os.environ.get("LLAMA_BIN") or (os.path.dirname(found) if found else "."),
                    help="folder with llama-server (default: $LLAMA_BIN or the one on PATH)")
    ap.add_argument("--slots", type=int, default=16, help="llama-server slots (2 per request, 2k with parallel k)")
    ap.add_argument("--slot-ctx", type=int, default=20480)
    ap.add_argument("--budget", type=int, default=16384, help="max thinking tokens per reasoning")
    ap.add_argument("--port", type=int, default=8090)
    ap.add_argument("--tokenizer", default=os.path.join(HERE, "tokenizer"))
    a = ap.parse_args()
    if a.start_backend:
        start_backend(a.start_backend, a.slots, a.slot_ctx, a.llama_bin, int(a.backend.rsplit(":", 1)[1]))
    props = json.loads(urllib.request.urlopen(a.backend + "/props", timeout=10).read())
    n = int(props.get("total_slots", a.slots))
    print(f"backend {a.backend}: {n} slots -> up to {n // 2} requests at once; serving on http://localhost:{a.port}/v1",
          flush=True)
    srv = ThreadingHTTPServer(("127.0.0.1", a.port), make_handler(Deliberator(a.backend, a.tokenizer, budget=a.budget),
                                                                  SlotPool(n), a))
    srv.serve_forever()