"""zed-fim-proxy — a keyless front door for Zed's edit-prediction feature.

Zed's edit_predictions.open_ai_compatible_api provider cannot send an
Authorization header, so it needs a route where POST /v1/completions succeeds
with no key. This proxy is that route, on a separate port from LiteLLM, with
four narrow guards + an injected model-scoped key:

  1. POST only, path exactly /v1/completions (GET /ping is an anonymous liveness).
  2. source-IP allowlist (best-effort — set ZED_ALLOWED_IPS; empty = allow all).
     NOTE: only meaningful if this process sees the real client IP — run the
     container with network_mode: host (docker port-publish would NAT the source
     to the bridge gateway and defeat it). Behind a site-to-site NAT it may still
     see a mesh IP, not the Mac — verify against the access log at deploy.
  3. request body "model" must equal ZED_ALLOWED_MODEL (default coder-fast).
  4. injects Authorization: Bearer $ZED_SCOPED_KEY (a LiteLLM virtual key scoped
     to that ONE model) and forwards to $ZED_UPSTREAM. The scoped key is the real
     blast-radius bound: even if guards 1-3 were bypassed, the key can reach
     nothing but coder-fast.

Stdlib only (no pip) — runs in a bare python:slim container.
"""
import json
import os
import sys
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

UPSTREAM = os.environ.get("ZED_UPSTREAM", "http://localhost:4000").rstrip("/")
SCOPED_KEY = os.environ["ZED_SCOPED_KEY"]
ALLOWED_MODEL = os.environ.get("ZED_ALLOWED_MODEL", "coder-fast")
ALLOWED_IPS = set(x.strip() for x in os.environ.get("ZED_ALLOWED_IPS", "").split(",") if x.strip())
PORT = int(os.environ.get("ZED_PORT", "4141"))
TIMEOUT = float(os.environ.get("ZED_TIMEOUT", "60"))


class Handler(BaseHTTPRequestHandler):
    server_version = "zed-fim-proxy/1.0"

    def _json(self, code, obj):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        if self.path.rstrip("/") == "/ping":
            return self._json(200, {"service": "ok"})
        return self._json(404, {"error": "not found"})

    def do_POST(self):
        src = self.client_address[0]
        if ALLOWED_IPS and src not in ALLOWED_IPS:
            return self._json(403, {"error": f"source {src} not allowed"})
        # Accept both /v1/completions and /completions so the Zed api_url can be
        # set to either http://host:4141/v1 or http://host:4141 (Zed appends
        # /completions). /v1/chat/completions is NOT in the set → stays rejected.
        if self.path.split("?")[0].rstrip("/") not in ("/v1/completions", "/completions"):
            return self._json(404, {"error": "only POST /v1/completions or /completions"})
        try:
            length = int(self.headers.get("Content-Length", 0))
        except ValueError:
            return self._json(400, {"error": "bad content-length"})
        raw = self.rfile.read(length)
        try:
            body = json.loads(raw)
        except Exception:
            return self._json(400, {"error": "invalid json body"})
        if body.get("model") != ALLOWED_MODEL:
            return self._json(403, {"error": f"model must be '{ALLOWED_MODEL}'"})
        req = urllib.request.Request(
            UPSTREAM + "/v1/completions",
            data=raw,
            method="POST",
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {SCOPED_KEY}",
            },
        )
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
                data, code, ctype = r.read(), r.status, r.headers.get("Content-Type", "application/json")
        except urllib.error.HTTPError as e:
            data, code, ctype = e.read(), e.code, e.headers.get("Content-Type", "application/json")
        except Exception as e:  # noqa: BLE001
            return self._json(502, {"error": f"upstream error: {e}"})
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, fmt, *args):
        sys.stderr.write("%s - %s\n" % (self.client_address[0], fmt % args))


if __name__ == "__main__":
    print(f"zed-fim-proxy on :{PORT} -> {UPSTREAM} (model={ALLOWED_MODEL}, "
          f"ip-allowlist={'set' if ALLOWED_IPS else 'OFF'})", flush=True)
    ThreadingHTTPServer(("0.0.0.0", PORT), Handler).serve_forever()