"""qev_image_server: a local System One HTTP server with images (qev_server.py's protocol plus an "images" field).

    python qev_image_server.py --model labqoat/squad-0.8b --port 8018

    curl http://127.0.0.1:8018/v1/systemone -H 'Content-Type: application/json' -d '{
      "state": "Customer screenshot: <image:shot>",
      "images": {"shot": "<base64 of a PNG or JPEG>"},
      "questions": {"err": {"type": "noul", "instructions": "Does the screenshot show an error dialog?"}}}'

The request format is the System One one of qev_server.py plus an optional "images": {name: base64 string} object
(a data: URL prefix is accepted). The state references an image with "<image:name>"; an image the state does not
mention is appended to the context. The answer has the same shape as qev_server.py's, with
"usage": {"input_tokens", "image_tokens", "output_tokens": 0}. Requests are serialised on one model (no
micro-batching here). Request validation and the Server class come from qev_server.py, unchanged.
"""

import argparse
import base64
import binascii
import io
import json
import os
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)  # qev_server.py and qev_image.py sit next to this file in the qev-0.8b repo

from qev_server import MAX_BODY, Server, _validate  # noqa: E402

from qev_image import QEVImage  # noqa: E402


def decode_images(raw) -> dict:
    from PIL import Image

    if raw is None:
        return {}
    if not isinstance(raw, dict):
        raise ValueError('"images" must be an object {name: base64 string}')
    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 make_handler(m: QEVImage, info: dict):
    lock = threading.Lock()

    class H(BaseHTTPRequestHandler):
        protocol_version = "HTTP/1.1"

        def _send(self, code, body):
            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):
            pass

        def do_GET(self):
            if self.path.split("?")[0] == "/health":
                return self._send(200, {"status": "ok", **info})
            return self._send(404, {"detail": "not found"})

        def do_POST(self):
            if self.path.split("?")[0] != "/v1/systemone":
                return self._send(404, {"detail": "not found"})
            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, _ = _validate(body)
                imgs = decode_images(body.get("images"))
            except (ValueError, json.JSONDecodeError) as e:
                return self._send(422, {"detail": str(e)})
            t = time.perf_counter()
            try:
                with lock:
                    (ans, toks, itoks), = m.decide_batch([(state, qs, imgs)], usage=True)
            except (ValueError, KeyError, TypeError, AssertionError) as e:
                return self._send(422, {"detail": str(e) or type(e).__name__})
            except Exception as e:  # noqa: BLE001
                return self._send(500, {"detail": f"{type(e).__name__}: {e}"})
            return self._send(200, {"model": info["model"], "answers": ans,
                                    "usage": {"input_tokens": toks, "image_tokens": itoks, "output_tokens": 0},
                                    "latency_ms": round((time.perf_counter() - t) * 1000, 2)})

    return H


def main():
    ap = argparse.ArgumentParser(description="qev-0.8b System One server with images (local)")
    ap.add_argument("--model", required=True)
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8018)
    ap.add_argument("--max-pixels", type=int, default=448 * 448)
    ap.add_argument("--device", default=None)
    a = ap.parse_args()
    m = QEVImage.from_pretrained(a.model, device=a.device, max_pixels=a.max_pixels)
    info = {"model": "qev-0.8b-image", "max_pixels": a.max_pixels, "device": str(m.device)}
    srv = Server((a.host, a.port), make_handler(m, info))
    print(f"qev_image_server on http://{a.host}:{a.port}", flush=True)
    srv.serve_forever()


if __name__ == "__main__":
    main()