from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import requests
import os

app = FastAPI()

# Optional: Add your HF_TOKEN to Space Secrets for private models
HF_TOKEN = os.environ.get("HF_TOKEN")

@app.get("/")
def home():
    return {"status": "running", "message": "HF Storage Proxy"}

@app.get("/resolve/{repo_id:path}/{filename:path}")
async def proxy_resolve(repo_id: str, filename: str):
    """
    Proxies requests to huggingface.co/repo/resolve/main/file
    """
    hf_url = f"https://huggingface.co/{repo_id}/resolve/main/{filename}"
    
    headers = {}
    if HF_TOKEN:
        headers["Authorization"] = f"Bearer {HF_TOKEN}"
    
    try:
        # We use stream=True to handle large model files without loading into memory
        response = requests.get(hf_url, headers=headers, stream=True, allow_redirects=True)
        response.raise_for_status()
        
        return StreamingResponse(
            response.iter_content(chunk_size=1024*1024),
            media_type=response.headers.get("content-type"),
            headers={
                "Content-Disposition": response.headers.get("Content-Disposition", f'attachment; filename="{filename}"'),
                "Content-Length": response.headers.get("Content-Length")
            }
        )
    except requests.exceptions.HTTPError as e:
        raise HTTPException(status_code=e.response.status_code, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))