# -*- coding: utf-8 -*-
"""POCKET-26B (Gemma4-based) CPU chat + live A/B vs Bonsai.
FastAPI frontend proxying to two local llama.cpp servers (upstream prebuilt release).
POCKET-26B speaks the Gemma-4 chat format; Bonsai speaks ChatML. Native /completion,
manual prompt build (no --jinja). No GPU — CPU + RAM only."""
import os, json, time, re
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse

BACKEND = os.environ.get("BACKEND", "http://127.0.0.1:8080")               # POCKET-26B (Gemma4)
BONSAI_BACKEND = os.environ.get("BONSAI_BACKEND", "http://127.0.0.1:8081")  # Bonsai (ChatML)
HERE = os.path.dirname(__file__)

app = FastAPI(title="POCKET-26B CPU chat")


def build_gemma(messages):
    """Gemma-4 chat prompt. llama.cpp adds <bos> from the GGUF metadata."""
    parts = []
    for m in messages:
        role = "user" if m.get("role") == "user" else "model"
        parts.append("<start_of_turn>%s\n%s<end_of_turn>\n" % (role, m.get("content", "").strip()))
    parts.append("<start_of_turn>model\n")
    return "".join(parts)


def build_chatml(messages):
    """Qwen-family ChatML prompt (Bonsai)."""
    parts = []
    if not any(m.get("role") == "system" for m in messages):
        parts.append("<|im_start|>system\nYou are a helpful assistant. Answer concisely.<|im_end|>\n")
    for m in messages:
        parts.append("<|im_start|>%s\n%s<|im_end|>\n" % (m.get("role", "user"), m.get("content", "")))
    parts.append("<|im_start|>assistant\n")
    return "".join(parts)


GEMMA_STOPS = ["<end_of_turn>", "<start_of_turn>", "<end_turn>", "<start_turn>"]
TEMPLATES = {
    "gemma":  (build_gemma,  GEMMA_STOPS),
    "chatml": (build_chatml, ["<|im_end|>", "<|im_start|>"]),
}


def _sse(obj):
    return "data: " + json.dumps(obj) + "\n\n"


async def _health(backend):
    try:
        async with httpx.AsyncClient(timeout=3.0) as c:
            r = await c.get(backend + "/health")
        return r.status_code == 200
    except Exception:
        return False


# Gemma4 wraps answers in a reasoning channel whose markers vary in this GGUF
# (sometimes <|channel>thought…<channel|>, sometimes <thought…\n). We extract the
# answer that follows the reasoning preamble so the UI shows clean text, never raw markers.
_JUNK = ["<|channel>", "<channel|>", "<|message>", "<tool_call|>", "<|start|>",
         "<|end|>", "<end_of_turn>", "<start_of_turn>", "<end_turn>", "<start_turn>",
         "<thought", "thought>"]


# catches partial/hallucinated turn & channel fragments the Q2 model sometimes emits
# (e.g. a stray "of_turn>" from a split <end_of_turn>), while leaving normal words alone.
_FRAG_RE = re.compile(r'[<|/a-z_]*_turn[|>]+|<\|?channel[|>]*|<thought')


def _scrub(s):
    for j in _JUNK:
        s = s.replace(j, "")
    return _FRAG_RE.sub("", s)


def _extract_gemma_answer(raw):
    """Return the answer portion of a (possibly partial) Gemma4 generation."""
    for close in ("<channel|>", "<|message>"):
        i = raw.rfind(close)
        if i != -1:
            return _scrub(raw[i + len(close):]).strip("\n ")
    head = raw.lstrip()
    looks_reasoning = ("<|channel>" in raw) or ("<thought" in raw) \
        or head[:10].lower().startswith(("channel", "thought"))
    if looks_reasoning:
        nl = raw.find("\n")            # answer begins after the thought label's newline
        if nl != -1:
            return _scrub(raw[nl + 1:]).strip("\n ")
        return ""                      # still inside the preamble
    return _scrub(raw).strip("\n ")    # plain answer, no reasoning wrapper


async def run_one(backend, tag, messages, max_tokens, template, temperature=0.7):
    """Stream one model's /completion output as SSE, tagged with `tag` (None = untagged).
    For Gemma (POCKET-26B) the reasoning wrapper is stripped so the UI shows a clean answer."""
    build, stops = TEMPLATES[template]
    payload = {
        "prompt": build(messages),
        "temperature": float(temperature),
        "top_p": 0.92,
        "min_p": 0.05,
        "repeat_penalty": 1.18,   # Q2 free-form generation falls into repetition loops without this
        "repeat_last_n": 256,
        "n_predict": int(max_tokens),
        "stream": True,
        "cache_prompt": True,
        "stop": stops,
    }
    n_tok = 0
    t_first = None
    tagd = {"m": tag} if tag else {}
    strip_channel = (template == "gemma")
    raw = ""
    sent_len = 0
    emitted = 0
    async with httpx.AsyncClient(timeout=None) as c:
        async with c.stream("POST", backend + "/completion", json=payload) as r:
            if r.status_code != 200:
                detail = (await r.aread()).decode("utf-8", "ignore")[:200]
                yield _sse(dict(tagd, error="backend %s: %s" % (r.status_code, detail)))
                return
            async for line in r.aiter_lines():
                if not line or not line.startswith("data:"):
                    continue
                try:
                    obj = json.loads(line[5:].strip())
                except Exception:
                    continue
                delta = obj.get("content")
                if delta:
                    if t_first is None:
                        t_first = time.monotonic()
                    n_tok += 1  # count every generated token (true generation speed)
                    if strip_channel:
                        raw += delta
                        clean = _extract_gemma_answer(raw)
                        if len(clean) > sent_len:
                            yield _sse(dict(tagd, t=clean[sent_len:]))
                            sent_len = len(clean)
                            emitted += 1
                    else:
                        emitted += 1
                        yield _sse(dict(tagd, t=delta))
                if obj.get("stop"):
                    if strip_channel and emitted == 0 and raw.strip():
                        fb = _scrub(raw).strip()   # never blank
                        if fb:
                            yield _sse(dict(tagd, t=fb))
                    tim = obj.get("timings") or {}
                    tok_s = tim.get("predicted_per_second")
                    if not tok_s and t_first is not None and n_tok > 0:
                        dt = time.monotonic() - t_first
                        tok_s = (n_tok / dt) if dt > 0 else 0.0
                    yield _sse(dict(tagd, done=True, tok_s=round(tok_s or 0.0, 1),
                                    n=int(tim.get("predicted_n") or n_tok)))
                    return
    yield _sse(dict(tagd, done=True, tok_s=0.0, n=n_tok))


@app.get("/api/status")
async def status():
    pk = await _health(BACKEND)
    bn = await _health(BONSAI_BACKEND)
    if pk:
        detail = "POCKET-26B on CPU · ready" if bn else "POCKET ready · loading Bonsai for the A/B…"
    else:
        detail = "downloading & loading models on CPU… (first boot can take a few minutes)"
    return JSONResponse({"state": "ready" if pk else "loading", "detail": detail,
                         "pocket": pk, "bonsai": bn})


@app.post("/api/chat")
async def chat(req: Request):
    body = await req.json()
    messages = body.get("messages") or [{"role": "user", "content": body.get("prompt", "Hi")}]
    mt = int(body.get("max_tokens", 512))

    async def gen():
        try:
            async for chunk in run_one(BACKEND, None, messages, mt, "gemma", body.get("temperature", 0.7)):
                yield chunk
        except Exception as e:
            yield _sse({"error": str(e)[:200]})

    return StreamingResponse(gen(), media_type="text/event-stream")


@app.post("/api/chat_ab")
async def chat_ab(req: Request):
    """Sequential race: Bonsai runs first, then POCKET-26B (each gets the full CPU)."""
    body = await req.json()
    messages = body.get("messages") or [{"role": "user", "content": body.get("prompt", "Hi")}]
    mt = int(body.get("max_tokens", 512))

    async def gen():
        try:
            async for chunk in run_one(BONSAI_BACKEND, "bonsai", messages, mt, "chatml"):
                yield chunk
            async for chunk in run_one(BACKEND, "pocket", messages, mt, "gemma"):
                yield chunk
            yield _sse({"ab_done": True})
        except Exception as e:
            yield _sse({"error": str(e)[:200]})

    return StreamingResponse(gen(), media_type="text/event-stream")


@app.post("/api/raw")
async def raw(req: Request):
    """DEBUG: return the model's UNstripped output. Accepts `raw_prompt` to bypass the
    template builder so prompt/prefill experiments need no rebuild."""
    body = await req.json()
    if body.get("raw_prompt") is not None:
        prompt = body["raw_prompt"]
        stops = body.get("stop", GEMMA_STOPS)
    else:
        messages = body.get("messages") or [{"role": "user", "content": body.get("prompt", "Hi")}]
        build, stops = TEMPLATES[body.get("template", "gemma")]
        prompt = build(messages)
    payload = {"prompt": prompt, "temperature": float(body.get("temperature", 0.7)),
               "top_p": 0.92, "min_p": 0.05,
               "repeat_penalty": float(body.get("repeat_penalty", 1.18)), "repeat_last_n": 256,
               "n_predict": int(body.get("max_tokens", 60)), "stream": False,
               "cache_prompt": True, "stop": stops}
    async with httpx.AsyncClient(timeout=None) as c:
        r = await c.post(BACKEND + "/completion", json=payload)
    j = r.json()
    return JSONResponse({"raw": j.get("content", ""),
                         "extracted": _extract_gemma_answer(j.get("content", "")),
                         "stop": j.get("stopping_word", "")})


@app.get("/", response_class=HTMLResponse)
def index():
    with open(os.path.join(HERE, "index.html"), encoding="utf-8") as f:
        return f.read()