cmshoppe/app/image_studio_images.py

"""Remote image helpers for the AI image studio."""

from __future__ import annotations

import hashlib
import io
import ipaddress
import os
import socket
import threading
import urllib.parse
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Callable, Optional

import requests

from . import appconfig, image_studio


THUMBNAIL_MAX_BYTES = 3 * 1024 * 1024
ORIGINAL_MAX_BYTES = 12 * 1024 * 1024
DEFAULT_CONNECT_TIMEOUT_SECONDS = 5
DEFAULT_READ_TIMEOUT_SECONDS = 30
DEFAULT_THUMBNAIL_SIZE = 220
DEFAULT_THUMBNAIL_WORKERS = 4
MAX_ORIGINAL_ASSETS = 16


class ImageStudioImageError(RuntimeError):
    """Raised when a remote image cannot be safely loaded or saved."""


class ImageStudioImageCancelled(ImageStudioImageError):
    """Raised when an original image download is cancelled safely."""


@dataclass(frozen=True)
class RemoteImage:
    url: str
    content: bytes
    content_type: str
    final_url: str
    redirected: bool


@dataclass(frozen=True)
class ThumbnailResult:
    key: str
    url: str
    image_bytes: bytes
    width: int
    height: int
    from_cache: bool = False


def _session():
    session = requests.Session()
    session.trust_env = False
    return session


def _assert_public_http_url(url):
    parts = urllib.parse.urlsplit(str(url or ""))
    if parts.scheme not in {"http", "https"}:
        raise ImageStudioImageError("远程图片地址只允许 http/https")
    host = parts.hostname
    if not host:
        raise ImageStudioImageError("远程图片地址缺少域名")
    if _is_local_hostname(host):
        raise ImageStudioImageError("远程图片地址不能指向本机或内网")
    try:
        _assert_public_ip(host)
        return
    except ValueError:
        pass
    try:
        addresses = socket.getaddrinfo(
            host,
            parts.port or (443 if parts.scheme == "https" else 80),
            type=socket.SOCK_STREAM,
        )
    except OSError as exc:
        raise ImageStudioImageError(f"远程图片地址无法解析：{exc}") from exc
    if not addresses:
        raise ImageStudioImageError("远程图片地址无法解析")
    for address in addresses:
        _assert_public_ip(address[4][0])


def _is_local_hostname(host):
    lowered = str(host or "").strip().lower().rstrip(".")
    return lowered in {"localhost"} or lowered.endswith(".localhost") or lowered.endswith(".local")


def _assert_public_ip(value):
    ip = ipaddress.ip_address(value)
    if (
        ip.is_private
        or ip.is_loopback
        or ip.is_link_local
        or ip.is_multicast
        or ip.is_reserved
        or ip.is_unspecified
    ):
        raise ImageStudioImageError("远程图片地址不能指向本机或内网")


def _safe_headers(referer=None):
    headers = {
        "User-Agent": "Mozilla/5.0",
        "Accept": "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8",
    }
    if referer:
        headers["Referer"] = str(referer)
    return headers


def _validate_content_type(value):
    content_type = str(value or "").split(";", 1)[0].strip().lower()
    if content_type and not content_type.startswith("image/"):
        raise ImageStudioImageError("远程地址返回的不是图片内容")
    return content_type


def download_remote_image(
    url,
    *,
    referer=None,
    max_bytes=ORIGINAL_MAX_BYTES,
    timeout=(DEFAULT_CONNECT_TIMEOUT_SECONDS, DEFAULT_READ_TIMEOUT_SECONDS),
    session=None,
    should_stop=None,
) -> RemoteImage:
    """Download remote image bytes with SSRF, timeout and size guards."""

    url = str(url or "").strip()
    should_stop = should_stop or (lambda: False)
    if should_stop():
        raise ImageStudioImageCancelled("用户停止下载商品原图")
    _assert_public_http_url(url)
    client = session or _session()
    try:
        response = client.get(
            url,
            stream=True,
            timeout=timeout,
            headers=_safe_headers(referer),
        )
    except requests.exceptions.RequestException as exc:
        raise ImageStudioImageError(f"远程图片下载失败：{exc}") from exc
    try:
        if should_stop():
            raise ImageStudioImageCancelled("用户停止下载商品原图")
        status = getattr(response, "status_code", 200)
        if status >= 400:
            raise ImageStudioImageError(f"远程图片下载失败：HTTP {status}")
        final_url = str(getattr(response, "url", "") or url)
        _assert_public_http_url(final_url)
        content_type = _validate_content_type(
            getattr(response, "headers", {}).get("Content-Type", "")
        )
        content_length = getattr(response, "headers", {}).get("Content-Length")
        if content_length:
            try:
                if int(content_length) > int(max_bytes):
                    raise ImageStudioImageError("远程图片超过大小上限")
            except ValueError:
                pass
        chunks = []
        total = 0
        iterator = (
            response.iter_content(chunk_size=65536)
            if hasattr(response, "iter_content")
            else [response.content]
        )
        for chunk in iterator:
            if should_stop():
                raise ImageStudioImageCancelled("用户停止下载商品原图")
            if not chunk:
                continue
            total += len(chunk)
            if total > int(max_bytes):
                raise ImageStudioImageError("远程图片超过大小上限")
            chunks.append(chunk)
        if should_stop():
            raise ImageStudioImageCancelled("用户停止下载商品原图")
        content = b"".join(chunks)
        _image_info(content)
        return RemoteImage(
            url=url,
            content=content,
            content_type=content_type,
            final_url=final_url,
            redirected=final_url != url,
        )
    finally:
        close = getattr(response, "close", None)
        if callable(close):
            close()


def _image_info(image_bytes):
    try:
        from PIL import Image

        with Image.open(io.BytesIO(image_bytes)) as image:
            image.verify()
        with Image.open(io.BytesIO(image_bytes)) as image:
            return {
                "format": str(image.format or "JPEG").upper(),
                "width": int(image.width),
                "height": int(image.height),
            }
    except Exception as exc:
        raise ImageStudioImageError(f"图片解码失败：{exc}") from exc


def _thumbnail_bytes(image_bytes, max_size=DEFAULT_THUMBNAIL_SIZE):
    try:
        from PIL import Image

        with Image.open(io.BytesIO(image_bytes)) as image:
            image.thumbnail((int(max_size), int(max_size)))
            output = io.BytesIO()
            image.convert("RGBA").save(output, format="PNG")
            return output.getvalue(), image.width, image.height
    except Exception as exc:
        raise ImageStudioImageError(f"缩略图解码失败：{exc}") from exc


def load_thumbnail(url, *, key=None, max_size=DEFAULT_THUMBNAIL_SIZE, referer=None, session=None):
    remote = download_remote_image(
        url,
        referer=referer,
        max_bytes=THUMBNAIL_MAX_BYTES,
        session=session,
    )
    image_bytes, width, height = _thumbnail_bytes(remote.content, max_size=max_size)
    return ThumbnailResult(
        key=str(key or url),
        url=remote.url,
        image_bytes=image_bytes,
        width=width,
        height=height,
        from_cache=False,
    )


class ThumbnailLoader:
    """Small thread-pool thumbnail loader with memory cache and cooperative cancel."""

    def __init__(self, max_workers=DEFAULT_THUMBNAIL_WORKERS, loader: Optional[Callable] = None):
        self._executor = ThreadPoolExecutor(max_workers=max(1, int(max_workers or 1)))
        self._loader = loader or load_thumbnail
        self._cache = {}
        self._cancelled = set()
        self._closed = False
        self._lock = threading.Lock()

    def submit(self, key, url, on_success=None, on_error=None):
        key = str(key)
        url = str(url or "").strip()
        with self._lock:
            if self._closed:
                raise ImageStudioImageError("缩略图加载器已关闭")
            cached = self._cache.get(url)
            if cached is not None:
                result = ThumbnailResult(key, url, cached.image_bytes, cached.width, cached.height, from_cache=True)
                if on_success is not None:
                    on_success(result)
                return None
            self._cancelled.discard(key)
            future = self._executor.submit(self._loader, url, key=key)

        def done_callback(done):
            try:
                result = done.result()
                with self._lock:
                    if key in self._cancelled or self._closed:
                        return
                    self._cache[url] = result
                if on_success is not None:
                    on_success(result)
            except Exception as exc:
                with self._lock:
                    if key in self._cancelled or self._closed:
                        return
                if on_error is not None:
                    on_error(key, exc)

        future.add_done_callback(done_callback)
        return future

    def cancel(self, key=None):
        with self._lock:
            if key is None:
                self._cancelled.add("*")
                self._closed = True
            else:
                self._cancelled.add(str(key))

    def close(self):
        with self._lock:
            self._closed = True
        self._executor.shutdown(wait=False, cancel_futures=True)


def _extension_for_format(format_name):
    normalized = str(format_name or "").strip().upper()
    if normalized in {"JPEG", "JPG"}:
        return ".jpg"
    if normalized == "PNG":
        return ".png"
    if normalized == "WEBP":
        return ".webp"
    return ".img"


def _original_file_path(project, asset, image_root):
    dirs = image_studio.project_image_dirs(image_root, project)
    digest = hashlib.sha1(str(asset.remote_url or asset.id).encode("utf-8")).hexdigest()[:10]
    stem = "original_%02d_%s" % (max(1, int(asset.source_order or 1)), digest)
    return dirs["originals"], stem


def _existing_local_asset(asset):
    local_path = str(asset.local_path or "").strip()
    if not local_path or not os.path.isfile(local_path):
        return None
    try:
        with open(local_path, "rb") as fh:
            _image_info(fh.read())
        return asset
    except ImageStudioImageError:
        return None


def import_original_files(
    project_id,
    file_paths,
    *,
    path=None,
    config=None,
    image_root=None,
    max_assets=MAX_ORIGINAL_ASSETS,
):
    """Validate and atomically copy local product images into one studio project."""

    imported = []
    errors = []
    for file_path in file_paths or []:
        try:
            with open(os.path.abspath(str(file_path)), "rb") as fh:
                content = fh.read(int(ORIGINAL_MAX_BYTES) + 1)
            imported.append(
                import_original_bytes(
                    project_id,
                    content,
                    filename_hint=os.path.basename(str(file_path)),
                    path=path,
                    config=config,
                    image_root=image_root,
                    max_assets=max_assets,
                )
            )
        except Exception as exc:
            errors.append({"path": os.path.abspath(str(file_path)), "error": str(exc)})
    return {"assets": imported, "errors": errors, "limit": int(max_assets)}


def import_original_bytes(
    project_id,
    content,
    *,
    filename_hint="clipboard.png",
    path=None,
    config=None,
    image_root=None,
    max_assets=MAX_ORIGINAL_ASSETS,
):
    cfg = appconfig.load_config() if config is None else config
    database_path = path or appconfig.db_path(cfg)
    project = image_studio.get_project(project_id, path=database_path)
    if project is None:
        raise ImageStudioImageError("商品套图任务不存在")
    image_bytes = bytes(content or b"")
    if not image_bytes:
        raise ImageStudioImageError("商品原图内容为空")
    if len(image_bytes) > int(ORIGINAL_MAX_BYTES):
        raise ImageStudioImageError("商品原图超过大小上限")
    info = _image_info(image_bytes)
    originals = image_studio.list_assets(
        project.id,
        kind=image_studio.ASSET_KIND_ORIGINAL,
        path=database_path,
    )
    active_originals = [
        asset for asset in originals if asset.status != image_studio.ASSET_STATUS_MISSING
    ]
    digest = hashlib.sha256(image_bytes).hexdigest()[:16]
    duplicate = next(
        (
            asset
            for asset in originals
            if digest in os.path.basename(str(asset.local_path or ""))
            and _existing_local_asset(asset) is not None
        ),
        None,
    )
    if duplicate is not None:
        return duplicate
    if len(active_originals) >= int(max_assets):
        raise ImageStudioImageError("商品原图最多只能添加%d张" % int(max_assets))

    directory = image_studio.project_image_dirs(
        image_root or appconfig.image_dir(cfg),
        project,
    )["originals"]
    os.makedirs(directory, exist_ok=True)
    order = max([int(asset.source_order or 0) for asset in originals] + [0]) + 1
    extension = _extension_for_format(info["format"])
    safe_hint = os.path.splitext(os.path.basename(str(filename_hint or "image")))[0]
    safe_hint = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in safe_hint)
    safe_hint = safe_hint.strip("_")[:32] or "image"
    final_path = os.path.join(
        directory,
        "original_%02d_%s_%s%s" % (order, safe_hint, digest, extension),
    )
    temp_path = final_path + ".tmp-" + uuid.uuid4().hex
    created = False
    try:
        with open(temp_path, "wb") as fh:
            fh.write(image_bytes)
        with open(temp_path, "rb") as fh:
            _image_info(fh.read())
        os.replace(temp_path, final_path)
        created = True
        return image_studio.add_asset(
            project.id,
            image_studio.ASSET_KIND_ORIGINAL,
            local_path=final_path,
            aspect_ratio="%d:%d" % (int(info["width"]), int(info["height"])),
            source_order=order,
            path=database_path,
        )
    except Exception as exc:
        for candidate in (temp_path, final_path if created else None):
            if candidate and os.path.exists(candidate):
                try:
                    os.remove(candidate)
                except OSError:
                    pass
        if isinstance(exc, ImageStudioImageError):
            raise
        raise ImageStudioImageError(f"保存商品原图失败：{exc}") from exc


def trash_generated_asset(asset_id, *, path=None, config=None, image_root=None):
    """Move one generated image to the app-managed trash without losing DB history."""

    cfg = appconfig.load_config() if config is None else config
    database_path = path or appconfig.db_path(cfg)
    asset = image_studio.get_asset(asset_id, path=database_path)
    if asset is None or not str(asset.kind or "").startswith("generated_"):
        raise ImageStudioImageError("只能删除商品套图生成结果")
    source_path = os.path.abspath(str(asset.local_path or ""))
    if not source_path or not os.path.isfile(source_path):
        raise ImageStudioImageError("生成图片文件不存在")
    project = image_studio.get_project(asset.project_id, path=database_path)
    if project is None:
        raise ImageStudioImageError("商品套图任务不存在")
    root = image_studio.project_image_dirs(
        image_root or appconfig.image_dir(cfg),
        project,
    )["root"]
    trash_dir = os.path.join(root, ".trash")
    os.makedirs(trash_dir, exist_ok=True)
    trash_path = os.path.join(
        trash_dir,
        "%s_%s" % (uuid.uuid4().hex, os.path.basename(source_path)),
    )
    try:
        os.replace(source_path, trash_path)
        image_studio.update_asset_local_path(
            asset.id,
            trash_path,
            status=image_studio.ASSET_STATUS_MISSING,
            path=database_path,
        )
    except Exception as exc:
        if os.path.isfile(trash_path) and not os.path.exists(source_path):
            try:
                os.replace(trash_path, source_path)
            except OSError:
                pass
        raise ImageStudioImageError(f"删除生成图片失败：{exc}") from exc
    return {
        "asset_id": asset.id,
        "original_path": source_path,
        "trash_path": trash_path,
    }


def restore_trashed_asset(record, *, path=None, config=None):
    cfg = appconfig.load_config() if config is None else config
    database_path = path or appconfig.db_path(cfg)
    payload = dict(record or {})
    asset = image_studio.get_asset(payload.get("asset_id"), path=database_path)
    if asset is None:
        raise ImageStudioImageError("待撤销的生成图片记录不存在")
    trash_path = os.path.abspath(str(payload.get("trash_path") or ""))
    original_path = os.path.abspath(str(payload.get("original_path") or ""))
    if not os.path.isfile(trash_path):
        raise ImageStudioImageError("废纸篓中的生成图片不存在")
    if os.path.exists(original_path):
        stem, extension = os.path.splitext(original_path)
        original_path = "%s_restored_%s%s" % (stem, uuid.uuid4().hex[:8], extension)
    os.makedirs(os.path.dirname(original_path), exist_ok=True)
    os.replace(trash_path, original_path)
    try:
        return image_studio.update_asset_local_path(
            asset.id,
            original_path,
            status=image_studio.ASSET_STATUS_AVAILABLE,
            path=database_path,
        )
    except Exception as exc:
        try:
            os.replace(original_path, trash_path)
        except OSError:
            pass
        raise ImageStudioImageError(f"撤销删除生成图片失败：{exc}") from exc


def download_original_asset(
    asset_id,
    *,
    path=None,
    config=None,
    image_root=None,
    session=None,
    should_stop=None,
):
    """Download one Shopee original image into originals/ and mark its asset available."""

    cfg = appconfig.load_config() if config is None else config
    should_stop = should_stop or (lambda: False)
    if should_stop():
        raise ImageStudioImageCancelled("用户停止下载商品原图")
    database_path = path or appconfig.db_path(cfg)
    asset = image_studio.get_asset(asset_id, path=database_path)
    if asset is None:
        raise ImageStudioImageError("原图资产不存在")
    if not asset.remote_url:
        raise ImageStudioImageError("原图资产缺少远程 URL")
    existing = _existing_local_asset(asset)
    if existing is not None:
        return existing
    project = image_studio.get_project(asset.project_id, path=database_path)
    if project is None:
        raise ImageStudioImageError("原图资产所属项目不存在")
    remote = download_remote_image(
        asset.remote_url,
        max_bytes=ORIGINAL_MAX_BYTES,
        session=session,
        should_stop=should_stop,
    )
    if should_stop():
        raise ImageStudioImageCancelled("用户停止下载商品原图")
    info = _image_info(remote.content)
    directory, stem = _original_file_path(project, asset, image_root or appconfig.image_dir(cfg))
    os.makedirs(directory, exist_ok=True)
    final_path = os.path.join(directory, stem + _extension_for_format(info["format"]))
    temp_path = final_path + ".tmp-" + uuid.uuid4().hex
    try:
        if should_stop():
            raise ImageStudioImageCancelled("用户停止下载商品原图")
        with open(temp_path, "wb") as fh:
            fh.write(remote.content)
        if should_stop():
            raise ImageStudioImageCancelled("用户停止下载商品原图")
        with open(temp_path, "rb") as fh:
            _image_info(fh.read())
        os.replace(temp_path, final_path)
    except Exception as exc:
        try:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        finally:
            if isinstance(exc, ImageStudioImageError):
                raise
            raise ImageStudioImageError(f"保存原图失败：{exc}") from exc
    return image_studio.update_asset_local_path(
        asset.id,
        final_path,
        status=image_studio.ASSET_STATUS_AVAILABLE,
        path=database_path,
    )