"""OpenAI-compatible API server for one model on the AMD Ryzen AI NPU (point other apps at http://HOST:PORT/v1).

The model runs in FastFlowLM (`flm serve`, the NPU-only engine Lemonade Server uses; see onw.lemonade), started
here as a child process on a free localhost port; this server forwards the OpenAI requests to it (as Lemonade's
WrappedServer does) and adds the chat UI, the check page and the API key.

  GET  /                      what this server is: model, OpenAI base URL, endpoints (JSON)
  GET  /v1/models
  GET  /v1/models/{id}
  POST /v1/chat/completions   (stream true/false)
  POST /v1/completions
  GET  /health                200 once the engine has loaded the model (503 while it loads)
  GET  /chat                  browser chat UI for trying the model
  GET  /check                 verification page: speed, image, multi-turn
  GET  /api/info              model, engine and NPU details

usage: onw serve MODEL [--port 8010] [--host 0.0.0.0] [--api-key KEY] [--context 4096] [--open]
MODEL is a model name from the onw window or `onw list` (llama3.2-1b-FLM, qwen3-4b-FLM, ...); it is downloaded
first when missing. CORS is open, so browser-based clients work too; with --api-key every /v1 request needs
"Authorization: Bearer KEY".
"""
import argparse, asyncio, atexit, json, os, subprocess, sys, threading, time

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse
from onw import __version__, i18n, lemonade

MODEL_ID = "npu-model"      # set in main()
app = FastAPI(title="LLM on the AMD NPU")
info = {}
WEB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web")
API_KEY = None              # --api-key: required as a Bearer token on /v1
BACKEND = {"url": None, "model": None, "proc": None}   # the engine: base URL, the model name it expects, its process
READY = threading.Event()
client = None               # httpx.AsyncClient, made on the first request (inside the event loop)


def oai_error(status, message, type_="invalid_request_error", code=None):
    """An error in OpenAI's shape: {"error": {"message", "type", "code"}}."""
    return JSONResponse({"error": {"message": message, "type": type_, "param": None, "code": code}}, status_code=status)


@app.exception_handler(HTTPException)
async def _http_error(request, e):
    return oai_error(e.status_code, str(e.detail), "invalid_request_error" if e.status_code < 500 else "server_error")


@app.exception_handler(RequestValidationError)
async def _validation_error(request, e):
    return oai_error(400, str(e))


@app.exception_handler(Exception)
async def _server_error(request, e):
    return oai_error(500, f"{type(e).__name__}: {e}", "server_error")


@app.middleware("http")
async def cors_and_key(request: Request, call_next):
    """Open CORS (browser-based OpenAI clients) and the optional API key on /v1."""
    if request.method == "OPTIONS":
        resp = JSONResponse({})
    elif API_KEY and request.url.path.startswith("/v1") and \
            request.headers.get("authorization", "") != f"Bearer {API_KEY}":
        resp = oai_error(401, "invalid API key", code="invalid_api_key")
    else:
        resp = await call_next(request)
    resp.headers["Access-Control-Allow-Origin"] = "*"
    resp.headers["Access-Control-Allow-Headers"] = "*"
    resp.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
    return resp


@app.get("/")
def index(request: Request):
    base = str(request.base_url).rstrip("/")
    return {"object": "onw.server", "name": "onw AMD - 俺のNPUがこんなに動くわけない", "version": __version__,
            "model": MODEL_ID, "engine": info.get("engine_name"), "openai_base_url": base + "/v1",
            "api_key_required": bool(API_KEY), "ready": READY.is_set(),
            "endpoints": {"models": "/v1/models", "chat_completions": "/v1/chat/completions",
                          "completions": "/v1/completions", "health": "/health", "chat_ui": "/chat",
                          "check_page": "/check", "info": "/api/info"},
            "context": info.get("context"), "vision": info.get("vision")}


@app.get("/health")
def health():
    if not READY.is_set():
        return JSONResponse({"status": "loading", "model": MODEL_ID}, status_code=503)
    return {"status": "ok", "model": MODEL_ID}


@app.get("/chat")
def chat_ui():
    return HTMLResponse(i18n.page(os.path.join(WEB, "index.html")))


@app.get("/check")
def check_page():
    return HTMLResponse(i18n.page(os.path.join(WEB, "check.html")))


@app.get("/api/info")
def api_info():
    return {**info, "ready": READY.is_set()}


STARTED = int(time.time())


def _model_card():
    return {"id": MODEL_ID, "object": "model", "created": STARTED, "owned_by": "local-npu",
            "engine": info.get("engine"), "labels": info.get("labels", [])}


@app.get("/v1/models")
def models():
    return {"object": "list", "data": [_model_card()]}


@app.get("/v1/models/{model_id:path}")
def model(model_id: str):
    if model_id != MODEL_ID:
        raise HTTPException(404, f"model {model_id!r} not found (this server runs {MODEL_ID!r})")
    return _model_card()


class ThinkSplit:
    """FastFlowLM writes the thinking into the content as <think>...</think>; OpenAI clients (and the chat UI) expect it
    in reasoning_content. Fed the content piece by piece (a tag may be cut between two pieces), returns
    (reasoning, content) for each piece."""
    # (opening, closing) tags: <think> (Qwen, DeepSeek, ...), Gemma 4's thought channel. flm splits the thinking off
    # itself for some models, but not from an answer cut off by max_tokens while still thinking
    TAGS = (("<think>", "</think>"), ("<|channel>thought", "<channel|>"))

    def __init__(self):
        self.state, self.buf = "start", ""            # start (nothing seen yet) -> think -> gap -> text, or start -> text
        self.CLOSE = self.TAGS[0][1]

    def feed(self, s, final=False):
        self.buf += s
        think = text = ""
        while True:
            if self.state == "start":                  # only a leading <think> opens the thinking
                head = self.buf.lstrip()
                hit = next(((o, c) for o, c in self.TAGS if head.startswith(o)), None)
                if hit:
                    self.CLOSE = hit[1]
                    self.state, self.buf = "think", head[len(hit[0]):]
                elif any(o.startswith(head) for o, _ in self.TAGS) and not final:
                    break                              # "", "<", "<thi": wait for more
                else:
                    self.state = "text"
            elif self.state == "think":
                i = self.buf.find(self.CLOSE)
                if i >= 0:
                    think += self.buf[:i]
                    self.state, self.buf = "gap", self.buf[i + len(self.CLOSE):]
                    continue
                keep = 0 if final else next((k for k in range(len(self.CLOSE) - 1, 0, -1)
                                             if self.buf.endswith(self.CLOSE[:k])), 0)
                think, self.buf = think + self.buf[:len(self.buf) - keep], self.buf[len(self.buf) - keep:]
                break
            elif self.state == "gap":                  # the line breaks between </think> and the answer
                self.buf = self.buf.lstrip("\n")
                if not self.buf:
                    break
                self.state = "text"
            else:
                text, self.buf = text + self.buf, ""
                break
        return think, text


def flm_request(body):
    """OpenAI / onw thinking switches -> FastFlowLM's "think". Always sent: flm keeps the last request's setting for
    the following ones, so a request without a switch would otherwise think when the one before asked for it."""
    if "think" in body:
        return
    on = (body.get("chat_template_kwargs") or {}).get("enable_thinking")
    effort = body.get("reasoning_effort")
    if isinstance(effort, str):
        on = effort not in ("none", "minimal")
    body["think"] = bool(on)


async def split_stream(lines):
    """SSE lines from FastFlowLM -> the same events with <think> moved into delta.reasoning_content."""
    sp = ThinkSplit()
    async for line in lines:
        if not line.startswith("data:") or line.strip() == "data: [DONE]":
            yield line + "\n"
            continue
        try:
            j = json.loads(line[5:])
        except ValueError:
            yield line + "\n"
            continue
        for c in j.get("choices") or []:
            d = c.setdefault("delta", {})
            final = c.get("finish_reason") is not None    # the last chunk: flush what waited for more
            if isinstance(d.get("content"), str) or final:
                think, text = sp.feed(d.get("content") or "", final=final)
                if text or isinstance(d.get("content"), str):
                    d["content"] = text
                if think:
                    d["reasoning_content"] = (d.get("reasoning_content") or "") + think
        yield "data: " + json.dumps(j, ensure_ascii=False) + "\n"


async def forward(request: Request, path):
    """Send the request on to the engine (the model name it knows in "model"); streams stay streams."""
    global client
    if not READY.is_set():
        return oai_error(503, "the model is still loading on the NPU", "server_error", "model_loading")
    try:
        body = await request.json()
    except ValueError:
        raise HTTPException(400, "the request body is not valid JSON")
    if not isinstance(body, dict):
        raise HTTPException(400, "the request body must be a JSON object")
    if (body.get("n") or 1) != 1:
        raise HTTPException(400, "n > 1 is not supported (one NPU engine)")
    body["model"] = BACKEND["model"]
    flm_request(body)
    if client is None:
        client = httpx.AsyncClient(base_url=BACKEND["url"], timeout=httpx.Timeout(None, connect=10))
    req = client.build_request("POST", path, json=body)
    try:
        r = await client.send(req, stream=True)
    except httpx.HTTPError as e:
        return oai_error(502, f"the NPU engine did not answer: {e}", "server_error")
    ctype = r.headers.get("content-type", "application/json")
    if body.get("stream") and r.status_code == 200:
        async def relay():
            try:
                if path.endswith("chat/completions"):
                    async for line in split_stream(r.aiter_lines()):
                        yield line
                else:
                    async for chunk in r.aiter_raw():
                        yield chunk
            finally:                                   # the client left or the answer ended: stop the engine's work
                await r.aclose()
        return StreamingResponse(relay(), status_code=r.status_code, media_type=ctype.split(";")[0],
                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
    data = await r.aread()
    await r.aclose()
    if r.status_code == 200 and "json" in ctype:       # the answer carries this server's model name
        try:
            j = json.loads(data)
            if isinstance(j, dict) and "model" in j:
                j["model"] = MODEL_ID
                for c in j.get("choices") or []:
                    msg = c.get("message") or {}
                    if isinstance(msg.get("content"), str):
                        think, text = ThinkSplit().feed(msg["content"], final=True)
                        if think:
                            msg["content"], msg["reasoning_content"] = text, think
                data = json.dumps(j, ensure_ascii=False).encode("utf8")
        except ValueError:
            pass
    if r.status_code >= 400 and "json" not in ctype:
        return oai_error(r.status_code, data.decode("utf8", "replace")[:1000] or "NPU engine error", "server_error")
    return Response(data, status_code=r.status_code, media_type=ctype.split(";")[0])


@app.post("/v1/chat/completions")
async def chat(request: Request):
    return await forward(request, "/v1/chat/completions")


@app.post("/v1/completions")
async def completions(request: Request):
    return await forward(request, "/v1/completions")


# ------------------------------------------------------------------ the model and its engine

def resolve(key):
    """(flm checkpoint, labels) of a model name ("llama3.2-1b-FLM", or flm's own "llama3.2:1b"). Downloads what is
    missing: FastFlowLM, then the model."""
    from onw import app as onwapp
    cfg = onwapp.load_config()
    if not lemonade.up_to_date():
        print(f"[setup] installing FastFlowLM {lemonade.VERSION} ...", flush=True)
        last = [0.0]

        def progress(frac, text):
            if time.time() - last[0] > 5:
                last[0] = time.time()
                print(f"[setup] {text}", flush=True)
        lemonade.install(progress)
    ok, msg = lemonade.flm_validate(cfg)
    if not ok:
        sys.exit(f"ERROR: FastFlowLM NPU check failed: {msg}")
    cat = lemonade.flm_models(cfg)
    m = cat.get(key) or cat.get(lemonade.flm_name(key))
    if m is None:
        sys.exit(f"ERROR: FastFlowLM has no model {key!r} (onw list shows the models).")
    if not m["installed"]:
        print(f"[setup] flm pull {m['checkpoint']} ...", flush=True)
        if subprocess.run(lemonade.flm_pull_command(m["checkpoint"]), env=lemonade.flm_env(cfg)).returncode != 0:
            sys.exit("ERROR: flm pull failed")
    bad = lemonade.flm_check_sizes(m["checkpoint"], cfg)
    if bad:                                            # a file missing or cut short: broken weights make garbage
        sys.exit(f"ERROR: the model's files are damaged ({', '.join(bad)}); download it again "
                 "(onw window: Models -> Download again)")
    return m["checkpoint"], m["labels"]


def start_engine(checkpoint, port, ctx):
    """Start FastFlowLM on a free port and mark READY once it answers; if it ends, this server ends too (the window
    then shows the failure and the log)."""
    bport = lemonade.free_port(8001, avoid=(port,))
    argv, ready_path, env = lemonade.serve_command(checkpoint, bport, ctx)
    print(f"[onw] FastFlowLM: {' '.join(argv)}", flush=True)
    kw = {"creationflags": subprocess.CREATE_NO_WINDOW} if os.name == "nt" else {}
    proc = subprocess.Popen(argv, env={**os.environ, **env}, cwd=os.path.dirname(argv[0]),
                            stdout=sys.stdout, stderr=subprocess.STDOUT, **kw)
    lemonade.tie_to_this_process(proc)
    BACKEND.update(url=f"http://127.0.0.1:{bport}", proc=proc)
    atexit.register(lambda: proc.poll() is None and proc.kill())

    def watch():
        t0 = time.time()
        if lemonade.wait_ready(proc, BACKEND["url"] + ready_path):
            print(f"[onw] model loaded on the NPU in {time.time() - t0:.0f} s", flush=True)
            READY.set()
        rc = proc.wait()
        print(f"ERROR: the NPU engine (FastFlowLM) ended (exit code {rc})", flush=True)
        os._exit(1)
    threading.Thread(target=watch, daemon=True).start()


def main(argv=None):
    global MODEL_ID, API_KEY
    ap = argparse.ArgumentParser(prog="onw serve", description="chat UI + OpenAI-compatible API on the AMD NPU")
    ap.add_argument("model", help="model name from `onw list` (llama3.2-1b-FLM, qwen3-4b-FLM, ...)")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8010)
    ap.add_argument("--context", type=int, default=4096, help="context length in tokens (default 4096)")
    ap.add_argument("--open", action="store_true", help="open the chat UI in the default browser once loaded")
    ap.add_argument("--api-key", help="require this key as a Bearer token on /v1 (e.g. with --host 0.0.0.0)")
    a = ap.parse_args(argv)
    API_KEY = a.api_key or os.environ.get("ONW_API_KEY") or None
    n = lemonade.npu()
    print(f"[onw] AMD {__version__} | {lemonade.npu_text()} | {n.get('name', '')}", flush=True)
    if n.get("arch") != "XDNA2":
        sys.exit("ERROR: no AMD XDNA2 NPU found (Ryzen AI 300 series or newer with the NPU driver installed). "
                 f"onw AMD runs models on the NPU only. Driver: {lemonade.DRIVER_URL}")
    from onw.app import _port_free
    if not _port_free(a.host, a.port):                 # before the engine starts and loads the model for nothing
        sys.exit(f"ERROR: port {a.port} is in use by another program (choose another one with --port)")
    checkpoint, labels = resolve(a.model)
    MODEL_ID = lemonade.flm_name(checkpoint)
    BACKEND["model"] = checkpoint                      # flm wants its checkpoint name in "model"
    info.update(model=MODEL_ID, title=MODEL_ID, engine="flm", engine_name=lemonade.ENGINE,
                engine_version=lemonade.VERSION, devices={"lm": "NPU"}, npu=lemonade.npu_text(),
                labels=labels, vision="vision" in labels, think="reasoning" in labels, context=a.context,
                max_context=a.context)
    start_engine(checkpoint, a.port, a.context)
    if a.open:
        import webbrowser
        host = "localhost" if a.host in ("0.0.0.0", "127.0.0.1") else a.host

        def open_when_ready():
            READY.wait()
            webbrowser.open(f"http://{host}:{a.port}/chat")
        threading.Thread(target=open_when_ready, daemon=True).start()
    print(f"\n  OpenAI API:  http://localhost:{a.port}/v1   (model id: {MODEL_ID})"
          f"\n  Chat UI:     http://localhost:{a.port}/chat\n  Check page:  http://localhost:{a.port}/check\n", flush=True)
    uvicorn.run(app, host=a.host, port=a.port, log_level="warning")


if __name__ == "__main__":
    main()