"""
Lightweight reverse proxy — single entry point for Tailscale Funnel.

Routes:
  /pr-tracker/*          → Streamlit  (localhost:8501)
  /course-dashboard.html → static HTML file
  /                      → redirect to /pr-tracker
"""
import http.server
import urllib.request
import urllib.error
import os
import sys

STREAMLIT_PORT  = 8501
STATIC_HTML     = os.path.expanduser("~/course-dashboard.html")
PROXY_PORT      = 8090


class Router(http.server.BaseHTTPRequestHandler):

    def log_message(self, fmt, *args):
        print(f"[proxy] {self.address_string()} - {fmt % args}")

    # ── helpers ───────────────────────────────────────────────────────────────

    def _send(self, code, content_type, body: bytes):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _redirect(self, location):
        self.send_response(302)
        self.send_header("Location", location)
        self.end_headers()

    def _proxy_to_streamlit(self):
        """Forward request to Streamlit and stream the response back."""
        target = f"http://localhost:{STREAMLIT_PORT}{self.path}"
        headers = {k: v for k, v in self.headers.items()
                   if k.lower() not in ("host", "connection")}
        headers["Host"] = f"localhost:{STREAMLIT_PORT}"

        # Read body for POST/PUT
        body = None
        length = int(self.headers.get("Content-Length", 0))
        if length:
            body = self.rfile.read(length)

        req = urllib.request.Request(target, data=body, headers=headers,
                                     method=self.command)
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                self.send_response(resp.status)
                for k, v in resp.headers.items():
                    if k.lower() in ("connection", "transfer-encoding"):
                        continue
                    self.send_header(k, v)
                self.end_headers()
                self.wfile.write(resp.read())
        except urllib.error.HTTPError as e:
            body = e.read()
            self.send_response(e.code)
            self.end_headers()
            self.wfile.write(body)
        except Exception as exc:
            msg = f"Proxy error: {exc}".encode()
            self._send(502, "text/plain", msg)

    # ── routing ───────────────────────────────────────────────────────────────

    def _handle(self):
        path = self.path.split("?")[0]

        if path == "/" or path == "":
            self._redirect("/pr-tracker")

        elif path == "/course-dashboard.html":
            if os.path.exists(STATIC_HTML):
                with open(STATIC_HTML, "rb") as f:
                    body = f.read()
                self._send(200, "text/html; charset=utf-8", body)
            else:
                self._send(404, "text/plain", b"course-dashboard.html not found")

        elif path.startswith("/pr-tracker") or path.startswith("/_stcore") \
                or path.startswith("/static") or path.startswith("/stream"):
            self._proxy_to_streamlit()

        else:
            # Fall through to Streamlit for everything else
            # (Streamlit assets, websocket upgrade handled by the browser directly)
            self._proxy_to_streamlit()

    def do_GET(self):    self._handle()
    def do_POST(self):   self._handle()
    def do_HEAD(self):   self._handle()
    def do_OPTIONS(self): self._handle()


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else PROXY_PORT
    server = http.server.ThreadingHTTPServer(("127.0.0.1", port), Router)
    print(f"[proxy] Listening on http://127.0.0.1:{port}")
    print(f"[proxy]  /pr-tracker           → Streamlit :{STREAMLIT_PORT}")
    print(f"[proxy]  /course-dashboard.html → {STATIC_HTML}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[proxy] Stopped.")