ovdashboard/api/ovkiosk/routers/image.py

import io
import logging
import re
import time
from typing import Iterator, Optional

from async_lru import alru_cache
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

_logger = logging.getLogger(__name__)

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


@router.on_event("startup")
async def on_startup():
    _logger.debug("image router started")


_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:
        pass


@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)


async def get_ttl_hash(seconds: int = 20) -> int:
    """
    Return the same value withing `seconds` time period
    https://stackoverflow.com/a/55900800
    """
    return round(time.time() / seconds)


@alru_cache(maxsize=20)
async def get_image_by_name(
    file_name: str,
    ttl_hash: Optional[int] = None,
) -> io.BytesIO:
    del ttl_hash

    file = DavFile(CLIENT.resource(f"img/{file_name}"), refresh=False)
    print(f"Downloading {file_name}")
    await file.download()
    img = Image.open(io.BytesIO(file.bytes)).convert("RGB")

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

    return img_buffer


@router.get(
    "/get/{prefix}",
    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),
    ttl_hash: int = Depends(get_ttl_hash),
) -> 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_buffer = await get_image_by_name(file_names[0], ttl_hash)
    img_buffer.seek(0)

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