import asyncio
import os
import json
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse, JSONResponse
import aiohttp

# ================= 配置区 =================
TARGET_BASE_URL = "https://api.deepinfra.com/v1/openai"
REAL_API_KEY = os.environ["REAL_API_KEY"]                # 从 Space Secrets 读取
FAKE_KEY = os.environ.get("FAKE_KEY", "123456")         # 你本地使用的假 key

# Space 容器可直接访问外网，无需代理
LOCAL_PROXY = None

# ==========================================

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

@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.session = aiohttp.ClientSession()
    yield
    if app.state.session:
        await app.state.session.close()

app = FastAPI(lifespan=lifespan)

@app.api_route("/v1", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"])
async def proxy_endpoint(path: str = "", request: Request = None):
    if request.method == "OPTIONS":
        return Response(status_code=200, headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
            "Access-Control-Allow-Headers": "Authorization, Content-Type"
        })

    # 验证本地传来的假 key
    auth_header = request.headers.get("authorization", "")
    if auth_header.lower() != f"bearer {FAKE_KEY}":
        return JSONResponse(
            status_code=401,
            content={"error": {"message": "Invalid local API key", "type": "invalid_request_error"}}
        )

    target_url = f"{TARGET_BASE_URL.rstrip('/')}/{path}" if path else TARGET_BASE_URL.rstrip('/')

    # 构建转发 headers
    headers = {}
    for key, value in request.headers.items():
        if key.lower() in ["host", "content-length", "connection", "authorization", "transfer-encoding"]:
            continue
        headers[key] = value
    headers["Authorization"] = f"Bearer {REAL_API_KEY}"

    body = await request.body()
    query_params = dict(request.query_params)

    # 判断是否流式请求
    is_stream = False
    if body:
        try:
            is_stream = json.loads(body).get("stream", False)
        except json.JSONDecodeError:
            pass

    session = request.app.state.session
    timeout = aiohttp.ClientTimeout(total=None, connect=10, sock_read=60)

    # 转发请求到 DeepInfra
    resp = await session.request(
        method=request.method,
        url=target_url,
        headers=headers,
        data=body,
        params=query_params,
        timeout=timeout,
        proxy=LOCAL_PROXY
    )

    # 流式响应
    if is_stream or "text/event-stream" in resp.headers.get("content-type", ""):
        async def stream_generator():
            try:
                async for chunk in resp.content.iter_any():
                    yield chunk
            except (aiohttp.ClientConnectionError, aiohttp.ServerDisconnectedError,
                    aiohttp.ClientPayloadError, TimeoutError, asyncio.TimeoutError,
                    asyncio.CancelledError):
                pass
            finally:
                resp.close()

        resp_headers = {
            "content-type": resp.headers.get("content-type", "text/event-stream"),
            "cache-control": "no-cache",
            "connection": "keep-alive",
            "access-control-allow-origin": "*"
        }
        return StreamingResponse(
            stream_generator(),
            media_type=resp_headers["content-type"],
            headers=resp_headers
        )

    # 非流式响应
    else:
        content = await resp.read()
        resp.close()

        # 调试日志（可保留，帮助排查）
        logger.warning(f"UPSTREAM STATUS: {resp.status}")
        logger.warning(f"UPSTREAM BODY LENGTH: {len(content)}")
        if content:
            logger.warning(f"UPSTREAM BODY PREVIEW: {content[:200]}")

        resp_headers = {
            "content-type": resp.headers.get("content-type", "application/json"),
            "access-control-allow-origin": "*"
        }
        # 避免某些编码头干扰客户端解析
        for k in ["content-encoding", "transfer-encoding"]:
            resp_headers.pop(k, None)

        return Response(content=content, status_code=resp.status, headers=resp_headers)