ForcePilot/backend/package/yuxi/channel/message/media.py

from __future__ import annotations

import asyncio
import base64
import io
import logging
from dataclasses import dataclass
from email.message import EmailMessage
from pathlib import Path, PurePosixPath
from urllib.parse import urlparse

import httpx
from PIL import Image

from yuxi.channel.message.media_cleaner import MediaCleaner
from yuxi.channel.message.media_store import MediaStore
from yuxi.config import config

logger = logging.getLogger(__name__)

_media_store: MediaStore | None = None
_media_store_initialized: bool = False
_media_cleaner: MediaCleaner | None = None
_background_tasks: set[asyncio.Task[object]] = set()

MAX_IMAGE_BYTES = 6 * 1024 * 1024  # 6MB
_MAX_STORE_CONCURRENCY = 5
_store_semaphore = asyncio.Semaphore(_MAX_STORE_CONCURRENCY)

_SVG_MIME_TYPES = frozenset({"image/svg+xml", "image/svg"})
_IMAGE_MAGIC_BYTES: dict[str, tuple[bytes, ...]] = {
    "image/jpeg": (b"\xff\xd8\xff",),
    "image/png": (b"\x89PNG\r\n\x1a\n",),
    "image/gif": (b"GIF87a", b"GIF89a"),
    "image/webp": (b"RIFF",),
    "image/bmp": (b"BM",),
    "image/tiff": (b"II*\x00", b"MM\x00*"),
    "image/x-icon": (b"\x00\x00\x01\x00",),
}
_SVG_SCRIPT_MARKER = b"<script"

_media_client: httpx.AsyncClient | None = None


def _get_media_client() -> httpx.AsyncClient:
    global _media_client
    if _media_client is None:
        _media_client = httpx.AsyncClient(limits=httpx.Limits(max_connections=10, max_keepalive_connections=5))
    return _media_client


async def _cleanup_media_client() -> None:
    global _media_client
    if _media_client is not None:
        await _media_client.aclose()
        _media_client = None

_MAX_RETRIES = 2
_RETRY_DELAY = 1.0
_RETRYABLE_STATUSES = frozenset({408, 429, 500, 502, 503, 504})


@dataclass
class MediaDownloadResult:
    data: bytes
    content_type: str | None
    filename: str | None


def _parse_content_disposition_filename(header: str | None) -> str | None:
    if not header:
        return None
    msg = EmailMessage()
    msg["Content-Disposition"] = header
    filename = msg.get_filename()
    if filename:
        return filename
    return None


def _is_transient_error(exc: Exception) -> bool:
    if isinstance(exc, httpx.HTTPStatusError):
        return exc.response.status_code in _RETRYABLE_STATUSES
    if isinstance(exc, httpx.TimeoutException):
        return True
    if isinstance(exc, httpx.NetworkError):
        return True
    return False


def _resolve_filename(final_url: str, response: httpx.Response, file_path_hint: str | None) -> str | None:
    header_filename = _parse_content_disposition_filename(response.headers.get("content-disposition"))
    if header_filename:
        return header_filename
    if file_path_hint:
        return PurePosixPath(file_path_hint).name or None
    try:
        path = urlparse(final_url).path
        name = PurePosixPath(path).name
        return name or None
    except Exception:
        return None


async def _read_with_limit(response: httpx.Response, max_bytes: int) -> bytes:
    chunks: list[bytes] = []
    total = 0
    async for chunk in response.aiter_bytes(chunk_size=65536):
        total += len(chunk)
        if total > max_bytes:
            raise ValueError(f"Payload exceeds maxBytes {max_bytes}")
        chunks.append(chunk)
    return b"".join(chunks)


def _validate_image_magic_bytes(data: bytes, content_type: str | None) -> bool:
    ct = (content_type or "").split(";")[0].strip().lower()
    if ct in _SVG_MIME_TYPES:
        if data.lower().find(_SVG_SCRIPT_MARKER) != -1:
            logger.warning("SVG contains script tag, rejecting for XSS safety")
            return False
        return True
    expected_magics = _IMAGE_MAGIC_BYTES.get(ct)
    if expected_magics is None:
        return True
    return any(data.startswith(magic) for magic in expected_magics)


async def _fetch_media(
    url: str,
    *,
    max_bytes: int,
    timeout: float,
    file_path_hint: str | None = None,
) -> MediaDownloadResult | None:
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        logger.warning("Unsupported media URL scheme: %s", parsed.scheme)
        return None

    for attempt in range(_MAX_RETRIES + 1):
        try:
            client = _get_media_client()
            response = await client.get(url, follow_redirects=True, timeout=timeout)

            content_length_raw = response.headers.get("content-length")
            if content_length_raw:
                try:
                    content_length = int(content_length_raw)
                    if content_length > max_bytes:
                        logger.error(
                            "Content-Length %d exceeds maxBytes %d for url=%s",
                            content_length,
                            max_bytes,
                            url,
                        )
                        return None
                except ValueError:
                    pass

            response.raise_for_status()

            content_type = response.headers.get("content-type")
            data = await _read_with_limit(response, max_bytes)

            if not _validate_image_magic_bytes(data, content_type):
                logger.warning(
                    "Media magic bytes mismatch: content-type=%s, url=%s",
                    content_type,
                    url,
                )
                return None

            filename = _resolve_filename(str(response.url), response, file_path_hint)

            return MediaDownloadResult(
                data=data,
                content_type=content_type,
                filename=filename,
            )

        except (httpx.HTTPStatusError, httpx.TimeoutException, httpx.NetworkError) as e:
            if attempt < _MAX_RETRIES and _is_transient_error(e):
                delay = _RETRY_DELAY * (2**attempt)
                logger.warning(
                    "Transient error fetching media (attempt %d/%d), retrying in %.1fs: url=%s, err=%s",
                    attempt + 1,
                    _MAX_RETRIES + 1,
                    delay,
                    url,
                    e,
                )
                await asyncio.sleep(delay)
                continue
            if isinstance(e, httpx.HTTPStatusError):
                logger.error(
                    "HTTP error downloading media: url=%s, status=%d",
                    url,
                    e.response.status_code,
                )
            elif isinstance(e, httpx.TimeoutException):
                logger.error("Timeout downloading media: url=%s", url)
            else:
                logger.error("Network error downloading media: url=%s, err=%s", url, e)
            return None

        except Exception:
            logger.exception("Failed to download media: url=%s", url)
            return None

    logger.error("All retry attempts exhausted for media: url=%s", url)
    return None


async def download_image_to_base64(
    url: str,
    timeout: float = 30.0,
) -> str | None:
    result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
    if result is None:
        return None

    if not (result.content_type or "").startswith("image/"):
        logger.warning(
            "Non-image response: content-type=%s, url=%s",
            result.content_type,
            url,
        )
        return None

    _maybe_store(result.data, result.content_type, result.filename, url)

    return base64.b64encode(result.data).decode("utf-8")


async def resolve_image_base64(media_urls: list[str]) -> str | None:
    if not media_urls:
        return None
    for url in media_urls:
        result = await download_image_to_base64(url)
        if result:
            return result
    return None


def get_media_store() -> MediaStore | None:
    global _media_store, _media_store_initialized
    if _media_store_initialized:
        return _media_store
    _media_store_initialized = True
    if not config.media_store_enabled:
        return None
    media_dir = Path(config.save_dir) / "data" / "media"
    _media_store = MediaStore(
        base_dir=media_dir,
        max_total_bytes=config.media_store_max_total_bytes,
        max_file_bytes=config.media_store_max_file_bytes,
        ttl_seconds=config.media_store_ttl_seconds,
    )
    return _media_store


def get_media_cleaner() -> MediaCleaner | None:
    global _media_cleaner
    store = get_media_store()
    if store is None:
        return None
    if _media_cleaner is None:
        _media_cleaner = MediaCleaner(store, config.media_cleanup_interval_seconds)
    return _media_cleaner


async def start_media_cleaner() -> None:
    cleaner = get_media_cleaner()
    if cleaner is not None:
        await cleaner.start()


async def stop_media_cleaner() -> None:
    cleaner = get_media_cleaner()
    if cleaner is not None:
        await cleaner.stop()
    await _cleanup_media_client()


async def shutdown_media() -> None:
    await stop_media_cleaner()
    await _cancel_background_tasks()
    await _cleanup_media_client()


async def _cancel_background_tasks() -> None:
    _cleanup_completed_tasks()
    pending = list(_background_tasks)
    if not pending:
        return
    for t in pending:
        t.cancel()
    results = await asyncio.gather(*pending, return_exceptions=True)
    for r in results:
        if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError):
            logger.warning("Background store task error during shutdown: %s", r)
    _background_tasks.clear()


def _cleanup_completed_tasks() -> None:
    done = {t for t in _background_tasks if t.done()}
    _background_tasks.difference_update(done)
    if done:
        logger.debug("Cleaned up %d completed background store tasks, %d remaining", len(done), len(_background_tasks))


async def _store_with_semaphore(store: MediaStore, data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
    async with _store_semaphore:
        await store.store(data, content_type, filename, source_url)


def _maybe_store(data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
    store = get_media_store()
    if store is None:
        return
    _cleanup_completed_tasks()
    if len(_background_tasks) >= _MAX_STORE_CONCURRENCY * 2:
        logger.warning(
            "Background store tasks piling up: %d pending, concurrency limit=%d",
            len(_background_tasks),
            _MAX_STORE_CONCURRENCY,
        )
    task = asyncio.create_task(_store_with_semaphore(store, data, content_type, filename, source_url))
    task.add_done_callback(_handle_store_task_done)
    _background_tasks.add(task)


def _handle_store_task_done(task: asyncio.Task[object]) -> None:
    try:
        task.result()
    except asyncio.CancelledError:
        pass
    except Exception:
        logger.exception("Failed to store media in background task")
    finally:
        _background_tasks.discard(task)


@dataclass
class ImageProcessResult:
    original_file_id: str | None
    thumbnail_file_id: str | None
    width: int
    height: int
    format: str
    mime_type: str
    size_bytes: int


async def download_and_process_image(
    url: str,
    *,
    timeout: float = 30.0,
    thumbnail_size: tuple[int, int] | None = None,
    convert_to: str | None = None,
) -> ImageProcessResult | None:
    result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
    if result is None:
        return None

    if not (result.content_type or "").startswith("image/"):
        logger.warning(
            "Non-image response for download_and_process_image: content-type=%s, url=%s",
            result.content_type,
            url,
        )
        return None

    from yuxi.utils.image_processor import image_processor as img_proc

    image_data = result.data
    if convert_to and convert_to in img_proc.CONVERTIBLE_FORMATS:
        try:
            image_data = img_proc.convert_format(result.data, convert_to)
        except Exception:
            logger.exception("Failed to convert image format to %s for url=%s", convert_to, url)

    store = get_media_store()
    original_file_id: str | None = None
    thumbnail_file_id: str | None = None

    if store is not None:
        original_file_id = await store.store(
            image_data,
            content_type=result.content_type,
            filename=result.filename,
            source_url=url,
        )

    try:
        thumbnail_data = img_proc.generate_thumbnail(image_data, size=thumbnail_size)
    except Exception:
        logger.exception("Failed to generate thumbnail for url=%s", url)
        thumbnail_data = None

    if thumbnail_data and store is not None:
        thumbnail_file_id = await store.store(
            thumbnail_data,
            content_type="image/jpeg",
            filename=None,
            source_url=url,
        )

    try:
        with io.BytesIO(image_data) as buf:
            with Image.open(buf) as img:
                width, height = img.size
                fmt = img.format or "JPEG"
    except Exception:
        width, height = 0, 0
        fmt = "JPEG"

    return ImageProcessResult(
        original_file_id=original_file_id,
        thumbnail_file_id=thumbnail_file_id,
        width=width,
        height=height,
        format=fmt,
        mime_type=f"image/{fmt.lower()}",
        size_bytes=len(image_data),
    )