import httpx
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
import json
import random # Import the random library

app = FastAPI()

# List of API URLs to be randomized
API_URLS = [
    "https://api.deepinfra.com/v1/openai/chat/completions",
    "https://stage.api.deepinfra.com/v1/openai/chat/completions",
]

@app.post("/v1/openai/chat/completions")
async def proxy_deepinfra(request: Request):
    """
    Proxies chat completion requests to the DeepInfra API.
    It randomizes the order of API URLs and uses the next as a fallback.
    """
    try:
        body = await request.json()
    except json.JSONDecodeError:
        raise HTTPException(status_code=400, detail="Invalid JSON in request body")

    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'),
    }

    # Create a randomized list of URLs for this specific request
    shuffled_urls = random.sample(API_URLS, len(API_URLS))

    async def stream_generator():
        last_error = None
        for url in shuffled_urls:
            print(f"Attempting to connect to: {url}")
            try:
                async with httpx.AsyncClient() as client:
                    # Setting timeout=None disables client-side timeouts
                    async with client.stream("POST", url, headers=headers, json=body, timeout=None) as response:
                        response.raise_for_status()
                        # If successful, print the URL being used
                        print(f"Successfully connected. Streaming from: {url}")
                        async for chunk in response.aiter_bytes():
                            yield chunk
                        return # Exit after successful stream
            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                last_error = e
                print(f"Failed to connect to {url}: {e}. Trying next URL.")
                continue

        # This part is reached only if all URLs in the shuffled list failed
        if last_error:
            raise HTTPException(status_code=502, detail=f"All API endpoints failed. Last error: {last_error}")

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

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