from fastapi import FastAPI, Request, Response, status, WebSocket, WebSocketDisconnect
import httpx
import uvicorn
import asyncio
import websockets

app = FastAPI()

TARGET_BASE = "http://127.0.0.1:9222"  # 目标服务器（Chrome DevTools）
TARGET_WS_BASE = "ws://127.0.0.1:9222"  # WebSocket 目标

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def proxy(request: Request, path: str):

    url = f"{TARGET_BASE}/{path}"
    headers = dict(request.headers)
    # 强制 Host 头为 127.0.0.1，避免 Chrome 拒绝
    headers["host"] = "127.0.0.1"
    body = await request.body()
    async with httpx.AsyncClient(follow_redirects=True) as client:
        resp = await client.request(
            request.method,
            url,
            headers=headers,
            content=body,
            params=request.query_params
        )
    return Response(
        content=resp.content,
        status_code=resp.status_code,
        headers=dict(resp.headers)
    )

@app.websocket("/{path:path}")
async def websocket_proxy(websocket: WebSocket, path: str):
    await websocket.accept()
    target_url = f"{TARGET_WS_BASE}/{path}"
    try:
        async with websockets.connect(target_url) as target_ws:
            async def client_to_server():
                while True:
                    data = await websocket.receive_text()
                    await target_ws.send(data)

            async def server_to_client():
                while True:
                    data = await target_ws.recv()
                    await websocket.send_text(data)

            await asyncio.gather(client_to_server(), server_to_client())
    except WebSocketDisconnect:
        pass
    except Exception:
        await websocket.close()

if __name__ == "__main__":
    uvicorn.run("fun:app", host="0.0.0.0", port=8000)