#!/usr/bin/env python3
"""Local relay between a Black Widow tab and anything else on this machine.

The browser cannot listen on a port, so this does: clients POST a question
here, the open tab polls /next, answers on the device's GPU, and POSTs the
reply back. Python 3 standard library only. Nothing leaves the machine.

    python3 bw_relay.py            # listens on http://localhost:8770
    python3 bw_relay.py 8800       # another port; set the same URL in the tab

Then, in the tab: set > api > "serve this tab". From a shell:

    curl -s localhost:8770/ask -d '{"text": "who designed the Sydney Opera House?"}'
    curl -s localhost:8770/v1/chat/completions \
         -d '{"messages": [{"role": "user", "content": "improve this script"}]}'

The OpenAI-shaped endpoint takes the last user message; history lives in the
tab, not here. GET /health reports whether a tab is polling.
"""
import json
import sys
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8770
LONG_POLL_S = 25
ANSWER_TIMEOUT_S = 180

lock = threading.Condition()
pending = []            # jobs the tab has not picked up yet
answers = {}            # id -> reply dict
last_poll = 0.0         # when a tab last asked for work


class H(BaseHTTPRequestHandler):
    server_version = "bw-relay/1"

    def log_message(self, fmt, *args):
        sys.stderr.write("%s %s\n" % (time.strftime("%H:%M:%S"), fmt % args))

    def _cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "content-type")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")

    def _json(self, code, obj=None):
        body = b"" if obj is None else json.dumps(obj).encode()
        self.send_response(code)
        self._cors()
        if body:
            self.send_header("Content-Type", "application/json")
        if code != 204:
            self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if body:
            self.wfile.write(body)

    def _body(self):
        n = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(n) if n else b""
        try:
            return json.loads(raw or b"{}")
        except json.JSONDecodeError:
            return {"text": raw.decode(errors="replace")}

    def do_OPTIONS(self):
        self._json(204)

    def do_GET(self):
        global last_poll
        if self.path.startswith("/next"):
            # The tab waits here for work; 204 when nothing arrived in time.
            with lock:
                last_poll = time.time()
                deadline = time.time() + LONG_POLL_S
                while not pending and time.time() < deadline:
                    lock.wait(deadline - time.time())
                job = pending.pop(0) if pending else None
            return self._json(200, job) if job else self._json(204)
        if self.path.startswith("/health"):
            return self._json(200, {"tab": time.time() - last_poll < LONG_POLL_S + 5, "pending": len(pending)})
        self._json(404, {"error": "GET /next, GET /health, POST /ask, POST /reply, POST /v1/chat/completions"})

    def do_POST(self):
        body = self._body()
        if self.path.startswith("/reply"):
            with lock:
                answers[body.get("id")] = body
                lock.notify_all()
            return self._json(200, {"ok": True})
        if self.path.startswith("/ask") or self.path.startswith("/v1/chat/completions"):
            text = body.get("text")
            if text is None:
                users = [m for m in body.get("messages", []) if m.get("role") == "user"]
                text = users[-1].get("content", "") if users else ""
            if not isinstance(text, str) or not text.strip():
                return self._json(400, {"error": "no message"})
            if time.time() - last_poll > LONG_POLL_S + 5:
                return self._json(503, {"error": "no Black Widow tab is polling this relay; open the page and enable set > api"})
            job = {"id": uuid.uuid4().hex[:8], "text": text, "t": time.time()}
            with lock:
                pending.append(job)
                lock.notify_all()
                deadline = time.time() + ANSWER_TIMEOUT_S
                while job["id"] not in answers and time.time() < deadline:
                    lock.wait(deadline - time.time())
                out = answers.pop(job["id"], None)
            if out is None:
                return self._json(504, {"error": "the tab did not answer in time"})
            if self.path.startswith("/v1/"):
                return self._json(200, {
                    "id": "bw-" + job["id"], "object": "chat.completion", "created": int(job["t"]),
                    "model": out.get("model", "black-widow"),
                    "choices": [{"index": 0, "finish_reason": "stop",
                                 "message": {"role": "assistant", "content": out.get("reply", out.get("error", ""))}}],
                    "usage": {"prompt_tokens": (out.get("stats") or {}).get("tokIn"), "completion_tokens": (out.get("stats") or {}).get("tokOut")},
                    "bw": {k: out.get(k) for k in ("notes", "hits", "attached", "unsourced", "verify", "stats", "events", "error")},
                })
            return self._json(200, out)
        self._json(404, {"error": "unknown path"})


if __name__ == "__main__":
    srv = ThreadingHTTPServer(("127.0.0.1", PORT), H)
    print(f"bw relay on http://localhost:{PORT}  (tab: set > api > serve this tab, relay http://localhost:{PORT})")
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass