"""Bounded reference thumbnails for the unified image-generator panel."""

from __future__ import annotations

import asyncio
from io import BytesIO
import os
from typing import Any

from PIL import Image, ImageOps


THUMBNAIL_MAX_SIZE = 256
THUMBNAIL_QUALITY = 82
THUMBNAIL_CACHE_SECONDS = 3600
_ALLOWED_FOLDER_TYPES = frozenset({"input", "output", "temp"})


def _is_within(root: str, candidate: str) -> bool:
    try:
        return os.path.commonpath(
            [os.path.normcase(root), os.path.normcase(candidate)]
        ) == os.path.normcase(root)
    except ValueError:
        return False


def resolve_thumbnail_source(
    folder_paths_module: Any,
    filename: str,
    subfolder: str = "",
    folder_type: str = "input",
) -> str:
    """Resolve one public ComfyUI image descriptor without allowing traversal."""

    folder_type = str(folder_type or "input").strip()
    if folder_type not in _ALLOWED_FOLDER_TYPES:
        raise ValueError("图片目录类型无效")

    filename = str(filename or "").strip()
    subfolder = str(subfolder or "").strip()
    if not filename or os.path.basename(filename) != filename:
        raise ValueError("图片文件名无效")
    if os.path.isabs(subfolder) or ".." in subfolder.replace("\\", "/").split("/"):
        raise ValueError("图片子目录无效")

    root_value = folder_paths_module.get_directory_by_type(folder_type)
    if not root_value:
        raise ValueError("图片目录不存在")
    root = os.path.realpath(os.path.abspath(root_value))
    source = os.path.realpath(os.path.abspath(os.path.join(root, subfolder, filename)))
    if not _is_within(root, source):
        raise ValueError("图片路径不安全")
    if not os.path.isfile(source):
        raise FileNotFoundError(filename)
    return source


def render_reference_thumbnail(
    source: str,
    max_size: int = THUMBNAIL_MAX_SIZE,
    quality: int = THUMBNAIL_QUALITY,
) -> bytes:
    """Render a small WebP without modifying or replacing the source image."""

    max_size = max(32, min(int(max_size), THUMBNAIL_MAX_SIZE))
    quality = max(1, min(int(quality), 100))
    with Image.open(source) as opened:
        if getattr(opened, "is_animated", False):
            opened.seek(0)
        # JPEG decoders can use draft mode to avoid materialising every source
        # pixel before the final thumbnail resize. Other formats still remain
        # bounded by the route-level semaphore.
        opened.draft("RGB", (max_size, max_size))
        image = None
        try:
            image = ImageOps.exif_transpose(opened)
            image.thumbnail(
                (max_size, max_size),
                Image.Resampling.LANCZOS,
                reducing_gap=3.0,
            )
            has_alpha = "A" in image.getbands() or "transparency" in image.info
            converted = image.convert("RGBA" if has_alpha else "RGB")
            try:
                output = BytesIO()
                converted.save(output, format="WEBP", quality=quality, method=4)
                return output.getvalue()
            finally:
                converted.close()
        finally:
            if image is not None and image is not opened:
                image.close()


def _thumbnail_etag(source: str) -> str:
    stat = os.stat(source)
    return f'"o1key-thumb-{stat.st_mtime_ns:x}-{stat.st_size:x}"'


def register_o1key_image_thumbnail_route(PromptServer, web, folder_paths_module):
    """Register the thumbnail endpoint used only by lightweight panel previews."""

    render_semaphore = asyncio.Semaphore(2)

    @PromptServer.instance.routes.get("/o1key/image/thumbnail")
    async def get_o1key_image_thumbnail(request):
        try:
            source = resolve_thumbnail_source(
                folder_paths_module,
                request.query.get("filename", ""),
                request.query.get("subfolder", ""),
                request.query.get("type", "input"),
            )
            etag = await asyncio.to_thread(_thumbnail_etag, source)
            headers = {
                "Cache-Control": f"private, max-age={THUMBNAIL_CACHE_SECONDS}",
                "ETag": etag,
                "X-Content-Type-Options": "nosniff",
            }
            if request.headers.get("If-None-Match") == etag:
                return web.Response(status=304, headers=headers)
            async with render_semaphore:
                body = await asyncio.to_thread(render_reference_thumbnail, source)
            return web.Response(body=body, content_type="image/webp", headers=headers)
        except FileNotFoundError:
            return web.Response(status=404)
        except ValueError:
            return web.Response(status=400)
        except Exception:
            # Do not expose local paths or decoder details to the browser.
            return web.Response(status=415)


__all__ = [
    "THUMBNAIL_MAX_SIZE",
    "register_o1key_image_thumbnail_route",
    "render_reference_thumbnail",
    "resolve_thumbnail_source",
]