from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse, Response
import httpx
import uvicorn
import asyncio

app = FastAPI()

@app.api_route("/v1/{url_path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
async def proxy_request(url_path: str, request: Request):
    print(f"接收到的 url_path: {url_path}")

    # url_path 示例: https/open.bigmodel.cn/api/paas/v4/chat/completions
    # 找到第一个 '/' 的位置，分隔协议和域名
    first_slash_idx = url_path.find('/')
    if first_slash_idx == -1:
        raise HTTPException(status_code=400, detail="无效的URL路径格式。期望协议/域名/路径。")
    
    protocol = url_path[:first_slash_idx] # 'https'
    print(f"解析出的协议: {protocol}")

    # 找到第二个 '/' 的位置，分隔域名和实际的路径
    second_slash_idx = url_path.find('/', first_slash_idx + 1)

    if second_slash_idx == -1:
        domain = url_path[first_slash_idx + 1:]
        remaining_path = ''
    else:
        domain = url_path[first_slash_idx + 1:second_slash_idx]
        remaining_path = url_path[second_slash_idx:]

    target_url = f"{protocol}://{domain}{remaining_path}"
    print(f"\n\n\n代理请求到 {target_url}")

    # 转发原始请求的头部，排除 'Host' 头部以避免冲突
    # FastAPI 的 request.headers 是不可变的，需要转换为字典
    headers = {key: value for key, value in request.headers.items() if key.lower() != 'host'}
    
    # 获取请求体
    request_body = await request.body()
    
    # 获取查询参数
    query_params = request.query_params

    async with httpx.AsyncClient(verify=True, follow_redirects=False) as client:
        try:
            # 使用 httpx 库向目标 URL 发送请求
            resp = await client.request(
                method=request.method,
                url=target_url,
                headers=headers,
                content=request_body, # 使用 content 传递请求体
                params=query_params,
                timeout=30.0, # 设置超时时间为30秒
            )

            # 打印目标 API 返回的实际状态码和响应体，用于调试
            print(f"目标API响应状态码: {resp.status_code}")
            print(f"目标API响应体: {resp.text[:500]}...") # 打印前500个字符，避免过长

            # 构建响应头部
            excluded_headers = ['content-encoding'] # 保持与 Flask 版本一致
            response_headers = {
                name: value for name, value in resp.headers.items()
                if name.lower() not in excluded_headers
            }
            
            # 返回流式响应内容
            # httpx 的 .aiter_bytes() 返回异步迭代器
            async def generate_response():
                async for chunk in resp.aiter_bytes(chunk_size=8192):
                    yield chunk

            return StreamingResponse(generate_response(), status_code=resp.status_code, headers=response_headers)

        except httpx.RequestError as e:
            error_detail = f"代理请求到 {target_url} 失败: {type(e).__name__} - {e}"
            print(f"代理请求失败: {error_detail}")
            if e.request:
                print(f"请求信息: {e.request.method} {e.request.url}")
            if hasattr(e, 'response') and e.response:
                print(f"响应信息: {e.response.status_code} {e.response.text[:200]}...")
            raise HTTPException(status_code=500, detail=error_detail)

# if __name__ == '__main__':
#     # 提示：请确保您已激活 conda 环境 'any-api' (conda activate any-api)
#     # 提示：请确保已安装 FastAPI, Uvicorn 和 httpx 库 (pip install fastapi uvicorn httpx)
#     print(f"代理服务器正在 0.0.0.0:7860 上启动")
#     uvicorn.run(app, host="0.0.0.0", port=7860)