"""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/qev-0.8b
    python qev_server.py --host 0.0.0.0 --port 8000 --encoding per_question

    curl http://127.0.0.1:8017/v1/systemone -H 'Content-Type: application/json' -d '{
      "model": "qev-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": "qev-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}} ->
                       {"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".
- GET  /health        readiness, backend, encoding, model files' SHA-256 (from SHA256SUMS when present).
- GET  /v1/models     the served 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):
        done, box = threading.Event(), {}
        self.q.put((state, questions, per_question, done, box, images))
        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)

    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 pq in (False, True):
                part = [j for j in jobs if j[2] == pq]
                if part:
                    self._run(part, pq)

    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)
        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[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


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"))
                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 = batcher.submit(state, qs, (enc or encoding) == "per_question", images=imgs)
            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
            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/qev-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("--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)
    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": "qev-0.8b", "source": a.model, "backend": m.backend, "device": str(m.device),
            "encoding": a.encoding, "images": images, "load_s": round(time.perf_counter() - t0, 1), "sha256": _file_hashes(local),
            "verbose": a.verbose}
    def run(reqs, pq):
        res = m.decide_batch(reqs, pq, usage=True)
        return [(a_, t_, k_) for (a_, t_), k_ in zip(res, m.last_image_tokens, 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()