#!/opt/chat-gateway/.venv/bin/python
"""chat-gateway — a thin authenticated front door to the model API."""
import http.server
import json
import os
import pathlib
import urllib.request

UPSTREAM = os.environ.get("UPSTREAM", "http://127.0.0.1:8110/complete")
REQUIRE_TOKEN = os.environ.get("REQUIRE_TOKEN", "0") == "1"
LOG_HEADERS = os.environ.get("LOG_HEADERS", "0") == "1"
CLIENTS = "/etc/chat-gateway/clients.token"


def upstream_key():
    creds = os.environ.get("CREDENTIALS_DIRECTORY")
    if creds and os.path.exists(os.path.join(creds, "upstream_key")):
        return pathlib.Path(creds, "upstream_key").read_text().strip()
    path = os.environ.get("UPSTREAM_KEY_FILE")
    if path and os.path.exists(path):
        return pathlib.Path(path).read_text().strip()
    return os.environ.get("UPSTREAM_KEY", "")


class Handler(http.server.BaseHTTPRequestHandler):
    def _json(self, status, body):
        payload = json.dumps(body).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def do_POST(self):
        if LOG_HEADERS:
            print(f"request headers: {dict(self.headers)}", flush=True)
        if REQUIRE_TOKEN:
            expected = pathlib.Path(CLIENTS).read_text().splitlines()[0].strip()
            if self.headers.get("Authorization") != f"Bearer {expected}":
                return self._json(401, {"error": "a client token is required"})
        length = int(self.headers.get("Content-Length") or 0)
        prompt = json.loads(self.rfile.read(length) or b"{}").get("prompt", "")
        req = urllib.request.Request(
            UPSTREAM,
            data=json.dumps({"prompt": prompt}).encode(),
            headers={"Content-Type": "application/json", "X-Api-Key": upstream_key()},
        )
        try:
            with urllib.request.urlopen(req, timeout=20) as r:
                answer = json.loads(r.read())
        except OSError as e:
            return self._json(502, {"error": f"upstream unavailable: {e}"})
        return self._json(200, {"answer": answer.get("answer", "")})

    def log_message(self, fmt, *args):
        print(fmt % args, flush=True)


http.server.ThreadingHTTPServer(("127.0.0.1", 8100), Handler).serve_forever()
