from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from urllib.parse import urlparse
import httpx

ALLOWED_HOSTS = {
    "huggingface.co",
    "cas-bridge.xethub.hf.co",
    # Add more if you need: "cdn-lfs.huggingface.co", etc.
}

app = FastAPI(title="cozip CORS proxy")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "HEAD", "OPTIONS"],
    allow_headers=["Range", "Content-Type"],
    expose_headers=["Content-Range", "Accept-Ranges", "Content-Length", "Content-Type"],
    max_age=86400,
)


@app.get("/")
async def root():
    return {
        "service": "cozip CORS proxy",
        "usage": "/proxy?url=<cozip_url>",
        "allowed_hosts": sorted(ALLOWED_HOSTS),
    }


@app.api_route("/proxy", methods=["GET", "HEAD"])
async def proxy(url: str, request: Request):
    parsed = urlparse(url)
    if parsed.hostname not in ALLOWED_HOSTS:
        raise HTTPException(
            status_code=403,
            detail=f"Host not allowed: {parsed.hostname}",
        )

    fwd_headers = {}
    if "range" in request.headers:
        fwd_headers["Range"] = request.headers["range"]

    client = httpx.AsyncClient(follow_redirects=True, timeout=60.0)
    try:
        req = client.build_request(request.method, url, headers=fwd_headers)
        upstream = await client.send(req, stream=True)
    except Exception as e:
        await client.aclose()
        raise HTTPException(status_code=502, detail=f"Upstream error: {e}")

    response_headers = {}
    for h in ("content-range", "accept-ranges", "content-length", "content-type"):
        if h in upstream.headers:
            response_headers[h] = upstream.headers[h]

    async def stream():
        try:
            async for chunk in upstream.aiter_bytes(chunk_size=65536):
                yield chunk
        finally:
            await upstream.aclose()
            await client.aclose()

    return StreamingResponse(
        stream(),
        status_code=upstream.status_code,
        headers=response_headers,
    )