"""A small server-side proxy that keeps an AI API key out of the browser. Python standard library only.

The browser sends a chat request to this server. The server checks it, adds the secret key, and
forwards it to the AI provider. The key is read from an environment variable and appears in one place
only: the header of the outgoing request. It is not in any response, error message or log line, and the
server's own access log (see PathFreeHandler) leaves out the request line, so a path that is the key is not logged.

Run it:  AI_API_KEY=... AI_UPSTREAM_URL=... AI_MODELS=demo-small python3 keyproxy.py
"""
import json
import logging
import math
import os
import sys
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from http import HTTPStatus
from wsgiref.simple_server import WSGIRequestHandler, make_server

log = logging.getLogger("keyproxy")
ROUTE = "/api/chat"
MAX_REPLY_BYTES = 1_000_000


class TokenBucket:
    """Holds up to `capacity` tickets and gains `rate` tickets per second. A request spends one.

    (The word "token" here has nothing to do with the tokens an AI model reads.)
    """

    def __init__(self, capacity, rate, now):
        self.capacity, self.rate = capacity, rate
        self.level, self.last = float(capacity), now

    def take(self, now):
        """Spend one ticket. Return 0.0 if allowed, else the seconds until a ticket is available."""
        self.level = min(self.capacity, self.level + (now - self.last) * self.rate)
        self.last = now
        if self.level >= 1:
            self.level -= 1
            return 0.0
        return (1 - self.level) / self.rate


@dataclass
class Config:
    api_key: str = field(repr=False)  # repr=False: printing the config must not print the key
    upstream_url: str
    allowed_models: frozenset
    max_body_bytes: int = 4096
    max_output_tokens: int = 300
    bucket_capacity: int = 5
    refill_per_second: float = 1.0
    timeout_seconds: float = 10.0


def config_from_env(env=os.environ):
    try:
        models = frozenset(name.strip() for name in env["AI_MODELS"].split(","))
        cfg = Config(env["AI_API_KEY"], env["AI_UPSTREAM_URL"], models)
    except KeyError:
        raise SystemExit("set AI_API_KEY, AI_UPSTREAM_URL and AI_MODELS") from None
    url = urllib.parse.urlsplit(cfg.upstream_url)
    if url.scheme != "https" and url.hostname not in ("127.0.0.1", "localhost"):  # a typo must not send the key in clear text
        raise SystemExit("AI_UPSTREAM_URL must start with https://")
    return cfg


class NoRedirect(urllib.request.HTTPRedirectHandler):
    """Python's urllib would resend the Authorization header to wherever a redirect points."""

    def redirect_request(self, *args, **kwargs):
        return None


def make_app(cfg, clock=time.monotonic):
    """Return a WSGI app. `clock` is a parameter so tests can move time without waiting."""
    buckets = {}  # one TokenBucket per client address
    opener = urllib.request.build_opener(NoRedirect)

    def app(environ, start_response):
        client = environ.get("REMOTE_ADDR", "unknown")  # never X-Forwarded-For: a client can write that

        def reply(status, body, extra=()):
            log.info("client=%s status=%d", client, status)  # no path, no headers, no body
            start_response(f"{status} {HTTPStatus(status).phrase}", [
                ("Content-Type", "application/json"), ("Content-Length", str(len(body))), *extra])
            return [body]

        def fail(status, message, extra=()):
            return reply(status, json.dumps({"error": message}).encode(), extra)  # fixed text, never input

        if environ.get("PATH_INFO") != ROUTE:
            return fail(404, "not found")
        if environ.get("REQUEST_METHOD") != "POST":
            return fail(405, "use POST", [("Allow", "POST")])

        if client not in buckets:
            buckets[client] = TokenBucket(cfg.bucket_capacity, cfg.refill_per_second, clock())
        wait = buckets[client].take(clock())
        if wait:
            return fail(429, "too many requests", [("Retry-After", str(math.ceil(wait)))])

        length = environ.get("CONTENT_LENGTH") or ""
        if not (length.isascii() and length.isdigit()):
            return fail(411, "length required")
        length = length.lstrip("0") or "0"  # leading zeros do not change the number
        if len(length) > 9 or int(length) > cfg.max_body_bytes:  # refuse before reading a byte of the body
            return fail(413, "request too large")
        raw = environ["wsgi.input"].read(int(length))

        try:
            data = json.loads(raw)
            model, messages = data["model"], data["messages"]
        except (ValueError, KeyError, TypeError, RecursionError):
            return fail(400, "send JSON with model and messages")
        if not isinstance(model, str) or not isinstance(messages, list) or not messages:
            return fail(400, "send JSON with model and messages")
        if model not in cfg.allowed_models:
            return fail(400, "model not allowed")

        # The server builds the upstream request: the client cannot add options such as a longer reply.
        upstream_body = json.dumps(
            {"model": model, "messages": messages, "max_tokens": cfg.max_output_tokens}).encode()
        request = urllib.request.Request(cfg.upstream_url, upstream_body, method="POST", headers={
            "Authorization": "Bearer " + cfg.api_key, "Content-Type": "application/json"})
        try:
            with opener.open(request, timeout=cfg.timeout_seconds) as upstream:
                answer = upstream.read(MAX_REPLY_BYTES)
        except Exception as exc:  # an error body can repeat what we sent, so log only a code or a class name
            log.warning("client=%s upstream=%s", client, getattr(exc, "code", type(exc).__name__))
            return fail(502, "upstream failed")
        return reply(200, answer.replace(cfg.api_key.encode(), b"[removed]"))

    return app


class PathFreeHandler(WSGIRequestHandler):
    """The standard handler prints each request line, path included, to stderr. This one prints the address and status."""

    def log_request(self, code="-", size="-"):
        sys.stderr.write(f"client={self.client_address[0]} status={code}\n")

    def log_message(self, format, *args):  # protocol errors quote the request line, so print none of it
        sys.stderr.write(f"client={self.client_address[0]} protocol error\n")


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
    with make_server("127.0.0.1", 8080, make_app(config_from_env()), handler_class=PathFreeHandler) as server:
        server.serve_forever()


if __name__ == "__main__":
    main()