Foxel/services/thumbnail.py

from __future__ import annotations
import asyncio
import inspect
import io
import hashlib
import tempfile
from contextlib import suppress
from pathlib import Path
from typing import Tuple
from fastapi import HTTPException

ALLOWED_EXT = {"jpg", "jpeg", "png", "webp", "gif", "bmp",
               "tiff", "arw", "cr2", "cr3", "nef", "rw2", "orf", "pef", "dng"}
RAW_EXT = {"arw", "cr2", "cr3", "nef", "rw2", "orf", "pef", "dng"}
VIDEO_EXT = {"mp4", "mov", "m4v", "avi", "mkv", "wmv", "flv", "webm", "mpg", "mpeg", "3gp"}
MAX_IMAGE_SOURCE_SIZE = 200 * 1024 * 1024
VIDEO_RANGE_LIMIT = 16 * 1024 * 1024  # 16MB
VIDEO_INITIAL_CHUNK = 4 * 1024 * 1024
CACHE_ROOT = Path('data/.thumb_cache')


def is_image_filename(name: str) -> bool:
    parts = name.rsplit('.', 1)
    if len(parts) < 2:
        return False
    return parts[1].lower() in ALLOWED_EXT


def is_raw_filename(name: str) -> bool:
    parts = name.rsplit('.', 1)
    if len(parts) < 2:
        return False
    return parts[1].lower() in RAW_EXT


def is_video_filename(name: str) -> bool:
    parts = name.rsplit('.', 1)
    if len(parts) < 2:
        return False
    return parts[1].lower() in VIDEO_EXT


def _cache_key(adapter_id: int, rel: str, size: int, mtime: int, w: int, h: int, fit: str) -> str:
    raw = f"{adapter_id}|{rel}|{size}|{mtime}|{w}x{h}|{fit}".encode()
    return hashlib.sha1(raw).hexdigest()


def _cache_path(key: str) -> Path:
    sub = Path(key[:2]) / key[2:4]
    return CACHE_ROOT / sub / f"{key}.webp"


def _ensure_cache_dir(p: Path):
    p.parent.mkdir(parents=True, exist_ok=True)


def _image_to_webp(im, w: int, h: int, fit: str) -> Tuple[bytes, str]:
    from PIL import Image
    if im.mode not in ("RGB", "RGBA"):
        im = im.convert("RGBA" if im.mode in ("P", "LA") else "RGB")
    if fit == 'cover':
        im_ratio = im.width / im.height
        target_ratio = w / h
        if im_ratio > target_ratio:
            new_h = h
            new_w = int(h * im_ratio)
        else:
            new_w = w
            new_h = int(w / im_ratio)
        im = im.resize((new_w, new_h))
        left = max(0, (im.width - w)//2)
        top = max(0, (im.height - h)//2)
        im = im.crop((left, top, left + w, top + h))
    else:
        im.thumbnail((w, h))
    buf = io.BytesIO()
    im.save(buf, 'WEBP', quality=80)
    return buf.getvalue(), 'image/webp'


def generate_thumb(data: bytes, w: int, h: int, fit: str, is_raw: bool = False) -> Tuple[bytes, str]:
    from PIL import Image
    if is_raw:
        try:
            import rawpy
            with rawpy.imread(io.BytesIO(data)) as raw:
                try:
                    thumb = raw.extract_thumb()
                except rawpy.LibRawNoThumbnailError:
                    thumb = None

                if thumb is not None and thumb.format in [rawpy.ThumbFormat.JPEG, rawpy.ThumbFormat.BITMAP]:
                    im = Image.open(io.BytesIO(thumb.data))
                else:
                    rgb = raw.postprocess(
                        use_camera_wb=False, use_auto_wb=True, output_bps=8)
                    im = Image.fromarray(rgb)
        except Exception as e:
            print(f"rawpy processing failed: {e}")
            raise e

    else:
        im = Image.open(io.BytesIO(data))

    return _image_to_webp(im, w, h, fit)


async def _collect_response_bytes(response, limit: int) -> bytes:
    if response is None:
        return b""

    try:
        if isinstance(response, (bytes, bytearray)):
            return bytes(response[:limit])

        body = getattr(response, "body", None)
        if body is not None:
            return bytes(body[:limit])

        iterator = getattr(response, "body_iterator", None)
        if iterator is not None:
            data = bytearray()
            async for chunk in iterator:
                if not chunk:
                    continue
                need = limit - len(data)
                if need <= 0:
                    break
                data.extend(chunk[:need])
                if len(data) >= limit:
                    break
            return bytes(data)

        if hasattr(response, "__aiter__"):
            data = bytearray()
            async for chunk in response:
                if not chunk:
                    continue
                need = limit - len(data)
                if need <= 0:
                    break
                data.extend(chunk[:need])
                if len(data) >= limit:
                    break
            return bytes(data)
    finally:
        close_func = getattr(response, "close", None)
        if callable(close_func):
            result = close_func()
            if inspect.isawaitable(result):
                await result

    return b""


async def _read_range_slice(adapter, root: str, rel: str, start: int, end: int) -> bytes:
    read_range = getattr(adapter, "read_file_range", None)
    if callable(read_range):
        try:
            return await read_range(root, rel, start, end)
        except TypeError:
            return await read_range(root, rel, start, end=end)

    stream_impl = getattr(adapter, "stream_file", None)
    if callable(stream_impl):
        range_header = f"bytes={start}-{end}"
        response = await stream_impl(root, rel, range_header)
        expected = end - start + 1
        return await _collect_response_bytes(response, expected)

    read_file = getattr(adapter, "read_file", None)
    if callable(read_file) and start == 0:
        data = await read_file(root, rel)
        slice_end = end + 1
        return data[:slice_end]

    return b""


async def _read_video_prefix(adapter, root: str, rel: str, size: int, limit: int = VIDEO_RANGE_LIMIT) -> bytes:
    chunk_size = min(VIDEO_INITIAL_CHUNK, limit)
    offset = 0
    collected = bytearray()

    while len(collected) < limit:
        end = offset + chunk_size - 1
        data = await _read_range_slice(adapter, root, rel, offset, end)
        if not data:
            break
        collected.extend(data)
        if len(data) < chunk_size:
            break
        offset += len(data)
        remaining = limit - len(collected)
        if remaining <= 0:
            break
        chunk_size = min(chunk_size * 2, remaining)

    if not collected and size <= limit:
        read_file = getattr(adapter, "read_file", None)
        if callable(read_file):
            blob = await read_file(root, rel)
            if blob:
                return bytes(blob[:limit])

    return bytes(collected[:limit])


async def _run_ffmpeg_extract_frame(src_path: str, dst_path: str):
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel", "error",
        "-i", src_path,
        "-frames:v", "1",
        dst_path,
    ]
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
    except FileNotFoundError as e:
        raise RuntimeError("未找到 ffmpeg，可执行文件需要在 PATH 中") from e

    stdout, stderr = await proc.communicate()
    if proc.returncode != 0:
        message = stderr.decode().strip() or stdout.decode().strip() or "ffmpeg 执行失败"
        raise RuntimeError(message)


async def _generate_video_thumb(video_bytes: bytes, rel: str, w: int, h: int, fit: str) -> Tuple[bytes, str]:
    from PIL import Image

    suffix = Path(rel).suffix or ".mp4"
    src_tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    src_path = src_tmp.name
    try:
        src_tmp.write(video_bytes)
        src_tmp.flush()
    finally:
        src_tmp.close()

    dst_tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    dst_path = dst_tmp.name
    dst_tmp.close()

    try:
        await _run_ffmpeg_extract_frame(src_path, dst_path)
        with Image.open(dst_path) as im:
            im.load()
            return _image_to_webp(im, w, h, fit)
    finally:
        with suppress(FileNotFoundError):
            Path(src_path).unlink()
        with suppress(FileNotFoundError):
            Path(dst_path).unlink()


async def get_or_create_thumb(adapter, adapter_id: int, root: str, rel: str, w: int, h: int, fit: str = 'cover'):
    stat = await adapter.stat_file(root, rel)
    size = int(stat.get('size') or 0)
    is_video = is_video_filename(rel)
    if not is_video and size > MAX_IMAGE_SOURCE_SIZE:
        raise HTTPException(400, detail="Image too large for thumbnail")

    key = _cache_key(adapter_id, rel, size, int(
        stat.get('mtime', 0)), w, h, fit)
    path = _cache_path(key)
    if path.exists():
        return path.read_bytes(), 'image/webp', key

    _ensure_cache_dir(path)
    thumb_bytes, mime = None, None

    get_thumb_impl = getattr(adapter, "get_thumbnail", None)
    if callable(get_thumb_impl):
        size_str = "large" if w > 400 else "medium" if w > 100 else "small"
        native_thumb_bytes = await get_thumb_impl(root, rel, size_str)

        if native_thumb_bytes:
            try:
                from PIL import Image
                im = Image.open(io.BytesIO(native_thumb_bytes))
                buf = io.BytesIO()
                im.save(buf, 'WEBP', quality=85)
                thumb_bytes = buf.getvalue()
                mime = 'image/webp'
            except Exception as e:
                print(
                    f"Failed to convert native thumbnail to WebP: {e}, falling back.")
                thumb_bytes, mime = None, None

    if not thumb_bytes:
        if is_video:
            try:
                video_bytes = await _read_video_prefix(adapter, root, rel, size)
            except HTTPException:
                raise
            except Exception as e:
                print(f"Video prefix read failed: {e}")
                raise HTTPException(500, detail=f"Video read failed: {e}")

            if not video_bytes:
                raise HTTPException(500, detail="Unable to read video data for thumbnail")

            try:
                thumb_bytes, mime = await _generate_video_thumb(video_bytes, rel, w, h, fit)
            except Exception as e:
                print(f"Video thumbnail generation failed: {e}")
                raise HTTPException(
                    500, detail=f"Video thumbnail generation failed: {e}")
        else:
            read_data = await adapter.read_file(root, rel)
            try:
                thumb_bytes, mime = generate_thumb(
                    read_data, w, h, fit, is_raw=is_raw_filename(rel))
            except Exception as e:
                print(e)
                raise HTTPException(
                    500, detail=f"Thumbnail generation failed: {e}")

    if thumb_bytes:
        path.write_bytes(thumb_bytes)
        return thumb_bytes, mime, key

    raise HTTPException(
        500, detail="Failed to generate thumbnail by any means")