"""qev_server: an HTTP server for qev-0.8b that speaks the System One wire format (POST /v1/systemone).

    pip install -r requirements.txt                 # + requirements-cuda.txt on an NVIDIA GPU
    python qev_server.py                            # http://127.0.0.1:8017, labqoat/squad-0.8b
    python qev_server.py --host 0.0.0.0 --port 8000 --encoding per_question
    python qev_server.py --specialists              # router + games / e2e-tester LoRA slots (fast CUDA backend)

    curl http://127.0.0.1:8017/v1/systemone -H 'Content-Type: application/json' -d '{
      "model": "squad-0.8b",
      "state": "My card was charged twice for the same order, please refund one.",
      "questions": {"billing": {"type": "noul", "instructions": "Is this a billing problem?"}}}'

    -> {"model": "squad-0.8b", "answers": {"billing": {"type": "noul", "noul": 0.97}},
        "usage": {"input_tokens": 51, "output_tokens": 0}}

Endpoints:
- POST /v1/systemone  {"state": str | object | array, "questions": {id: question}, "model": optional, "encoding":
                       optional "joint" | "per_question", "images": optional {name: base64 PNG / JPEG},
                       "image_detail": optional "low" | "default" | "high" (256², 448², 1024² pixels; or
                       "max_pixels": int)} ->
                       {"model", "answers", "usage"}. Code written for the System One API works unchanged; "model" in
                       the request is accepted and ignored. Images (fast CUDA backend; pillow + torchvision): the
                       state places one with "<image:name>", an image it does not mention goes at the end of the
                       context; usage then also reports "image_tokens".
                       With --specialists, usage also reports "route": the slot the router chose ("trunk",
                       "games", "e2e-tester"; per_question: one per question when they differ).
- GET  /health        readiness, backend, encoding, model files' SHA-256 (from SHA256SUMS when present).
- GET  /v1/models     the served model.

State cache and memo (fast backend, on by default): an agent asking one question at a time about the same page
sends the same state again and again. Model states (Gated DeltaNet states and attention K / V at the end of each
request's state, a few tokens before it and every 2048 tokens) are kept on the GPU, LRU within --prefix-cache-mb
(default 2048; 0: off), and a request (calls of up to 4) continues from the deepest one its tokens share: the same
state computes only its questions (L40S, e2e pages: ~7.5 ms instead of ~20 ms), a state that grew computes only what
is new. A new state costs more than without the cache (~+40%, it builds the checkpoints) and is computed the plain
way when the call has several rows or other calls wait. Answers equal the uncached ones up to bf16 round-off (the same
argmax and route; on bf16-sensitive requests the probabilities move as much as between the engine's own CUDA-graph
and eager paths). A request repeated whole (state, questions, encoding) is answered from a memo of the last --memo
requests (default 1024; 0: off) without the model.
Concurrent requests are micro-batched: the first request in an idle server runs at once; requests that arrive while
a batch is running are answered together in the next one (up to --max-batch). Answers are the same as one at a time
up to floating-point round-off. With QEV_API_KEY set, every call but /health needs "Authorization: Bearer <key>".
Only the Python standard library is used on top of qev.py.
"""

import argparse
import hmac
import json
import os
import queue
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

MAX_BODY = 32 * 1024 * 1024
ENCODINGS = ("joint", "per_question")


class Batcher:
    """One worker thread owns the model; requests queue up and are answered in batches."""

    def __init__(self, decide_batch, max_batch: int = 8):
        self.decide_batch, self.max_batch = decide_batch, max_batch
        self.q: queue.Queue = queue.Queue()
        threading.Thread(target=self._loop, daemon=True).start()

    def submit(self, state, questions: dict, per_question: bool, timeout: float = 300.0, images: dict | None = None,
               max_pixels: int | None = None):
        done, box = threading.Event(), {}
        self.q.put((state, questions, per_question, done, box, images, max_pixels))
        if not done.wait(timeout):
            raise TimeoutError("the model did not answer in time")
        if "error" in box:
            raise box["error"]
        return box["answers"], box["tokens"], box.get("image_tokens", 0), box.get("route")

    def _loop(self):
        while True:
            jobs = [self.q.get()]
            while len(jobs) < self.max_batch:
                try:
                    jobs.append(self.q.get_nowait())
                except queue.Empty:
                    break
            for key in dict.fromkeys((j[2], j[6]) for j in jobs):  # one batch per (encoding, image budget)
                part = [j for j in jobs if (j[2], j[6]) == key]
                self._run(part, key[0])

    def _run(self, jobs: list, per_question: bool):
        try:
            res = self.decide_batch([(j[0], j[1], j[5]) if j[5] else (j[0], j[1]) for j in jobs], per_question,
                                    jobs[0][6])
        except Exception as e:  # noqa: BLE001 - one bad request must not fail the others: answer each on its own
            if len(jobs) > 1:
                for j in jobs:
                    self._run([j], per_question)
            else:
                jobs[0][4]["error"] = e
                jobs[0][3].set()
            return
        for j, r in zip(jobs, res, strict=True):
            j[4]["answers"], j[4]["tokens"] = r[0], r[1]
            j[4]["image_tokens"] = r[2] if len(r) > 2 else 0
            j[4]["route"] = r[3] if len(r) > 3 else None
            j[3].set()


def _validate(req) -> tuple:
    if not isinstance(req, dict):
        raise ValueError('the body must be a JSON object: {"state": ..., "questions": {id: question}}')
    if "state" not in req:
        raise ValueError('missing "state"')
    qs = req.get("questions")
    if not isinstance(qs, dict) or not qs:
        raise ValueError('"questions" must be a non-empty object {id: question}')
    for k, q in qs.items():
        if not isinstance(q, dict) or q.get("type") not in ("noul", "choice", "score"):
            raise ValueError(f'question "{k}": "type" must be "noul", "choice" or "score"')
        if not isinstance(q.get("instructions", ""), str):
            raise ValueError(f'question "{k}": "instructions" must be a string')
        if q["type"] in ("choice", "score") and not q.get("criteria"):
            raise ValueError(f'question "{k}": a {q["type"]} question needs "criteria"')
    enc = req.get("encoding")
    if enc is not None and enc not in ENCODINGS:
        raise ValueError('"encoding" must be "joint" or "per_question"')
    return req["state"], qs, enc


IMAGE_DETAIL = {"low": 256 * 256, "default": 448 * 448, "high": 1024 * 1024}


def image_budget(req: dict) -> int | None:
    """The request's image budget: "max_pixels" (int, 1024 to 4194304) or "image_detail" ("low" 256x256, "default"
    448x448, "high" 1024x1024); None: the server's --max-pixels."""
    mp, det = req.get("max_pixels"), req.get("image_detail")
    if mp is not None:
        if isinstance(mp, bool) or not isinstance(mp, int) or not 1024 <= mp <= 2048 * 2048:
            raise ValueError('"max_pixels" must be an integer from 1024 to 4194304')
        return mp
    if det is not None:
        if det not in IMAGE_DETAIL:
            raise ValueError('"image_detail" must be "low", "default" or "high"')
        return IMAGE_DETAIL[det]
    return None


def decode_images(raw) -> dict | None:
    """The "images" field: {name: base64 string (a data: URL prefix is accepted)} -> {name: PIL image}."""
    import base64
    import binascii
    import io

    if raw is None or raw == {}:
        return None
    if not isinstance(raw, dict):
        raise ValueError('"images" must be an object {name: base64 string}')
    try:
        from PIL import Image
    except ImportError:
        raise ValueError("images need pillow (pip install pillow torchvision)") from None
    out = {}
    for k, v in raw.items():
        if not isinstance(v, str):
            raise ValueError(f'image "{k}": a base64 string')
        if v.startswith("data:"):
            v = v.split(",", 1)[-1]
        try:
            out[k] = Image.open(io.BytesIO(base64.b64decode(v, validate=True))).convert("RGB")
        except (binascii.Error, OSError) as e:
            raise ValueError(f'image "{k}": not a decodable base64 image ({e})') from None
    return out


def _file_hashes(path: str | None) -> dict:
    if path and os.path.exists(os.path.join(path, "SHA256SUMS")):
        out = {}
        for line in open(os.path.join(path, "SHA256SUMS")):
            parts = line.split()
            if len(parts) == 2 and parts[1].endswith((".safetensors", "tokenizer.json")):
                out[parts[1]] = parts[0]
        return out
    return {}


class Server(ThreadingHTTPServer):
    daemon_threads = True
    request_queue_size = 1024  # the listen backlog (the stdlib default of 5 resets bursts of concurrent clients)


def make_handler(batcher: Batcher, info: dict, encoding: str, api_key: str | None):
    class H(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def _send(self, code: int, body: dict):
            data = json.dumps(body, 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, *args):
            if info.get("verbose"):
                super().log_message(fmt, *args)

        def _authorized(self) -> bool:
            if not api_key:
                return True
            got = self.headers.get("Authorization", "")
            return hmac.compare_digest(got, f"Bearer {api_key}")

        def do_GET(self):
            path = self.path.split("?", 1)[0]
            if path == "/health":
                return self._send(200, {"status": "ok", **{k: v for k, v in info.items() if k != "verbose"}})
            if not self._authorized():
                return self._send(401, {"detail": "missing or wrong API key"})
            if path == "/v1/models":
                return self._send(200, {"object": "list", "data": [{"id": info["model"], "object": "model",
                                                                    "owned_by": "labqoat"}]})
            return self._send(404, {"detail": "not found"})

        def do_POST(self):
            path = self.path.split("?", 1)[0]
            if path != "/v1/systemone":
                return self._send(404, {"detail": "not found"})
            if not self._authorized():
                return self._send(401, {"detail": "missing or wrong API key"})
            try:
                n = int(self.headers.get("Content-Length", 0))
                if n > MAX_BODY:
                    return self._send(413, {"detail": f"the body is larger than {MAX_BODY} bytes"})
                body = json.loads(self.rfile.read(n) or b"{}")
                state, qs, enc = _validate(body)
                imgs = decode_images(body.get("images"))
                budget = image_budget(body)
                if imgs and not info.get("images"):
                    raise ValueError("this server does not read images (the fast CUDA backend is needed)")
            except (ValueError, json.JSONDecodeError) as e:
                return self._send(422, {"detail": str(e)})
            t = time.perf_counter()
            try:
                ans, toks, itoks, route = batcher.submit(state, qs, (enc or encoding) == "per_question", images=imgs,
                                                         max_pixels=budget if imgs else None)
            except (ValueError, KeyError, TypeError, AssertionError) as e:
                return self._send(422, {"detail": str(e) or type(e).__name__})
            except TimeoutError as e:
                return self._send(503, {"detail": str(e)})
            except Exception as e:  # noqa: BLE001
                return self._send(500, {"detail": f"{type(e).__name__}: {e}"})
            usage = {"input_tokens": toks, "output_tokens": 0}
            if imgs:
                usage["image_tokens"] = itoks
            if route is not None:
                usage["route"] = route
            return self._send(200, {"model": info["model"], "answers": ans, "usage": usage,
                                    "latency_ms": round((time.perf_counter() - t) * 1000, 2)})

    return H


def main():
    ap = argparse.ArgumentParser(description="qev-0.8b System One HTTP server (POST /v1/systemone)")
    ap.add_argument("--model", default="labqoat/squad-0.8b", help="local directory or Hugging Face repo id")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8017)
    ap.add_argument("--backend", default="auto", choices=["auto", "fast", "transformers"])
    ap.add_argument("--device", default=None)
    ap.add_argument("--encoding", default="joint", choices=ENCODINGS,
                    help="default for requests without \"encoding\": joint (every question behind one copy of the "
                         "state) or per_question (one row per question; often more coherent answers)")
    ap.add_argument("--max-batch", type=int, default=8,
                    help="most requests answered in one batch (8 measured best on an L40S: ~120 req/s)")
    ap.add_argument("--fast-start", action="store_true",
                    help="fast backend: listen once the common CUDA graphs are captured and capture the rest in the "
                         "background (requests during that time can be slow); default: capture all before listening")
    ap.add_argument("--no-images", action="store_true",
                    help="do not load the vision tower (requests with \"images\" then get HTTP 422)")
    ap.add_argument("--max-pixels", type=int, default=448 * 448, help="images are resized to at most this many pixels")
    ap.add_argument("--prefix-cache-mb", type=int, default=2048,
                    help="fast backend: the cross-request state cache, MB of GPU memory (LRU): a request whose state "
                         "shares tokens with an earlier one continues from its cached model state; 0: off")
    ap.add_argument("--memo", type=int, default=1024,
                    help="fast backend: whole requests remembered (the same request again: the same answer, no model "
                         "call); 0: off")
    ap.add_argument("--specialists", action="store_true",
                    help="router + LoRA slots from the model's specialists/ directory (fast CUDA backend)")
    ap.add_argument("--verbose", action="store_true", help="log every request")
    a = ap.parse_args()

    from qev import QEV

    t0 = time.perf_counter()
    m = QEV.from_pretrained(a.model, device=a.device, backend=a.backend, precapture=False, max_pixels=a.max_pixels,
                            specialists=a.specialists, prefix_cache_mb=a.prefix_cache_mb, memo=a.memo)
    images = False
    if m.fast is not None and not a.no_images:
        try:
            m.fast.enable_images(a.max_pixels)
            images = True
        except Exception as e:  # noqa: BLE001 - no pillow / torchvision / processor config: a text-only server
            print(f"qev_server: images off ({type(e).__name__}: {e})", flush=True)
    if m.fast is not None:  # CUDA graphs for every batch shape, so no live request waits for a capture
        print("qev_server: capturing CUDA graphs (about 2 min the first time, faster with a warm kernel cache)",
              flush=True)
        m.fast.precapture(background=a.fast_start, per_question=a.encoding == "per_question", images=images)
    m.decide("warm-up", {"q": {"type": "noul", "instructions": "Is this a warm-up?"}})
    local = a.model if os.path.isdir(a.model) else None
    if local is None:
        try:
            from huggingface_hub import snapshot_download

            local = snapshot_download(a.model, allow_patterns=["SHA256SUMS"])
        except Exception:  # noqa: BLE001
            local = None
    info = {"model": "squad-0.8b", "source": a.model, "backend": m.backend, "device": str(m.device),
            "encoding": a.encoding, "images": images, "specialists": m.fast.slots.b.describe() if (
                m.fast is not None and m.fast.slots is not None) else None, "load_s": round(time.perf_counter() - t0, 1), "sha256": _file_hashes(local),
            "verbose": a.verbose}
    def run(reqs, pq, max_pixels=None):
        res = m.decide_batch(reqs, pq, usage=True, max_pixels=max_pixels)
        routes = m.last_routes if len(m.last_routes) == len(res) else [None] * len(res)
        return [(a_, t_, k_, r_) for (a_, t_), k_, r_ in zip(res, m.last_image_tokens, routes, strict=True)]

    batcher = Batcher(run, max_batch=a.max_batch)
    key = os.environ.get("QEV_API_KEY") or None
    httpd = Server((a.host, a.port), make_handler(batcher, info, a.encoding, key))
    print(f"qev_server: {info['backend']} backend on {info['device']}, ready in {info['load_s']} s; "
          f"POST http://{a.host}:{a.port}/v1/systemone" + (" (API key required)" if key else ""), flush=True)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()