from fastapi import HTTPException, Request, Response
from nicegui import run

from xmas_app.util.db import get_doc_by_name


async def get_document_from_db(
    plan_id: str, filename: str, request: Request
) -> Response:
    """Retrieves a document from the database.

    Supports GET, HEAD as well as range requests, e.g. from GDAL regarding COG.
    """

    def _parse_range(range_header: str, file_size: int) -> tuple[int, int] | None:
        try:
            units, range_spec = range_header.strip().split("=")
            if units != "bytes":
                return None

            if "," in range_spec:
                return None

            start_str, end_str = range_spec.split("-")

            start = int(start_str) if start_str else 0
            end = int(end_str) if end_str else file_size - 1

            if start < 0 or end >= file_size or start > end:
                return None

            return start, end
        except Exception:
            return None

    doc = await run.io_bound(get_doc_by_name, plan_id, filename)
    if not doc:
        return Response(status_code=404)

    image_bytes = doc.filedata
    file_size = len(image_bytes)

    headers = {
        "Content-Type": doc.filetype,
        "Accept-Ranges": "bytes",
    }

    range_header = request.headers.get("range")

    if range_header:
        byte_range = _parse_range(range_header, file_size)
        if byte_range is None:
            raise HTTPException(status_code=416)

        start, end = byte_range

        chunk = memoryview(image_bytes)[start : end + 1]

        headers.update(
            {
                "Content-Range": f"bytes {start}-{end}/{file_size}",
                "Content-Length": str(end - start + 1),
            }
        )

        if request.method == "HEAD":
            return Response(status_code=206, headers=headers)

        return Response(
            content=chunk,
            status_code=206,
            headers=headers,
        )

    headers["Content-Length"] = str(file_size)

    if request.method == "HEAD":
        return Response(status_code=200, headers=headers)

    return Response(
        content=image_bytes,
        status_code=200,
        headers=headers,
    )
