#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
本地代理桥接：在本机起一个 HTTP 代理端口，把请求转发到本机真实代理（如 127.0.0.1:7890）。
再用 ngrok 暴露该端口，把 ngrok 公网地址填给 7006 的 Proxy，7006 即可通过你的本地代理出网。

用法:
  python scripts/local_proxy_bridge.py
  # 默认监听 0.0.0.0:8899，上游代理 127.0.0.1:7890（本机 Clash/V2Ray 等）
  # 再开一个终端: ngrok http 8899  （或 ngrok tcp 8899，则填 http://0.tcp.ngrok.io:端口）
  # 把 ngrok 生成的公网 URL 填到 .env 的 HTTP_PROXY / HTTPS_PROXY，7006 即可通过你的本机代理出网

环境变量（可选）:
  PROXY_BRIDGE_LISTEN=0.0.0.0:8899   # 监听地址
  PROXY_BRIDGE_UPSTREAM=127.0.0.1:7890  # 上游代理（本机 Clash/V2Ray 等）
"""

import asyncio
import os
import sys

# 可选：把项目根加入 path
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_ROOT = os.path.dirname(_SCRIPT_DIR)
if _ROOT not in sys.path:
    sys.path.insert(0, _ROOT)


def _parse_addr(s: str, default_host: str, default_port: int):
    s = (s or "").strip()
    if not s:
        return default_host, default_port
    if ":" in s:
        host, _, port = s.rpartition(":")
        return host or default_host, int(port) if port else default_port
    return default_host, int(s) if s.isdigit() else default_port


async def _relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    """双向转发：reader -> writer 直到 EOF。"""
    try:
        while True:
            data = await reader.read(65536)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError):
        pass
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


async def _handle_client(
    client_reader: asyncio.StreamReader,
    client_writer: asyncio.StreamWriter,
    upstream_host: str,
    upstream_port: int,
):
    """处理一个客户端连接：把首包（CONNECT/GET 等）转发到上游代理，再双向 relay。"""
    try:
        # 读首行 + headers（到 \r\n\r\n）
        first_line = await client_reader.readline()
        if not first_line:
            return
        header_lines = []
        while True:
            line = await client_reader.readline()
            if line in (b"\r\n", b"\n"):
                break
            header_lines.append(line)

        request_head = first_line + b"".join(header_lines) + b"\r\n"

        # 连上游代理
        try:
            up_reader, up_writer = await asyncio.wait_for(
                asyncio.open_connection(upstream_host, upstream_port), timeout=10.0
            )
        except Exception as e:
            print(f"[proxy-bridge] upstream connect failed: {e}", flush=True)
            client_writer.write(
                b"HTTP/1.1 502 Bad Gateway\r\nConnection: close\r\n\r\n"
                b"Upstream proxy connect failed"
            )
            await client_writer.drain()
            client_writer.close()
            return

        # 把请求头发给上游
        up_writer.write(request_head)
        await up_writer.drain()

        # CONNECT 时上游会先回 200 Connection Established，需要把这部分先读完并回给客户端，再双向 relay
        # 非 CONNECT 时上游直接回响应，也要先读完并回给客户端
        # 为简单起见：先读上游的响应头（到 \r\n\r\n），转发给客户端，然后双向 relay 剩余 body/隧道
        up_buf = b""
        while b"\r\n\r\n" not in up_buf and len(up_buf) < 65536:
            chunk = await up_reader.read(4096)
            if not chunk:
                break
            up_buf += chunk
        if up_buf:
            client_writer.write(up_buf)
            await client_writer.drain()

        # 双向转发剩余数据（CONNECT 隧道或响应 body）
        await asyncio.gather(
            _relay(client_reader, up_writer),
            _relay(up_reader, client_writer),
        )
    except Exception as e:
        print(f"[proxy-bridge] handle error: {e}", flush=True)
    finally:
        try:
            client_writer.close()
            await client_writer.wait_closed()
        except Exception:
            pass


async def _run(listen_host: str, listen_port: int, upstream_host: str, upstream_port: int):
    server = await asyncio.start_server(
        lambda r, w: _handle_client(r, w, upstream_host, upstream_port),
        listen_host,
        listen_port,
    )
    addrs = ", ".join(str(s.getsockname()) for s in server.sockets)
    print(f"[proxy-bridge] listening on {addrs}, upstream={upstream_host}:{upstream_port}", flush=True)
    print(f"[proxy-bridge] expose with: ngrok http {listen_port}", flush=True)
    async with server:
        await server.serve_forever()


def main():
    listen_spec = os.environ.get("PROXY_BRIDGE_LISTEN", "0.0.0.0:8899")
    upstream_spec = os.environ.get("PROXY_BRIDGE_UPSTREAM", "127.0.0.1:7890")
    listen_host, listen_port = _parse_addr(listen_spec, "0.0.0.0", 8899)
    upstream_host, upstream_port = _parse_addr(upstream_spec, "127.0.0.1", 7890)
    asyncio.run(_run(listen_host, listen_port, upstream_host, upstream_port))


if __name__ == "__main__":
    main()