#!/usr/bin/env python3
import os, sys, socket, ssl, urllib.request, urllib.parse, struct, base64, select, threading, time

WORKER_HOST = os.environ.get("WORKER_HOST", "egress-proxy.arjav-3003jain.workers.dev")
PROXY_PORT = int(os.environ.get("LOCAL_PROXY_PORT", "8087"))
BUF_SIZE = 65536

def log(msg):
    print(f"[proxy] {time.strftime('%H:%M:%S')} {msg}", flush=True)

def http_proxy_req(data):
    try:
        lines = data.split(b"\r\n")
        parts = lines[0].decode("utf-8", errors="replace").split(" ")
        if len(parts) < 3:
            return None
        method, target_url = parts[0], parts[1]
        worker_url = f"https://{WORKER_HOST}/http?url={urllib.parse.quote(target_url, safe='')}"
        req = urllib.request.Request(worker_url, method=method)
        for line in lines[1:]:
            if b":" in line:
                k, v = line.split(b":", 1)
                k = k.strip().decode("utf-8", errors="replace").lower()
                if k in ("host", "proxy-connection", "proxy-authorization", "transfer-encoding"):
                    continue
                req.add_header(k, v.strip().decode("utf-8", errors="replace"))
        if method in ("POST", "PUT", "PATCH"):
            req.data = data.split(b"\r\n\r\n", 1)[1] if b"\r\n\r\n" in data else b""
        resp = urllib.request.urlopen(req, timeout=30)
        hdrs = f"HTTP/1.1 {resp.status} OK\r\n".encode()
        for k, v in resp.headers.items():
            if k.lower() not in ("transfer-encoding", "content-encoding"):
                hdrs += f"{k}: {v}\r\n".encode()
        return hdrs + b"\r\n" + resp.read()
    except urllib.error.HTTPError as e:
        return f"HTTP/1.1 {e.code} Error\r\n\r\n{e.read().decode(errors='replace')}".encode()
    except Exception as e:
        return f"HTTP/1.1 502 Bad Gateway\r\n\r\nProxy error: {e}".encode()

def recv_exact(sock, n):
    data = b""
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            return None
        data += chunk
    return data

def ws_send_frame(sock, data):
    mask_key = os.urandom(4)
    masked = bytes(b ^ mask_key[i % 4] for i, b in enumerate(data))
    length = len(data)
    hdr = bytearray([0x82])
    if length < 126:
        hdr.append(0x80 | length)
    elif length < 65536:
        hdr += bytearray([0x80 | 126]) + bytearray(struct.pack(">H", length))
    else:
        hdr += bytearray([0x80 | 127]) + bytearray(struct.pack(">Q", length))
    sock.sendall(bytes(hdr) + mask_key + masked)

def ws_recv_frame(sock):
    while True:
        b1 = sock.recv(1)
        if not b1:
            return None
        op = b1[0] & 0x0F
        if op == 0x08:
            return None
        if op in (0x01, 0x02):
            break
        if op == 0x09:
            ws_send_frame(sock, b"")
            continue
        if op == 0x0A:
            continue
    b2 = sock.recv(1)[0]
    length = b2 & 0x7F
    if length == 126:
        raw = recv_exact(sock, 2)
        if raw is None: return None
        length = struct.unpack(">H", raw)[0]
    elif length == 127:
        raw = recv_exact(sock, 8)
        if raw is None: return None
        length = struct.unpack(">Q", raw)[0]
    return recv_exact(sock, length)

def ws_connect_raw(host, port):
    ctx = ssl.create_default_context()
    sock = socket.create_connection((WORKER_HOST, 443), timeout=15)
    ssock = ctx.wrap_socket(sock, server_hostname=WORKER_HOST)
    key = base64.b64encode(os.urandom(16)).decode()
    req = (
        f"GET /ws-tunnel?host={host}&port={port} HTTP/1.1\r\n"
        f"Host: {WORKER_HOST}\r\n"
        f"Upgrade: websocket\r\n"
        f"Connection: Upgrade\r\n"
        f"Sec-WebSocket-Key: {key}\r\n"
        f"Sec-WebSocket-Version: 13\r\n"
        f"\r\n"
    )
    ssock.sendall(req.encode())
    resp = b""
    while b"\r\n\r\n" not in resp:
        chunk = ssock.recv(BUF_SIZE)
        if not chunk:
            raise ConnectionError("Worker hung up")
        resp += chunk
    if b"101" not in resp:
        raise ConnectionError(f"Upgrade failed: {resp.split(b'\r\n')[0].decode()}")
    return key, ssock

def handle_connect(client_sock, addr, host, port):
    try:
        key, ws_sock = ws_connect_raw(host, port)
    except Exception as e:
        log(f"CONNECT {host}:{port} tunnel error: {e}")
        try:
            client_sock.sendall(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
        except:
            pass
        return

    try:
        client_sock.sendall(b"HTTP/1.1 200 Connection Established\r\n\r\n")
    except:
        try:
            ws_sock.close()
        except:
            pass
        return

    stop = threading.Event()

    def client_to_ws():
        try:
            while not stop.is_set():
                ready, _, _ = select.select([client_sock], [], [], 1)
                if not ready:
                    continue
                data = client_sock.recv(BUF_SIZE)
                if not data:
                    break
                ws_send_frame(ws_sock, data)
        except:
            pass
        finally:
            stop.set()

    def ws_to_client():
        try:
            while not stop.is_set():
                data = ws_recv_frame(ws_sock)
                if data is None:
                    break
                client_sock.sendall(data)
        except:
            pass
        finally:
            stop.set()

    t1 = threading.Thread(target=client_to_ws, daemon=True)
    t2 = threading.Thread(target=ws_to_client, daemon=True)
    t1.start()
    t2.start()
    t1.join()
    t2.join()

    try:
        ws_sock.close()
    except:
        pass

def handle_client(client_sock, addr):
    try:
        data = client_sock.recv(BUF_SIZE)
        if not data:
            return
        parts = data.split(b"\r\n")[0].decode("utf-8", errors="replace").split(" ")
        if parts[0] == "CONNECT":
            target = parts[1]
            host, port = (target.split(":") + ["443"])[:2]
            handle_connect(client_sock, addr, host, int(port))
        else:
            response = http_proxy_req(data)
            if response:
                client_sock.sendall(response)
    except:
        pass
    finally:
        try:
            client_sock.close()
        except:
            pass

def main():
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("127.0.0.1", PROXY_PORT))
    srv.listen(128)
    log(f"Listening 127.0.0.1:{PROXY_PORT} → {WORKER_HOST}")
    while True:
        c, a = srv.accept()
        threading.Thread(target=handle_client, args=(c, a), daemon=True).start()

if __name__ == "__main__":
    main()