Foxel/domain/virtual_fs/thumbnail.py

import asyncio
import inspect
import io
import hashlib
import subprocess
import tempfile
from contextlib import suppress
from pathlib import Path
from typing import Tuple

from PIL import Image
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_TAIL_LIMIT = 2 * 1024 * 1024  # 2MB
VIDEO_TAIL_FALLBACK_LIMIT = 4 * 1024 * 1024  # 4MB
VIDEO_HEAD_LIMIT = 2 * 1024 * 1024  # 2MB
VIDEO_HEAD_FALLBACK_LIMIT = 4 * 1024 * 1024  # 4MB
VIDEO_THUMB_SEEK_SECONDS = (15, 10, 5, 3, 1, 0)
VIDEO_BLACK_FRAME_MEAN_THRESHOLD = 12.0
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]:
    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 _load_image_with_pillow(data: bytes):
    im = Image.open(io.BytesIO(data))
    im.load()
    return im


def _load_raw_with_ffmpeg(data: bytes, filename: str | None) -> "Image.Image":
    src_path: str | None = None
    dst_path: str | None = None
    try:
        with tempfile.NamedTemporaryFile(suffix=Path(filename or "").suffix or ".raw", delete=False) as src_tmp:
            src_tmp.write(data)
            src_path = src_tmp.name
        dst_tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        dst_path = dst_tmp.name
        dst_tmp.close()

        cmd = [
            "ffmpeg",
            "-y",
            "-hide_banner",
            "-loglevel", "error",
            "-i", src_path,
            "-frames:v", "1",
            dst_path,
        ]
        try:
            result = subprocess.run(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
            )
        except FileNotFoundError as e:
            raise RuntimeError("未找到 ffmpeg，可执行文件需要在 PATH 中") from e
        except subprocess.CalledProcessError as e:
            stderr = (e.stderr or b"").decode().strip()
            stdout = (e.stdout or b"").decode().strip()
            message = stderr or stdout or "ffmpeg 转换 RAW 失败"
            raise RuntimeError(message) from e

        with open(dst_path, "rb") as f:
            img_bytes = f.read()
        im = Image.open(io.BytesIO(img_bytes))
        im.load()
        return im
    finally:
        if dst_path:
            with suppress(FileNotFoundError):
                Path(dst_path).unlink()
        if src_path:
            with suppress(FileNotFoundError):
                Path(src_path).unlink()


def load_image_from_bytes(data: bytes, *, filename: str | None = None, is_raw: bool = False):
    if not is_raw:
        return _load_image_with_pillow(data)

    first_error: Exception | None = None
    try:
        return _load_image_with_pillow(data)
    except Exception as exc:
        first_error = exc

    try:
        return _load_raw_with_ffmpeg(data, filename)
    except Exception as exc:
        msg = f"RAW 解码失败: ffmpeg 处理异常 {exc}"
        if first_error:
            msg = f"RAW 解码失败: Pillow 异常 {first_error}; ffmpeg 异常 {exc}"
        raise RuntimeError(msg) from exc


def raw_bytes_to_jpeg(data: bytes, filename: str | None = None) -> bytes:
    im = load_image_from_bytes(data, filename=filename, is_raw=True)
    if im.mode != "RGB":
        im = im.convert("RGB")
    buf = io.BytesIO()
    im.save(buf, "JPEG", quality=90)
    return buf.getvalue()


def generate_thumb(data: bytes, w: int, h: int, fit: str, is_raw: bool = False, filename: str | None = None) -> Tuple[bytes, str]:
    im = load_image_from_bytes(data, filename=filename, is_raw=is_raw)
    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_head(adapter, root: str, rel: str, size: int, limit: int = VIDEO_HEAD_LIMIT) -> bytes:
    end = limit - 1
    if size > 0:
        end = min(end, size - 1)
    if end < 0:
        return b""
    return await _read_range_slice(adapter, root, rel, 0, end)


async def _read_video_tail(adapter, root: str, rel: str, size: int, limit: int) -> Tuple[bytes, int]:
    if size <= 0:
        return b"", 0
    start = max(0, size - limit)
    end = size - 1
    data = await _read_range_slice(adapter, root, rel, start, end)
    return data, start


def _write_video_sparse_file(rel: str, head_bytes: bytes, tail_bytes: bytes, tail_offset: int) -> str:
    suffix = Path(rel).suffix or ".mp4"
    src_tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    src_path = src_tmp.name
    try:
        if head_bytes:
            src_tmp.write(head_bytes)
            src_tmp.flush()
    finally:
        src_tmp.close()

    if tail_bytes:
        with open(src_path, "r+b") as f:
            f.seek(max(0, int(tail_offset)))
            f.write(tail_bytes)
            f.flush()
    return src_path


async def _run_ffmpeg_extract_frame(src_path: str, dst_path: str, *, seek_seconds: float | None = None):
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel", "error",
    ]
    is_http_input = src_path.startswith(("http://", "https://"))
    if is_http_input and seek_seconds is not None:
        cmd += ["-ss", str(seek_seconds), "-i", src_path]
    else:
        cmd += ["-i", src_path]
        if seek_seconds is not None:
            cmd += ["-ss", str(seek_seconds)]
    cmd += [
        "-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)


def _frame_mean_luma(im) -> float:
    from PIL import ImageStat
    gray = im.convert("L").resize((64, 64))
    return float(ImageStat.Stat(gray).mean[0])


def _is_black_image_bytes(image_bytes: bytes) -> bool:
    from PIL import Image
    with Image.open(io.BytesIO(image_bytes)) as im:
        im.load()
        return _frame_mean_luma(im) < VIDEO_BLACK_FRAME_MEAN_THRESHOLD


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

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

    best: tuple[float, bytes, str] | None = None
    last_error: Exception | None = None
    try:
        for seek_seconds in VIDEO_THUMB_SEEK_SECONDS:
            try:
                with suppress(FileNotFoundError):
                    Path(dst_path).unlink()
                await _run_ffmpeg_extract_frame(src_path, dst_path, seek_seconds=seek_seconds)
                with Image.open(dst_path) as im:
                    im.load()
                    mean = _frame_mean_luma(im)
                    webp_bytes, mime = _image_to_webp(im, w, h, fit)

                if best is None or mean > best[0]:
                    best = (mean, webp_bytes, mime)
                if mean >= VIDEO_BLACK_FRAME_MEAN_THRESHOLD:
                    return webp_bytes, mime
            except Exception as e:
                last_error = e
                continue

        if best is not None:
            return best[1], best[2]
        if last_error is not None:
            raise last_error
        raise RuntimeError("ffmpeg 截帧失败")
    finally:
        with suppress(FileNotFoundError):
            Path(dst_path).unlink()


async def _generate_video_thumb_from_segments(
    head_bytes: bytes,
    tail_bytes: bytes,
    tail_offset: int,
    rel: str,
    w: int,
    h: int,
    fit: str,
) -> Tuple[bytes, str]:
    src_path = _write_video_sparse_file(rel, head_bytes, tail_bytes, tail_offset)
    try:
        return await _generate_video_thumb_from_src_path(src_path, w, h, fit)
    finally:
        with suppress(FileNotFoundError):
            Path(src_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:
            async def _maybe_transcoding_thumb() -> Tuple[bytes, str] | None:
                fid = (stat or {}).get("fid") if isinstance(stat, dict) else None
                get_url = getattr(adapter, "get_video_transcoding_url", None)
                if not fid or not callable(get_url):
                    return None
                try:
                    url = await get_url(str(fid))
                except Exception as e:
                    print(f"Video transcoding url fetch failed: {e}")
                    return None
                if not url:
                    return None
                try:
                    return await _generate_video_thumb_from_src_path(url, w, h, fit)
                except Exception as e:
                    print(f"Video transcoding thumbnail generation failed: {e}")
                    return None

            def _is_hevc_decoder_missing(exc: Exception) -> bool:
                msg = str(exc).lower()
                return ("no decoder found" in msg) and ("hevc" in msg or "h265" in msg)

            async def _read_head(limit: int) -> bytes:
                try:
                    return await _read_video_head(adapter, root, rel, size, limit=limit)
                except HTTPException:
                    raise
                except Exception as e:
                    print(f"Video head read failed: {e}")
                    raise HTTPException(500, detail=f"Video read failed: {e}")

            async def _read_tail(limit: int) -> Tuple[bytes, int]:
                try:
                    return await _read_video_tail(adapter, root, rel, size, limit=limit)
                except HTTPException:
                    raise
                except Exception as e:
                    print(f"Video tail read failed: {e}")
                    raise HTTPException(500, detail=f"Video read failed: {e}")

            head_bytes = await _read_head(VIDEO_HEAD_LIMIT)
            tail_bytes, tail_offset = await _read_tail(VIDEO_TAIL_LIMIT)
            if not head_bytes and not tail_bytes:
                raise HTTPException(500, detail="Unable to read video data for thumbnail")

            try:
                thumb_bytes, mime = await _generate_video_thumb_from_segments(
                    head_bytes, tail_bytes, tail_offset, rel, w, h, fit
                )
            except Exception as e1:
                if _is_hevc_decoder_missing(e1):
                    got = await _maybe_transcoding_thumb()
                    if got is not None:
                        thumb_bytes, mime = got
                if not thumb_bytes:
                    try:
                        tail_bytes, tail_offset = await _read_tail(VIDEO_TAIL_FALLBACK_LIMIT)
                        thumb_bytes, mime = await _generate_video_thumb_from_segments(
                            head_bytes, tail_bytes, tail_offset, rel, w, h, fit
                        )
                    except HTTPException:
                        raise
                    except Exception as e2:
                        print(f"Video thumbnail generation failed: {e2}")
                        raise HTTPException(500, detail=f"Video thumbnail generation failed: {e2}")

            if thumb_bytes and _is_black_image_bytes(thumb_bytes):
                try:
                    head_bytes = await _read_head(VIDEO_HEAD_FALLBACK_LIMIT)
                    retry_thumb, retry_mime = await _generate_video_thumb_from_segments(
                        head_bytes, tail_bytes, tail_offset, rel, w, h, fit
                    )
                    if retry_thumb and not _is_black_image_bytes(retry_thumb):
                        thumb_bytes, mime = retry_thumb, retry_mime
                except Exception:
                    pass
        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), 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")