# -*- coding: utf-8 -*-
"""POCKET-35B CPU chat + live A/B vs Bonsai.
FastAPI frontend proxying to two local llama.cpp servers (upstream prebuilt release,
supports POCKET's qwen35moe arch). Native /completion + manual ChatML. No GPU."""
import os, json, time
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
BONSAI_BACKEND = os.environ.get("BONSAI_BACKEND", "http://127.0.0.1:8081")  # Bonsai
HERE = os.path.dirname(__file__)

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


def build_chatml(messages, no_think=False):
    """Qwen-family ChatML prompt (works for POCKET and Bonsai).
    When no_think is set, the assistant turn is prefilled with a closed (empty)
    <think></think> block so the model skips reasoning and answers directly — the
    reliable way to get short, fast replies for the A/B race (the /no_think soft
    switch is not honored through a raw prompt)."""
    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", "")))
    if no_think:
        parts.append("<|im_start|>assistant\n<think>\n\n</think>\n\n")
    else:
        parts.append("<|im_start|>assistant\n")
    return "".join(parts)


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


async def run_one(backend, tag, messages, max_tokens, temperature=0.7, no_think=False):
    """Stream one model's /completion output as SSE, tagged with `tag` (None = untagged)."""
    payload = {
        "prompt": build_chatml(messages, no_think=no_think),
        "temperature": float(temperature),
        "top_p": 0.95,
        "n_predict": int(max_tokens),
        "stream": True,
        "cache_prompt": True,
        "stop": ["<|im_end|>", "<|im_start|>"],
    }
    n_tok = 0
    t_first = None
    tagd = {"m": tag} if tag else {}
    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
                    yield _sse(dict(tagd, t=delta))
                if obj.get("stop"):
                    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-35B on CPU · ready" if bn else "POCKET ready · loading Bonsai for the A/B…"
    else:
        detail = "downloading & loading models… (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, 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 (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", 200))

    async def gen():
        try:
            async for chunk in run_one(BONSAI_BACKEND, "bonsai", messages, mt):
                yield chunk
            async for chunk in run_one(BACKEND, "pocket", messages, mt):
                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.get("/", response_class=HTMLResponse)
def index():
    with open(os.path.join(HERE, "index.html"), encoding="utf-8") as f:
        return f.read()