app/services/conversion.py

from __future__ import annotations

import asyncio
import io
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import List, Optional, Tuple

import structlog
from PIL import Image

from app.core.config import settings
from app.core.exceptions import ConversionError, InvalidParameterError, TooManyPagesError
from app.models.schemas import ConversionParams, PageResult

logger = structlog.get_logger(__name__)

_SUPPORTED_FORMATS = {"JPEG", "PNG"}
_MAX_MEMORY_ESTIMATE_MB = 512

_render_pool = ThreadPoolExecutor(
    max_workers=settings.PAGE_RENDER_WORKERS,
    thread_name_prefix="page_renderer",
)


def _format_str(fmt) -> str:
    """Normalise an ImageFormat enum / string to its raw uppercase value."""
    return fmt.value if hasattr(fmt, "value") else str(fmt)


def _save_kwargs(fmt: str, quality: int) -> dict:
    """PIL save kwargs for the requested output format."""
    if fmt == "JPEG":
        return {"quality": quality, "optimize": True}
    if fmt == "PNG":
        return {"optimize": True}
    return {}


def _ext_for(fmt: str) -> str:
    return "jpg" if fmt == "JPEG" else "png"


class ConversionService:
    def __init__(self) -> None:
        self._output_dir = Path(settings.OUTPUT_DIR)
        self._output_dir.mkdir(parents=True, exist_ok=True)

    async def convert(
        self,
        pdf_bytes: bytes,
        params: ConversionParams,
        job_id: Optional[str] = None,
    ) -> Tuple[List[PageResult], float]:
        """
        Full pipeline:
          1. Validate PDF + determine page indices
          2. Render pages (always individually for parallelism)
          3a. split_page=True / omitted  → return one PageResult per page
          3b. split_page=False           → stitch all pages into one tall image,
                                           return a single PageResult
        """
        job_id = job_id or str(uuid.uuid4())
        job_dir = self._output_dir / job_id
        job_dir.mkdir(parents=True, exist_ok=True)

        start = time.perf_counter()
        loop = asyncio.get_event_loop()

        try:
            page_indices, total_pages = await loop.run_in_executor(
                _render_pool, _validate_and_get_indices, pdf_bytes, params
            )
        except (ConversionError, InvalidParameterError, TooManyPagesError):
            raise
        except Exception as exc:
            logger.exception("pdf_validation_failed", job_id=job_id)
            raise ConversionError(str(exc)) from exc

        logger.info(
            "conversion_started",
            job_id=job_id,
            total_pages=total_pages,
            rendering=len(page_indices),
            dpi=params.dpi,
            format=params.format,
            split_page=params.split_page,
        )

        try:
            page_blobs: dict[int, bytes] = await loop.run_in_executor(
                _render_pool, _split_pdf_pages, pdf_bytes, page_indices
            )
        except Exception as exc:
            logger.exception("pdf_split_failed", job_id=job_id)
            raise ConversionError(f"Failed to split PDF into pages: {exc}") from exc

        page_results = await self._render_parallel(
            page_blobs, params, job_dir, job_id, loop
        )

        if params.split_page is False:
            try:
                stitched = await loop.run_in_executor(
                    _render_pool, _stitch_pages, page_results, params, job_dir
                )
                page_results = [stitched]
                logger.info("pages_stitched", job_id=job_id, total_pages=len(page_blobs))
            except Exception as exc:
                logger.exception("stitch_failed", job_id=job_id)
                raise ConversionError(f"Failed to stitch pages into single image: {exc}") from exc

        duration_ms = round((time.perf_counter() - start) * 1000, 2)
        logger.info("conversion_complete", job_id=job_id, pages=len(page_results), duration_ms=duration_ms)
        return page_results, duration_ms

    async def _render_parallel(
        self,
        page_blobs: dict[int, bytes],
        params: ConversionParams,
        job_dir: Path,
        job_id: str,
        loop: asyncio.AbstractEventLoop,
    ) -> List[PageResult]:
        ordered_indices = sorted(page_blobs.keys())
        render_coros = [
            loop.run_in_executor(
                _render_pool,
                _render_single_page,
                page_idx,
                page_blobs[page_idx],
                params,
                job_dir,
            )
            for page_idx in ordered_indices
        ]

        raw_results = await asyncio.gather(*render_coros, return_exceptions=True)

        results_by_idx: dict[int, PageResult] = {}
        errors: list[str] = []

        for page_idx, outcome in zip(ordered_indices, raw_results):
            if isinstance(outcome, Exception):
                logger.error("page_render_failed", job_id=job_id, page=page_idx + 1, error=str(outcome))
                errors.append(f"Page {page_idx + 1}: {outcome}")
            else:
                results_by_idx[page_idx] = outcome
                logger.debug("page_rendered", job_id=job_id, page=page_idx + 1, size_bytes=outcome.size_bytes)

        if not results_by_idx:
            raise ConversionError(f"All pages failed to render. Errors: {'; '.join(errors)}")

        if errors:
            logger.warning("some_pages_failed", job_id=job_id, failed=len(errors), succeeded=len(results_by_idx))

        return [results_by_idx[i] for i in sorted(results_by_idx)]


def _validate_and_get_indices(
    pdf_bytes: bytes,
    params: ConversionParams,
) -> Tuple[List[int], int]:
    from pypdf import PdfReader
    from pypdf.errors import PdfReadError

    fmt = _format_str(params.format)
    if fmt not in _SUPPORTED_FORMATS:
        raise InvalidParameterError(
            f"Unsupported format '{fmt}'. Only JPEG (JPG) and PNG are supported."
        )

    try:
        reader = PdfReader(io.BytesIO(pdf_bytes))
        if reader.is_encrypted:
            from pypdf import PasswordType
            if reader.decrypt("") == PasswordType.NOT_DECRYPTED:
                raise ConversionError("PDF is password-protected. Please provide an unlocked PDF.")
        total_pages = len(reader.pages)
    except ConversionError:
        raise
    except PdfReadError as exc:
        raise ConversionError(f"Malformed PDF: {exc}") from exc

    if total_pages == 0:
        raise ConversionError("PDF contains no pages")
    if total_pages > settings.MAX_PAGES:
        raise TooManyPagesError(settings.MAX_PAGES)

    page_indices = params.parse_pages(total_pages)
    _guard_memory(len(page_indices), params.dpi)

    return page_indices, total_pages


def _split_pdf_pages(pdf_bytes: bytes, page_indices: List[int]) -> dict[int, bytes]:
    """Extract each requested page into its own single-page PDF blob for parallel rendering."""
    from pypdf import PdfReader, PdfWriter

    reader = PdfReader(io.BytesIO(pdf_bytes))
    blobs: dict[int, bytes] = {}

    for page_idx in page_indices:
        writer = PdfWriter()
        writer.add_page(reader.pages[page_idx])
        buf = io.BytesIO()
        writer.write(buf)
        blobs[page_idx] = buf.getvalue()

    logger.debug("pdf_split_complete", pages=len(blobs))
    return blobs


def _normalise_for_jpeg(img: Image.Image) -> Image.Image:
    """Return an RGB image suitable for JPEG output, flattening alpha onto white."""
    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: Image.Image) -> Image.Image:
    """Return an image whose mode PIL accepts for PNG (RGB / RGBA / L)."""
    if img.mode in ("RGB", "RGBA", "L"):
        return img
    return img.convert("RGB")


def _render_single_page(
    page_idx: int,
    page_pdf_bytes: bytes,
    params: ConversionParams,
    job_dir: Path,
) -> PageResult:
    """Render one single-page PDF blob to an image file on disk."""
    from pdf2image import convert_from_bytes
    from pdf2image.exceptions import PDFInfoNotInstalledError, PDFPageCountError, PDFSyntaxError

    fmt = _format_str(params.format)

    try:
        pil_images = convert_from_bytes(
            page_pdf_bytes,
            dpi=params.dpi,
            fmt="ppm",
            thread_count=1,
            transparent=params.transparent_bg and fmt == "PNG",
            use_pdftocairo=True,
        )
    except (PDFInfoNotInstalledError, PDFPageCountError, PDFSyntaxError) as exc:
        raise ConversionError(f"Page {page_idx + 1} render failed: {exc}") from exc

    if not pil_images:
        raise ConversionError(f"Page {page_idx + 1}: renderer returned no image")

    img = pil_images[0]
    if fmt == "JPEG":
        img = _normalise_for_jpeg(img)
    if params.grayscale:
        img = img.convert("L")

    ext = _ext_for(fmt)
    filename = f"page_{page_idx + 1:04d}.{ext}"
    out_path = job_dir / filename

    try:
        img.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))
    except OSError as exc:
        raise ConversionError(f"Failed to write {filename}: {exc}") from exc

    stat = out_path.stat()
    return PageResult(
        page_number=page_idx + 1,
        download_url=f"/api/v1/files/{job_dir.name}/{filename}",
        width=img.width,
        height=img.height,
        size_bytes=stat.st_size,
        format=fmt,
    )


def _stitch_pages(
    page_results: List[PageResult],
    params: ConversionParams,
    job_dir: Path,
) -> PageResult:
    """
    Stitch all rendered page images vertically into a single tall image.
    Called only when split_page=False.
    Pages are placed top-to-bottom in document order with no gaps.
    """
    fmt = _format_str(params.format)
    ext = _ext_for(fmt)

    images: list[Image.Image] = []
    for pr in sorted(page_results, key=lambda p: p.page_number):
        path = job_dir / Path(pr.download_url).name
        img = Image.open(str(path))
        if fmt == "JPEG":
            img = _normalise_for_jpeg(img)
        else:
            img = _normalise_for_png(img)
        images.append(img)

    if not images:
        raise ConversionError("No rendered page images found to stitch")

    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

    out_filename = f"stitched.{ext}"
    out_path = job_dir / out_filename
    stitched.save(str(out_path), format=fmt, **_save_kwargs(fmt, params.quality))

    stat = out_path.stat()
    return PageResult(
        page_number=1,
        download_url=f"/api/v1/files/{job_dir.name}/{out_filename}",
        width=stitched.width,
        height=stitched.height,
        size_bytes=stat.st_size,
        format=fmt,
    )


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 > _MAX_MEMORY_ESTIMATE_MB:
        raise InvalidParameterError(
            f"Requested conversion would require approximately {estimated_mb:.0f} MB of memory "
            f"({page_count} pages at {dpi} DPI). "
            f"Reduce DPI or use the 'pages' parameter to select fewer pages."
        )


conversion_service = ConversionService()