from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
import httpx
import random

app = FastAPI()

TELEGRAPH_URL = 'https://generativelanguage.googleapis.com/v1beta'

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"])
async def proxy(path: str, request: Request):
    try:
        url = f"{TELEGRAPH_URL}/{path}"
        url += f"?{request.query_params}" if request.query_params else ""

        provided_api_keys = request.query_params.get('key')

        if not provided_api_keys:
            raise HTTPException(status_code=400, detail='API key is missing.')

        api_key_array = [key.strip() for key in provided_api_keys.split(';') if key.strip()]
        if not api_key_array:
            raise HTTPException(status_code=400, detail='Valid API key is missing.')

        selected_api_key = random.choice(api_key_array)
        url = url.replace(provided_api_keys, selected_api_key)

        headers = {key: value for key, value in request.headers.items() if key.lower() != 'host'}
        headers['content-type'] = 'application/json'

        async with httpx.AsyncClient() as client:
            content = await request.body()
            
            if content:
                try:
                    original_body = await request.json()
                except:
                    original_body = {}
                new_body = {**original_body}
            else:
                new_body = {}

            new_body["safetySettings"] = [
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_NONE"
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH",
                    "threshold": "BLOCK_NONE"
                },
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_NONE"
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_NONE"
                },
                {
                    "category": "HARM_CATEGORY_CIVIC_INTEGRITY",
                    "threshold": "BLOCK_NONE"
                }
            ]
            
            # 使用 httpx 发起代理请求
            r = await client.request(
                method=request.method,
                url=url,
                headers=headers,
                content=content if not new_body else JSONResponse(content=new_body).body, # 将 new_body 转换为 JSON 字符串
                timeout=60.0,
            )

            # 流式传输响应
            return StreamingResponse(
                content=r.aiter_bytes(),
                status_code=r.status_code,
                headers=r.headers,
                media_type=r.headers.get("content-type")
            )

    except HTTPException as e:
        return JSONResponse(content={"detail": e.detail}, status_code=e.status_code)

    except Exception as e:
        return JSONResponse(content={"detail": "An error occurred: " + str(e)}, status_code=500)