"""A small local HTTP server for Dewpoint, using only the standard library.

    python serve.py                       # ensemble, port 8000
    python serve.py --single --port 9000  # the 307M member only
    python serve.py --backend onnx        # no torch needed

    curl -s localhost:8000/restore -H "content-type: application/json" \\
         -d '{"text": "what time is it in tokyo right now", "lang": "en"}'
    -> {"text": "What time is it in Tokyo right now?"}

POST /restore   {"text": str | [str], "lang": "en", "labels": false}
GET  /health    {"ok": true, "backend": ..., "members": [...]}

It serves one request at a time, which is right for a single model on one
machine; put it behind a process manager for anything larger.
"""
import argparse
import json
import os
import sys
from http.server import BaseHTTPRequestHandler, HTTPServer

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dewpoint import Punctuator, split_words  # noqa: E402


def make_handler(p):
    class H(BaseHTTPRequestHandler):
        def _send(self, code, obj):
            body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
            self.send_response(code)
            self.send_header("content-type", "application/json; charset=utf-8")
            self.send_header("content-length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def do_GET(self):
            if self.path.rstrip("/") == "/health":
                self._send(200, {"ok": True, "backend": p.backend,
                                 "members": len(p.members)})
            else:
                self._send(404, {"error": "try POST /restore or GET /health"})

        def do_POST(self):
            if self.path.rstrip("/") != "/restore":
                return self._send(404, {"error": "POST /restore"})
            try:
                n = int(self.headers.get("content-length", 0))
                req = json.loads(self.rfile.read(n) or b"{}")
                text, lang = req.get("text", ""), req.get("lang", "en")
                texts = [text] if isinstance(text, str) else list(text)
                out = []
                for t in texts:
                    item = {"text": p.restore(t, lang)}
                    if req.get("labels"):
                        r = p.predict(split_words(t, lang), lang)
                        item.update(words=r["words"], punct=r["punct"], case=r["case"])
                    out.append(item)
                self._send(200, out[0] if isinstance(text, str) else out)
            except Exception as e:  # report, don't crash the server
                self._send(400, {"error": f"{type(e).__name__}: {e}"})

        def log_message(self, fmt, *args):
            pass

    return H


def main():
    ap = argparse.ArgumentParser(description="Serve Dewpoint over HTTP.")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--model", default=os.path.dirname(os.path.abspath(__file__)))
    ap.add_argument("--single", action="store_true")
    ap.add_argument("--backend", default="auto", choices=["auto", "torch", "onnx"])
    a = ap.parse_args()
    p = Punctuator.from_pretrained(a.model, backend=a.backend,
                                   members=["mmbert-base"] if a.single else None)
    print(f"Dewpoint ({p.backend}, {len(p.members)} member(s)) on http://{a.host}:{a.port}")
    HTTPServer((a.host, a.port), make_handler(p)).serve_forever()


if __name__ == "__main__":
    main()