"""Stream authenticated storage video ranges without downloading the full file first."""

import httpx
from fastapi import HTTPException, Request
from huggingface_hub import hf_hub_url
from starlette.background import BackgroundTask
from starlette.responses import Response, StreamingResponse


async def video_response(request: Request, repo_id: str, path: str, token: str):
    # Browser cookies and user headers never go to the storage service.
    headers = {"Authorization": f"Bearer {token}", "Accept-Encoding": "identity"}
    for name in ("range", "if-range"):
        if request.headers.get(name):
            headers[name] = request.headers[name]
    client = httpx.AsyncClient(follow_redirects=True, timeout=httpx.Timeout(60, connect=15))
    try:
        upstream = await client.send(client.build_request(
            request.method, hf_hub_url(repo_id, path, repo_type="dataset"), headers=headers), stream=True)
    except httpx.HTTPError:
        await client.aclose()
        raise HTTPException(502, "Video storage is temporarily unavailable") from None

    async def close():
        try:
            await upstream.aclose()
        finally:
            await client.aclose()

    if upstream.status_code not in (200, 206, 416):
        await close()
        raise HTTPException(404 if upstream.status_code == 404 else 502,
                            "Video not found" if upstream.status_code == 404 else "Video storage is temporarily unavailable")
    response_headers = {name: upstream.headers[name] for name in (
        "content-type", "content-length", "content-range", "accept-ranges", "etag", "last-modified"
    ) if name in upstream.headers}
    response_headers["Cache-Control"] = "private, no-store"
    if request.method == "HEAD" or upstream.status_code == 416:
        await close()
        # A 416 may have a storage error body: don't forward its body length with an empty response.
        if upstream.status_code == 416:
            response_headers.pop("content-length", None)
        return Response(status_code=upstream.status_code, headers=response_headers)

    async def chunks():
        try:
            async for chunk in upstream.aiter_raw():
                yield chunk
        finally:
            await close()

    return StreamingResponse(chunks(), status_code=upstream.status_code, headers=response_headers,
                             background=BackgroundTask(close))