"""
tiny_server.py — wrap tiny_engine in a minimal OpenAI-compatible
/v1/completions endpoint (streaming only) so measure_client.py can be
tested without a GPU. Stdlib only.   python tiny_server.py  (port 8000)
"""
import json
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from tiny_engine import ByteTokenizer, TinyDecoder, generate

TOK, MODEL = ByteTokenizer(), TinyDecoder()


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *a):
        pass

    def do_POST(self):
        req = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.end_headers()

        def send(obj):
            self.wfile.write(f"data: {json.dumps(obj)}\n\n".encode())
            self.wfile.flush()                       # push each token NOW

        n_out = 0
        for event, text, _ in generate(MODEL, TOK, req["prompt"], req["max_tokens"]):
            if event == "token":
                n_out += 1
                send({"id": "c", "object": "text_completion", "model": "tiny",
                      "created": 0, "choices": [{"index": 0, "text": text,
                                                 "finish_reason": None}]})
        send({"id": "c", "object": "text_completion", "model": "tiny", "created": 0,
              "choices": [], "usage": {"prompt_tokens": len(TOK.encode(req["prompt"])),
                                       "completion_tokens": n_out,
                                       "total_tokens": n_out + len(TOK.encode(req["prompt"]))}})
        self.wfile.write(b"data: [DONE]\n\n")


if __name__ == "__main__":
    ThreadingHTTPServer(("127.0.0.1", 8000), Handler).serve_forever()