import itertools
from fastapi import FastAPI, Request, Response, UploadFile, File, Form
from fastapi.responses import StreamingResponse, JSONResponse
from starlette.background import BackgroundTask
import httpx
import uvicorn
import os
import io
import uuid
import asyncio

app = FastAPI(title="Featherless Proxy")

# --- Featherless API Keys ---
api_keys_env = os.environ.get("FEATHERLESS_API_KEYS", "")
if api_keys_env:
    # Support comma-separated or newline-separated keys in env var
    api_keys = [k.strip() for k in api_keys_env.replace('\\n', ',').replace('\n', ',').split(",") if k.strip()]
else:
    API_KEYS_FILE = "api_keys.txt"
    if not os.path.exists(API_KEYS_FILE):
        API_KEYS_FILE = os.path.join(os.path.dirname(__file__), "api_keys.txt")
    
    if os.path.exists(API_KEYS_FILE):
        with open(API_KEYS_FILE, "r") as f:
            api_keys = [line.strip() for line in f if line.strip()]
    else:
        api_keys = []

if not api_keys:
    raise ValueError("No API keys found. Please set the FEATHERLESS_API_KEYS environment variable.")

key_pool = itertools.cycle(api_keys)
TARGET_BASE_URL = "https://api.featherless.ai"
client = httpx.AsyncClient(base_url=TARGET_BASE_URL, timeout=None)

# --- NIM API Keys ---
nim_keys_env = os.environ.get("NIM_API_KEYS", "")
if nim_keys_env:
    nim_keys = [k.strip() for k in nim_keys_env.replace('\\n', ',').replace('\n', ',').split(",") if k.strip()]
else:
    NIM_KEYS_FILE = os.path.join(os.path.dirname(__file__), "nim_keys.txt")
    nim_keys = []
    if os.path.exists(NIM_KEYS_FILE):
        with open(NIM_KEYS_FILE, "r") as f:
            nim_keys = [line.strip() for line in f if line.strip()]

nim_key_pool = itertools.cycle(nim_keys) if nim_keys else None
NIM_BASE_URL = "https://integrate.api.nvidia.com"
nim_client = httpx.AsyncClient(base_url=NIM_BASE_URL, timeout=120.0)

@app.on_event("shutdown")
async def shutdown_event():
    await client.aclose()
    if nim_client:
        await nim_client.aclose()

@app.get("/health")
async def health_check():
    return {"status": "ok"}

@app.post("/analyze")
async def analyze_content(
    text: str = Form(...)
):
    if not nim_key_pool:
        return JSONResponse(status_code=500, content={"error": "No NIM keys available"})
        
    next_key = next(nim_key_pool)
    headers = {
        "Authorization": f"Bearer {next_key}",
        "Content-Type": "application/json"
    }
    
    messages = []
    
    try:
        model = "deepseek-ai/deepseek-v4-flash-0731"
        if not text:
            return JSONResponse(status_code=400, content={"error": "Must provide text."})
        messages.append({"role": "user", "content": text})
            
        payload = {
            "model": model,
            "messages": messages,
            "max_tokens": 4096
        }
        
        try:
            resp = await nim_client.post("/v1/chat/completions", json=payload, headers=headers, timeout=60.0)
            resp.raise_for_status()
            data = resp.json()
        except Exception as primary_e:
            fallback_model = "meta/llama-3.1-8b-instruct"
            print(f"Primary model {model} failed/timed out ({type(primary_e).__name__}). Falling back to {fallback_model}...")
            payload["model"] = fallback_model
            resp = await nim_client.post("/v1/chat/completions", json=payload, headers=headers, timeout=30.0)
            resp.raise_for_status()
            data = resp.json()
        
        return JSONResponse(content={"result": data["choices"][0]["message"]["content"]})
        
    except Exception as e:
        err_msg = str(e) or repr(e)
        print(f"Analyze error: {err_msg}")
        return JSONResponse(status_code=500, content={"error": err_msg})

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def proxy(path: str, request: Request):
    next_key = next(key_pool)
    print(f"Using API Key: {next_key[:8]}... for path: {path}")

    # Forward headers, but replace Authorization with the selected key
    headers = dict(request.headers)
    headers.pop("host", None)  # Let httpx handle host
    headers["authorization"] = f"Bearer {next_key}"

    url = f"/{path}"
    
    # Forward query params
    params = dict(request.query_params)
    
    # Read the body if any
    body = await request.body()
    
    req = client.build_request(
        method=request.method,
        url=url,
        content=body,
        headers=headers,
        params=params,
    )
    
    try:
        response = await client.send(req, stream=True)
        
        response_headers = dict(response.headers)
        # Remove encodings like gzip/deflate if we are proxying, 
        # or just pass through. It's safer to delete 'content-encoding' and 'content-length' 
        # when streaming because Starlette will set them.
        response_headers.pop("content-encoding", None)
        response_headers.pop("content-length", None)
        return StreamingResponse(
            response.aiter_bytes(), 
            status_code=response.status_code, 
            headers=response_headers,
            background=BackgroundTask(response.aclose)
        )
    except Exception as e:
        print(f"Proxy error: {e}")
        return Response(content=f"Proxy error: {str(e)}", status_code=500)

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