#!/usr/bin/env python3
# llama-queue-proxy: 排队 + 超时快速失败 + 健康透传（方案 A）
# 上游: llama-server 127.0.0.1:8080（双槽）；本代理监听 :8081
# 行为: 上游并行槽位已满时排队，排队超过 QUEUE_TIMEOUT 秒 → 503（调用方退避重试）
import os, threading, urllib.request, urllib.error, http.server

UPSTREAM = os.environ.get("UPSTREAM_URL", "http://127.0.0.1:8080")
PORT = int(os.environ.get("PROXY_PORT", "8081"))
SLOTS = int(os.environ.get("SLOTS", "2"))
QUEUE_TIMEOUT = float(os.environ.get("QUEUE_TIMEOUT", "15"))
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "180"))

sem = threading.Semaphore(SLOTS)

class H(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    def log_message(self, *a):
        pass

    def _send(self, code, body, ct="application/json"):
        b = body.encode() if isinstance(body, str) else body
        self.send_response(code)
        self.send_header("Content-Type", ct)
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        try:
            self.wfile.write(b)
        except Exception:
            pass

    def _forward(self, body, headers):
        req = urllib.request.Request(UPSTREAM + self.path, data=body,
                                     method=self.command, headers=headers)
        try:
            up = urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT)
        except urllib.error.HTTPError as e:
            up = e
        except Exception as e:
            self._send(502, '{"error":{"message":"upstream unreachable: %s"}}' % str(e)[:120])
            return
        data = up.read()
        ct = up.headers.get("Content-Type", "application/json")
        self.send_response(up.getcode())
        self.send_header("Content-Type", ct)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        try:
            self.wfile.write(data)
        except Exception:
            pass

    def do_POST(self):
        ln = int(self.headers.get("Content-Length", 0) or 0)
        body = self.rfile.read(ln) if ln else b""
        if not sem.acquire(timeout=QUEUE_TIMEOUT):
            self._send(503, '{"error":{"message":"llama busy: queue timeout after %ds, retry later"}}' % int(QUEUE_TIMEOUT))
            return
        try:
            hdrs = {k: v for k, v in self.headers.items()
                    if k.lower() not in ("content-length", "host", "connection", "accept-encoding")}
            hdrs["Content-Length"] = str(len(body))
            self._forward(body, hdrs)
        finally:
            sem.release()

    def do_GET(self):
        self._forward(None, {})

if __name__ == "__main__":
    srv = http.server.ThreadingHTTPServer(("127.0.0.1", PORT), H)
    print("[llama-queue-proxy] listen :%d -> %s slots=%d queue_timeout=%ds request_timeout=%ds"
          % (PORT, UPSTREAM, SLOTS, QUEUE_TIMEOUT, REQUEST_TIMEOUT), flush=True)
    srv.serve_forever()