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

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

def get_proxy_response(path: str, query: str = "", method: str = "GET", headers: dict = {}, body: bytes = None):
    url = f"{BASE_URL}/{path}"
    if query:
        url += f"?{query}"
    
    # Filter headers to avoid conflicts
    filtered_headers = {k: v for k, v in headers.items() if k.lower() not in ["host", "content-length", "connection"]}
    if HF_TOKEN and "authorization" not in [k.lower() for k in filtered_headers.keys()]:
        filtered_headers["Authorization"] = f"Bearer {HF_TOKEN}"
        
    try:
        resp = requests.request(
            method=method,
            url=url,
            headers=filtered_headers,
            data=body,
            stream=True,
            allow_redirects=True
        )
        
        # Exclude headers that FastAPI/Uvicorn might handle or conflict with
        excluded = ["content-encoding", "transfer-encoding", "connection", "keep-alive"]
        resp_headers = {k: v for k, v in resp.headers.items() if k.lower() not in excluded}
        
        return StreamingResponse(
            resp.iter_content(chunk_size=128*1024),
            status_code=resp.status_code,
            headers=resp_headers,
            media_type=resp.headers.get("content-type")
        )
    except Exception as e:
        return Response(content=str(e), status_code=500)

@app.get("/resolve/{repo_id:path}/{filename:path}")
async def old_resolve_compat(repo_id: str, filename: str):
    """Keep compatibility with the previous /resolve structure."""
    return get_proxy_response(f"{repo_id}/resolve/main/{filename}")

@app.api_route("/{path:path}", methods=["GET", "HEAD", "OPTIONS", "POST"])
async def catch_all(request: Request, path: str):
    """Universal proxy to Hugging Face."""
    body = await request.body() if request.method == "POST" else None
    return get_proxy_response(path, request.url.query, request.method, dict(request.headers), body)