app/services/media_conversion_service.py

"""Media-to-Media conversion service.

Converts between media formats without blocking the FastAPI event loop:

  * PDF -> images (JPEG / PNG / WEBP) using pypdfium2 + Pillow.
  * Image -> image (JPEG / PNG / WEBP / BMP / GIF / TIFF) using Pillow.

All CPU-bound pixel work runs in the shared thread pool
(:mod:`app.core.thread_pool`), mirroring the architecture of the document
converter and the reference PDF-conversion service.

Output files are written to a per-job directory and then either:

  * returned as data URLs (``SUPABASE_UPLOAD_ENABLED=false``),
    or
  * uploaded to Supabase Storage with 24-hour signed URLs
    (``SUPABASE_UPLOAD_ENABLED=true``) and a warning stating the expiry.

If a Supabase upload fails, the file falls back to a data URL.
"""

from __future__ import annotations

import asyncio
import base64
import io
import re
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from app.config import get_settings
from app.core.logger import get_logger
from app.core.thread_pool import thread_pool as _thread_pool
from app.models.schemas import (
    ImageConversionParams,
    MediaConversionData,
    MediaOutputFile,
    MediaUploadSummary,
    PDFConversionParams,
)

_logger = get_logger(__name__)
_settings = get_settings()

_PAGE_SPEC_RE = re.compile(r"^\s*(\d+(-\d+)?)(\s*,\s*(\d+(-\d+)?))*\s*$")

_PDF_OUTPUT_FORMATS = frozenset({"JPEG", "PNG", "WEBP"})
_IMAGE_OUTPUT_FORMATS = frozenset({"JPEG", "PNG", "WEBP", "BMP", "GIF", "TIFF"})

_EXT_BY_FORMAT: Dict[str, str] = {
    "JPEG": "jpg", "PNG": "png", "WEBP": "webp",
    "BMP": "bmp", "GIF": "gif", "TIFF": "tiff",
}
_MIME_BY_FORMAT: Dict[str, str] = {
    "JPEG": "image/jpeg", "PNG": "image/png", "WEBP": "image/webp",
    "BMP": "image/bmp", "GIF": "image/gif", "TIFF": "image/tiff",
}


class MediaConversionError(Exception):
    """Raised for invalid input or failed conversions, mapped to HTTP errors."""

    def __init__(self, message: str, status_code: int = 400) -> None:
        super().__init__(message)
        self.message = message
        self.status_code = status_code


def _fmt_str(fmt) -> str:
    return fmt.value if hasattr(fmt, "value") else str(fmt)


def _ext_for(fmt: str) -> str:
    return _EXT_BY_FORMAT.get(fmt, "bin")


def _save_kwargs(fmt: str, quality: int) -> Dict[str, Any]:
    if fmt == "JPEG":
        return {"quality": quality, "optimize": True}
    if fmt == "PNG":
        return {"optimize": True}
    if fmt == "WEBP":
        return {"quality": quality, "method": 4}
    if fmt == "TIFF":
        return {"compression": "tiff_lzw"}
    return {}


def _normalise_for_jpeg(img):
    """Return an RGB image suitable for JPEG, flattening alpha onto white."""
    from PIL import Image

    mode = img.mode
    if mode == "RGB":
        return img
    if mode in ("RGBA", "LA", "P"):
        if mode == "P":
            img = img.convert("RGBA")
        bg = Image.new("RGB", img.size, (255, 255, 255))
        mask = img.split()[-1] if mode in ("RGBA", "LA") else None
        bg.paste(img, mask=mask)
        return bg
    return img.convert("RGB")


def _normalise_for_png(img):
    """Return an image whose mode PIL accepts for PNG / WEBP."""
    if img.mode in ("RGB", "RGBA", "L", "LA"):
        return img
    if img.mode == "P":
        return img.convert("RGBA")
    return img.convert("RGB")


def _parse_pages(spec: Optional[str], total_pages: int) -> List[int]:
    """Expand a page spec ('1', '1-3', '1,3,5-7') into 0-indexed page indices."""
    if not spec or not spec.strip():
        return list(range(total_pages))
    indices: set[int] = set()
    for token in str(spec).split(","):
        token = token.strip()
        if "-" in token:
            start, end = (int(x) for x in token.split("-", 1))
            indices.update(range(start - 1, end))
        else:
            indices.add(int(token) - 1)
    valid = sorted(i for i in indices if 0 <= i < total_pages)
    if not valid:
        raise MediaConversionError(
            f"Page spec '{spec}' contains no pages within the document's {total_pages} page(s).",
            status_code=422,
        )
    return valid


def _guard_memory(page_count: int, dpi: int) -> None:
    bytes_per_page = (dpi * 8.5) * (dpi * 11) * 3
    estimated_mb = (bytes_per_page * page_count) / (1024 * 1024)
    if estimated_mb > _settings.media_max_memory_mb:
        raise MediaConversionError(
            f"Requested conversion would need ~{estimated_mb:.0f} MB of memory "
            f"({page_count} pages at {dpi} DPI). Reduce DPI or select fewer pages.",
            status_code=422,
        )


# ---------------------------------------------------------------------------
# Thread-pool worker functions (must stay top-level for pickling)
# ---------------------------------------------------------------------------

def _pdf_total_pages(data: bytes) -> int:
    import pypdfium2 as pdfium

    try:
        doc = pdfium.PdfDocument(data)
    except Exception as exc:
        raise MediaConversionError(
            f"Input is not a valid PDF: {exc}", status_code=422
        ) from exc
    try:
        total = len(doc)
    finally:
        doc.close()
    if total == 0:
        raise MediaConversionError("PDF contains no pages.", status_code=422)
    return total


def _split_pdf_pages(data: bytes, page_indices: List[int]) -> Dict[int, bytes]:
    """Extract each requested page into its own single-page PDF blob."""
    import pypdfium2 as pdfium

    blobs: Dict[int, bytes] = {}
    try:
        src = pdfium.PdfDocument(data)
        for page_idx in page_indices:
            dst = pdfium.PdfDocument.new()
            try:
                dst.import_pages(src, pages=[page_idx])
                with io.BytesIO() as buf:
                    dst.save(buf)
                    blobs[page_idx] = buf.getvalue()
            finally:
                dst.close()
    finally:
        src.close()
    return blobs


def _render_pdf_page(blob: bytes, params: PDFConversionParams, out_path: Path) -> Tuple[int, int, int]:
    """Render a single-page PDF blob to an image file. Returns (width, height, size_bytes)."""
    from PIL import Image
    import pypdfium2 as pdfium

    fmt = _fmt_str(params.format)
    doc = None
    bitmap = None
    try:
        doc = pdfium.PdfDocument(blob)
        scale = params.dpi / 72.0
        bitmap = doc[0].render(scale=scale)
        img = bitmap.to_pil()
    except Exception as exc:
        raise MediaConversionError(f"Failed to render PDF page: {exc}", status_code=500) from exc
    finally:
        if bitmap is not None:
            try:
                bitmap.close()
            except Exception:
                pass
        if doc is not None:
            try:
                doc.close()
            except Exception:
                pass

    if fmt == "JPEG":
        img = _normalise_for_jpeg(img)
    else:
        if params.transparent_bg:
            img = _normalise_for_png(img)
        else:
            img = _normalise_for_png(img)
            if img.mode in ("RGBA", "LA"):
                bg = Image.new("RGB", img.size, (255, 255, 255))
                bg.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
                img = bg
    if params.grayscale:
        img = img.convert("L")

    try:
        img.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))
    except OSError as exc:
        raise MediaConversionError(f"Failed to write output file '{out_path.name}': {exc}", status_code=500) from exc

    stat = out_path.stat()
    return img.width, img.height, stat.st_size


def _stitch_images(page_files: List[Tuple[int, Path]], fmt: str, out_path: Path, quality: int) -> Tuple[int, int, int]:
    """Stitch rendered page images vertically into one tall image."""
    from PIL import Image

    images: List[Image.Image] = []
    for _, path in sorted(page_files, key=lambda t: t[0]):
        img = Image.open(path)
        if fmt == "JPEG":
            img = _normalise_for_jpeg(img)
        else:
            img = _normalise_for_png(img)
        images.append(img)

    total_width = max(im.width for im in images)
    total_height = sum(im.height for im in images)
    mode = images[0].mode
    fill = (255, 255, 255) if mode == "RGB" else 255
    stitched = Image.new(mode, (total_width, total_height), color=fill)
    y_offset = 0
    for img in images:
        stitched.paste(img, (0, y_offset))
        y_offset += img.height

    stitched.save(str(out_path), format=fmt, **_save_kwargs(fmt, quality))
    stat = out_path.stat()
    return stitched.width, stitched.height, stat.st_size


def _convert_image_bytes(data: bytes, params: ImageConversionParams, out_path: Path) -> Tuple[int, int, int, str]:
    """Convert raw image bytes to the requested format. Returns (width, height, size_bytes, detected_format)."""
    from PIL import Image, UnidentifiedImageError

    try:
        img = Image.open(io.BytesIO(data))
        detected = (img.format or "UNKNOWN").upper()
        img.load()
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise MediaConversionError(
            f"Input is not a valid or supported image: {exc}", status_code=422
        ) from exc

    if params.rotate:
        img = img.rotate(params.rotate % 360, expand=True)
    if params.flip == "horizontal":
        img = img.transpose(Image.FLIP_LEFT_RIGHT)
    elif params.flip == "vertical":
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    if params.grayscale:
        img = img.convert("L")

    if params.width is not None or params.height is not None:
        orig_w, orig_h = img.size
        if params.width is not None and params.height is not None:
            target = (params.width, params.height)
        elif params.width is not None:
            ratio = params.width / orig_w
            target = (params.width, max(1, round(orig_h * ratio)))
        else:
            ratio = params.height / orig_h
            target = (max(1, round(orig_w * ratio)), params.height)
        if target[0] * target[1] > _settings.media_max_image_pixels:
            raise MediaConversionError(
                f"Resized image ({target[0]}x{target[1]}) exceeds the "
                f"{_settings.media_max_image_pixels} pixel limit.",
                status_code=422,
            )
        img = img.resize(target, Image.LANCZOS)

    fmt = _fmt_str(params.format)
    if detected == fmt:
        raise MediaConversionError(
            f"Input is already {fmt}. Same-format conversion is not supported; "
            "choose a different output format.",
            status_code=422,
        )
    if fmt == "JPEG":
        img = _normalise_for_jpeg(img)
    elif fmt in ("PNG", "WEBP") and img.mode not in ("RGB", "RGBA", "L", "LA"):
        img = _normalise_for_png(img)

    try:
        img.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))
    except OSError as exc:
        raise MediaConversionError(f"Failed to write output file '{out_path.name}': {exc}", status_code=500) from exc

    stat = out_path.stat()
    return img.width, img.height, stat.st_size, detected


# ---------------------------------------------------------------------------
# Service
# ---------------------------------------------------------------------------

class MediaConversionService:
    """Orchestrates media conversion plus local/storage exposure of results."""

    def __init__(self, storage=None) -> None:
        self._storage = storage

    async def _resolve_storage(self):
        if self._storage is not None:
            return self._storage
        from app.services.media_storage_service import get_storage_service

        return await get_storage_service()

    async def convert_pdf(
        self,
        data: bytes,
        params: PDFConversionParams,
        source: str,
        job_id: Optional[str] = None,
    ) -> MediaConversionData:
        """Convert a PDF to images. Runs validation + rendering in the thread pool."""
        job_id = job_id or str(uuid.uuid4())
        fmt = _fmt_str(params.format)
        if fmt not in _PDF_OUTPUT_FORMATS:
            raise MediaConversionError(
                f"Unsupported output format '{fmt}'. PDF can only be converted to JPEG, PNG or WEBP.",
                status_code=422,
            )

        loop = asyncio.get_running_loop()
        total_pages = await loop.run_in_executor(_thread_pool, _pdf_total_pages, data)
        if total_pages > _settings.media_max_pages:
            raise MediaConversionError(
                f"PDF has {total_pages} pages, exceeding the {_settings.media_max_pages} page limit.",
                status_code=422,
            )
        page_indices = _parse_pages(params.pages, total_pages)
        _guard_memory(len(page_indices), params.dpi)

        start = time.perf_counter()
        job_dir = self._job_dir(job_id)
        out_dir = job_dir / "out"
        out_dir.mkdir(parents=True, exist_ok=True)

        page_blobs = await loop.run_in_executor(
            _thread_pool, _split_pdf_pages, data, page_indices
        )

        if params.split_page:
            render_coros = [
                loop.run_in_executor(
                    _thread_pool,
                    _render_pdf_page,
                    page_blobs[idx],
                    params,
                    out_dir / f"page_{idx + 1:04d}.{_ext_for(fmt)}",
                )
                for idx in sorted(page_blobs)
            ]
            outcomes = await asyncio.gather(*render_coros, return_exceptions=True)
            files: List[MediaOutputFile] = []
            for idx, outcome in zip(sorted(page_blobs), outcomes):
                if isinstance(outcome, Exception):
                    _logger.error("page_render_failed job=%s page=%d error=%s", job_id, idx + 1, outcome)
                    raise MediaConversionError(
                        f"Failed to render page {idx + 1}: {outcome}", status_code=500
                    ) from outcome
                width, height, size = outcome
                files.append(self._output_file(
                    f"page_{idx + 1:04d}.{_ext_for(fmt)}", idx + 1,
                    width, height, size, fmt,
                ))
        else:
            if len(page_blobs) == 1:
                width, height, size = await loop.run_in_executor(
                    _thread_pool,
                    _render_pdf_page,
                    page_blobs[sorted(page_blobs)[0]],
                    params,
                    out_dir / f"stitched.{_ext_for(fmt)}",
                )
            else:
                render_coros = [
                    loop.run_in_executor(
                        _thread_pool,
                        _render_pdf_page,
                        page_blobs[idx],
                        params,
                        out_dir / f"_page_{idx + 1:04d}.{_ext_for(fmt)}",
                    )
                    for idx in sorted(page_blobs)
                ]
                outcomes = await asyncio.gather(*render_coros, return_exceptions=True)
                for idx, outcome in zip(sorted(page_blobs), outcomes):
                    if isinstance(outcome, Exception):
                        raise MediaConversionError(
                            f"Failed to render page {idx + 1}: {outcome}", status_code=500
                        ) from outcome
                page_files = [(idx, out_dir / f"_page_{idx + 1:04d}.{_ext_for(fmt)}") for idx in sorted(page_blobs)]
                width, height, size = await loop.run_in_executor(
                    _thread_pool,
                    _stitch_images,
                    page_files,
                    fmt,
                    out_dir / f"stitched.{_ext_for(fmt)}",
                    params.quality,
                )
                for _, tmp_path in page_files:
                    tmp_path.unlink(missing_ok=True)
            files = [self._output_file(
                f"stitched.{_ext_for(fmt)}", None, width, height, size, fmt,
            )]

        upload, warning = await self._expose(files, job_id, out_dir, job_dir)
        duration_ms = round((time.perf_counter() - start) * 1000, 2)
        _logger.info(
            "pdf_conversion_complete job=%s pages=%d duration_ms=%s mode=%s",
            job_id, total_pages, duration_ms, upload.mode,
        )
        return MediaConversionData(
            job_id=job_id,
            source=source,
            input_format="PDF",
            output_format=fmt,
            total_pages=total_pages,
            converted_files=len(files),
            outputs=files,
            upload=upload,
            warning=warning,
        )

    async def convert_image(
        self,
        data: bytes,
        filename: str,
        params: ImageConversionParams,
        job_id: Optional[str] = None,
    ) -> MediaConversionData:
        """Convert an image to a target image format."""
        job_id = job_id or str(uuid.uuid4())
        fmt = _fmt_str(params.format)
        if fmt not in _IMAGE_OUTPUT_FORMATS:
            raise MediaConversionError(
                f"Unsupported output format '{fmt}'.",
                status_code=422,
            )

        start = time.perf_counter()
        job_dir = self._job_dir(job_id)
        out_dir = job_dir / "out"
        out_dir.mkdir(parents=True, exist_ok=True)
        stem = Path(filename or "image").stem or "image"
        out_path = out_dir / f"{stem}.{_ext_for(fmt)}"

        loop = asyncio.get_running_loop()
        width, height, size, detected = await loop.run_in_executor(
            _thread_pool, _convert_image_bytes, data, params, out_path
        )

        file = self._output_file(out_path.name, None, width, height, size, fmt)
        upload, warning = await self._expose([file], job_id, out_dir, job_dir)
        duration_ms = round((time.perf_counter() - start) * 1000, 2)
        _logger.info(
            "image_conversion_complete job=%s input=%s duration_ms=%s mode=%s",
            job_id, detected, duration_ms, upload.mode,
        )
        return MediaConversionData(
            job_id=job_id,
            source=filename,
            input_format=detected or "IMAGE",
            output_format=fmt,
            total_pages=1,
            converted_files=1,
            outputs=[file],
            upload=upload,
            warning=warning,
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _job_dir(job_id: str) -> Path:
        root = Path(_settings.media_output_dir).resolve()
        root.mkdir(parents=True, exist_ok=True)
        return root / job_id

    @staticmethod
    def _output_file(
        filename: str,
        page_number: Optional[int],
        width: int,
        height: int,
        size_bytes: int,
        fmt: str,
    ) -> MediaOutputFile:
        return MediaOutputFile(
            filename=filename,
            page_number=page_number,
            width=width,
            height=height,
            size_bytes=size_bytes,
            format=fmt,
            content_type=_MIME_BY_FORMAT.get(fmt, "application/octet-stream"),
            url="",  # filled in by _expose
        )

    async def _expose(
        self,
        files: List[MediaOutputFile],
        job_id: str,
        out_dir: Path,
        job_dir: Path,
    ) -> Tuple[MediaUploadSummary, Optional[str]]:
        """Expose output files via Supabase signed URLs or local download endpoints."""
        if _settings.supabase_upload_enabled:
            storage = await self._resolve_storage()
            bucket = await storage.ensure_bucket(_settings.supabase_storage_bucket)
            ttl = _settings.supabase_signed_url_ttl_seconds
            semaphore = asyncio.Semaphore(max(1, _settings.media_upload_concurrency))

            from app.services.media_storage_service import iso_expiry

            async def _upload_one(f: MediaOutputFile) -> MediaOutputFile:
                storage_path = f"{job_id}/{f.filename}"
                async with semaphore:
                    await storage.upload_file(
                        bucket, storage_path,
                        (out_dir / f.filename).read_bytes(), f.content_type,
                    )
                    f.url = await storage.create_signed_url(bucket, storage_path, ttl)
                return f

            outcomes = await asyncio.gather(*[_upload_one(f) for f in files], return_exceptions=True)
            failed = 0
            for f, outcome in zip(files, outcomes):
                if isinstance(outcome, Exception):
                    failed += 1
                    _logger.error("storage_upload_failed job=%s file=%s error=%s", job_id, f.filename, outcome)
                    f.url = f"data:{f.content_type};base64," + base64.b64encode((out_dir / f.filename).read_bytes()).decode("ascii")

            expires_at = iso_expiry(ttl)
            warning = (
                f"Converted files were uploaded to Supabase Storage bucket "
                f"'{bucket}'. The returned signed URLs are valid for 24 hours "
                f"(expire at {expires_at}). Regenerate by re-running the conversion."
            )
            return (
                MediaUploadSummary(
                    mode="storage",
                    bucket=bucket,
                    total_files=len(files),
                    failed_uploads=failed,
                    url_ttl_seconds=ttl,
                    expires_at=expires_at,
                ),
                warning,
            )

        for f in files:
            f.url = f"data:{f.content_type};base64," + base64.b64encode((out_dir / f.filename).read_bytes()).decode("ascii")
        return MediaUploadSummary(mode="local", total_files=len(files), failed_uploads=0), None


media_conversion_service = MediaConversionService()