import json
import os
from http.server import BaseHTTPRequestHandler, HTTPServer
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

TARGET = os.getenv("GROBID_TARGET", "http://127.0.0.1:8070").rstrip("/")
PORT = int(os.getenv("PORT", "7860"))

HOP_BY_HOP = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}


class ProxyHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def _forward(self):
        content_length = int(self.headers.get("Content-Length", "0") or "0")
        body = self.rfile.read(content_length) if content_length > 0 else None
        upstream_url = f"{TARGET}{self.path}"

        headers = {}
        for k, v in self.headers.items():
            lk = k.lower()
            if lk == "host" or lk in HOP_BY_HOP:
                continue
            headers[k] = v

        req = Request(upstream_url, data=body, headers=headers, method=self.command)

        try:
            with urlopen(req, timeout=180) as resp:
                data = resp.read()
                self.send_response(resp.getcode())
                for k, v in resp.getheaders():
                    if k.lower() in HOP_BY_HOP:
                        continue
                    self.send_header(k, v)
                self.send_header("Content-Length", str(len(data)))
                self.end_headers()
                if data:
                    self.wfile.write(data)
        except HTTPError as e:
            data = e.read() if hasattr(e, "read") else b""
            self.send_response(e.code)
            for k, v in e.headers.items():
                if k.lower() in HOP_BY_HOP:
                    continue
                self.send_header(k, v)
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            if data:
                self.wfile.write(data)
        except URLError:
            data = json.dumps({"error": "Upstream GROBID unavailable"}).encode("utf-8")
            self.send_response(502)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(data)))
            self.end_headers()
            self.wfile.write(data)

    def do_GET(self): self._forward()
    def do_POST(self): self._forward()
    def do_PUT(self): self._forward()
    def do_DELETE(self): self._forward()
    def do_PATCH(self): self._forward()
    def do_OPTIONS(self): self._forward()


if __name__ == "__main__":
    server = HTTPServer(("0.0.0.0", PORT), ProxyHandler)
    print(f"Proxy listening on :{PORT}, forwarding to {TARGET}")
    server.serve_forever()