# 开发服务器，用于在不构建vite应用的情况下仍然像生产环境中那样使用相对路径调用api，兼容Vite HMR
# 如果出现Bug请注意先检查是不是proxy导致的问题，再检查是不是代码写错了！！！

import asyncio
import aiohttp
from aiohttp import web
import logging

# 全局变量配置
VITE_ADD = "http://localhost:5173"
PY_ADD = "http://127.0.0.1:5000"
PORT = 8080
MAX_BODY_MB = 50

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("dev_proxy")

async def proxy_handler(request):
    path = request.path
    query = request.query_string
    method = request.method
    headers = dict(request.headers)
    
    # 判定转发目标
    if path.startswith('/api/'):
        target_url = f"{PY_ADD}{path}"
    else:
        target_url = f"{VITE_ADD}{path}"
        
    if query:
        target_url += f"?{query}"

    # 处理 WebSocket 升级请求 (Vite HMR)
    if request.headers.get('Upgrade', '').lower() == 'websocket':
        return await ws_proxy_handler(request, target_url)

    try:
        async with aiohttp.ClientSession() as session:
            # 读取请求体
            data = await request.read()
            
            async with session.request(
                method=method,
                url=target_url,
                headers=headers,
                data=data,
                allow_redirects=False
            ) as resp:
                # 构建响应头
                resp_headers = dict(resp.headers)
                
                # 读取响应体
                body = await resp.read()
                return web.Response(
                    body=body,
                    status=resp.status,
                    headers=resp_headers
                )
    except Exception as e:
        logger.error(f"Error proxying {method} {path}: {e}")
        return web.Response(text=str(e), status=502)

async def ws_proxy_handler(request, target_url):
    """处理 WebSocket 转发，主要用于 Vite HMR"""
    # 将 http:// 转换为 ws://
    ws_target_url = target_url.replace('http://', 'ws://').replace('https://', 'wss://')
    
    # 获取并转发 WebSocket 子协议
    protocol = request.headers.get('Sec-WebSocket-Protocol', '')
    protocols = [p.strip() for p in protocol.split(',')] if protocol else []
    
    ws_server = web.WebSocketResponse(protocols=protocols)
    await ws_server.prepare(request)
    
    logger.info(f"WebSocket connection upgraded: {request.path} (Protocols: {protocols})")

    async with aiohttp.ClientSession() as session:
        async with session.ws_connect(ws_target_url, protocols=protocols) as ws_client:
            
            async def forward(src, dst):
                async for msg in src:
                    if msg.type == aiohttp.WSMsgType.TEXT:
                        await dst.send_str(msg.data)
                    elif msg.type == aiohttp.WSMsgType.BINARY:
                        await dst.send_bytes(msg.data)
                    elif msg.type == aiohttp.WSMsgType.CLOSE:
                        await dst.close()
                        break
                    elif msg.type == aiohttp.WSMsgType.ERROR:
                        break

            # 双向转发
            await asyncio.gather(
                forward(ws_server, ws_client),
                forward(ws_client, ws_server)
            )

    return ws_server

async def main():
    app = web.Application(client_max_size=MAX_BODY_MB * 1024 * 1024)
    # 捕获所有路径
    app.router.add_route('*', '/{path:.*}', proxy_handler)
    
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '127.0.0.1', PORT)
    
    logger.info(f"Dev Proxy started at http://127.0.0.1:{PORT}")
    logger.info(f"Routing /api/* to {PY_ADD}")
    logger.info(f"Routing others to {VITE_ADD}")
    
    await site.start()
    
    # 保持运行
    try:
        while True:
            await asyncio.sleep(3600)
    except KeyboardInterrupt:
        pass
    finally:
        await runner.cleanup()

if __name__ == "__main__":
    asyncio.run(main())