#!/usr/bin/env python3
"""Trusted-LAN observer bridge for agent-browser's localhost-only dashboard/stream."""
from http.client import HTTPConnection
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import re
import socket
import threading

DASHBOARD_LISTEN = ("0.0.0.0", 4849)
DASHBOARD_TARGET = ("127.0.0.1", 4848)
STREAM_LISTEN = ("0.0.0.0", 4851)
STREAM_TARGET = ("127.0.0.1", 4850)

# The bundled dashboard renders this literal and otherwise asks a laptop to
# connect to its own localhost. Route it back through this LAN bridge instead.
LOCAL_STREAM_LITERAL = b"ws://localhost:${e}"
LAN_STREAM_LITERAL = b"ws://${window.location.hostname}:4851"


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

    def _forward(self):
        body_length = int(self.headers.get("Content-Length", "0"))
        body = self.rfile.read(body_length) if body_length else None
        headers = {key: value for key, value in self.headers.items() if key.lower() not in {"host", "connection"}}
        upstream = HTTPConnection(*DASHBOARD_TARGET, timeout=15)
        try:
            upstream.request(self.command, self.path, body=body, headers=headers)
            response = upstream.getresponse()
            payload = response.read()
            content_type = response.getheader("Content-Type", "")
            if "javascript" in content_type or "text/html" in content_type:
                payload = payload.replace(LOCAL_STREAM_LITERAL, LAN_STREAM_LITERAL)
            self.send_response(response.status, response.reason)
            for key, value in response.getheaders():
                if key.lower() not in {"content-length", "connection", "transfer-encoding"}:
                    self.send_header(key, value)
            self.send_header("Content-Length", str(len(payload)))
            self.send_header("Connection", "close")
            self.end_headers()
            self.wfile.write(payload)
        finally:
            upstream.close()

    do_GET = _forward
    do_POST = _forward
    do_PUT = _forward
    do_DELETE = _forward

    def log_message(self, *_):
        pass


def pipe(source, destination):
    try:
        while data := source.recv(65536):
            destination.sendall(data)
    except OSError:
        pass
    finally:
        try:
            destination.shutdown(socket.SHUT_WR)
        except OSError:
            pass


def handle_stream(client):
    try:
        upstream = socket.create_connection(STREAM_TARGET, timeout=5)
        # agent-browser accepts its own localhost dashboard as the WebSocket
        # origin. Preserve that trust boundary when a LAN viewer connects via
        # this narrow proxy, rather than exposing the stream server directly.
        request = bytearray()
        while b"\r\n\r\n" not in request and len(request) < 65536:
            chunk = client.recv(4096)
            if not chunk:
                client.close()
                upstream.close()
                return
            request.extend(chunk)
        rewritten = re.sub(
            rb"(?im)^Origin: [^\r\n]*\r?$",
            b"Origin: http://localhost:4848\r",
            bytes(request),
        )
        upstream.sendall(rewritten)
    except OSError:
        client.close()
        return
    threading.Thread(target=pipe, args=(client, upstream), daemon=True).start()
    pipe(upstream, client)
    client.close()
    upstream.close()


def run_stream_proxy():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
        server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server.bind(STREAM_LISTEN)
        server.listen(50)
        while True:
            client, _ = server.accept()
            threading.Thread(target=handle_stream, args=(client,), daemon=True).start()


threading.Thread(target=run_stream_proxy, daemon=True).start()
ThreadingHTTPServer(DASHBOARD_LISTEN, DashboardProxy).serve_forever()