import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import json

app = FastAPI()

@app.post("/v1/openai/chat/completions")
async def proxy_deepinfra(request: Request):
    """
    Proxies chat completion requests to the DeepInfra API.
    """
    client = httpx.AsyncClient()
    url = "https://api.deepinfra.com/v1/openai/chat/completions"

    # Extract the raw body from the request
    body = await request.body()
    # Decode the body to a string and then load it as JSON
    try:
        data = json.loads(body.decode('utf-8'))
    except json.JSONDecodeError:
        return {"error": "Invalid JSON in request body"}, 400

    headers = {
        'sec-ch-ua-platform': request.headers.get('sec-ch-ua-platform', '"Windows"'),
        'Referer': request.headers.get('Referer', 'https://deepinfra.com/'),
        'sec-ch-ua': request.headers.get('sec-ch-ua', '"Chromium";v="140", "Not=A?Brand";v="24", "Google Chrome";v="140"'),
        'sec-ch-ua-mobile': request.headers.get('sec-ch-ua-mobile', '?0'),
        'User-Agent': request.headers.get('User-Agent', 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36'),
        'accept': request.headers.get('accept', 'text/event-stream'),
        'X-Deepinfra-Source': request.headers.get('X-Deepinfra-Source', 'web-embed'),
        'Content-Type': request.headers.get('Content-Type', 'application/json'),
    }

    async def stream_response():
        async with client.stream("POST", url, headers=headers, json=data, timeout=None) as response:
            async for chunk in response.aiter_bytes():
                yield chunk

    return StreamingResponse(stream_response(), media_type="text/event-stream")

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