ovdashboard/api/ovdashboard_api/routers/image.py

import re
from io import BytesIO
from typing import Iterator

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.responses import StreamingResponse
from PIL import Image
from webdav3.exceptions import RemoteResourceNotFound

from .. import CLIENT
from ..dav_file import DavFile

router = APIRouter(prefix="/image", tags=["image"])


_re_image_file = re.compile(
    r"\.(gif|jpe?g|tiff?|png|bmp)$",
    flags=re.IGNORECASE,
)


async def get_image_file_names() -> Iterator[str]:
    try:
        file_names = CLIENT.list("img")

        return (
            name
            for name in file_names
            if _re_image_file.search(name)
        )

    except RemoteResourceNotFound:
        return iter(())


@router.get("/list", response_model=list[str])
async def list_images(
    image_file_names: Iterator[str] = Depends(get_image_file_names),
) -> list[str]:
    return list(image_file_names)


async def find_file_names(
    prefix: str = "",
    image_file_names: Iterator[str] = Depends(get_image_file_names),
) -> Iterator[str]:
    return (
        file_name
        for file_name in image_file_names
        if file_name.lower().startswith(prefix.lower())
    )


@router.get("/find/{prefix}", response_model=list[str])
async def find_images(
    file_names: Iterator[str] = Depends(find_file_names),
) -> list[str]:
    return list(file_names)


@router.get(
    "/get/{prefix}",
    response_class=StreamingResponse,
    responses={
        status.HTTP_200_OK: {
            "description": "Operation successful",
        },
        status.HTTP_404_NOT_FOUND: {
            "description": "image file not found",
            "content": None,
        },
        status.HTTP_409_CONFLICT: {
            "description": "ambiguous image file name",
            "content": None,
        },
    },
)
async def get_image(
    prefix: str,
    file_names: Iterator[str] = Depends(find_file_names),
) -> StreamingResponse:
    file_names = list(file_names)

    if not (file_names := list(file_names)):
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND)

    elif len(file_names) > 1:
        raise HTTPException(status_code=status.HTTP_409_CONFLICT)

    img_file = DavFile(f"img/{file_names[0]}")
    img = Image.open(BytesIO(await img_file.bytes)).convert("RGB")

    img_buffer = BytesIO()
    img.save(img_buffer, format='JPEG', quality=85)
    img_buffer.seek(0)

    return StreamingResponse(
        content=img_buffer,
        media_type="image/jpeg",
        headers={
            "Content-Disposition": f"filename={prefix}.jpg"
        },
    )