from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
import httpx
import os

app = FastAPI()
HF_TOKEN = os.environ.get("HF_TOKEN")
BASE_URL = "https://huggingface.co"

# Create a shared async client
client = httpx.AsyncClient(base_url=BASE_URL, follow_redirects=True)

@app.get("/")
def home():
    return {"status": "running", "message": "Universal HF Proxy Active", "token_detected": bool(HF_TOKEN)}

@app.api_route("/{path:path}", methods=["GET", "HEAD", "OPTIONS", "POST"])
async def proxy(request: Request, path: str):
    url = f"{BASE_URL}/{path}"
    if request.url.query:
        url += f"?{request.url.query}"
    
    print(f"Proxying: {request.method} {url}")
    
    # Filter headers
    headers = {k: v for k, v in request.headers.items() if k.lower() not in ["host", "content-length", "connection"]}
    if HF_TOKEN and "authorization" not in [k.lower() for k in headers.keys()]:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"
    
    try:
        # Build request
        body = await request.body()
        rp_req = client.build_request(request.method, url, headers=headers, content=body)
        rp_resp = await client.send(rp_req, stream=True)
        
        # Filter response headers
        excluded = ["content-encoding", "transfer-encoding", "connection", "keep-alive"]
        resp_headers = {k: v for k, v in rp_resp.headers.items() if k.lower() not in excluded}
        
        return StreamingResponse(
            rp_resp.aiter_raw(),
            status_code=rp_resp.status_code,
            headers=resp_headers,
            media_type=rp_resp.headers.get("content-type")
        )
    except Exception as e:
        print(f"Proxy Error: {e}")
        return Response(content=str(e), status_code=500)