"""
PDF and TIFF Rasterizer for Ink Estimation Engine
Renders vector PDF and image inputs directly to 8-bit CMYK buffers.
Eliminates Adobe Photoshop dependency entirely.
"""

import os
import pymupdf as fitz
from PIL import Image
import numpy as np

Image.MAX_IMAGE_PIXELS = None


class PageRaster:
    def __init__(self, page_num: int, width: int, height: int, dpi: float, width_mm: float, height_mm: float, cmyk_bytes: bytes):
        self.page_num = page_num
        self.width = width
        self.height = height
        self.dpi = dpi
        self.width_mm = width_mm
        self.height_mm = height_mm
        self.cmyk_bytes = cmyk_bytes
        self.pixel_count = width * height


class DocumentRasterizer:
    @staticmethod
    def inspect_document(filepath: str):
        """Returns metadata about the document (pages, dimensions, format)."""
        ext = os.path.splitext(filepath)[1].lower()
        if ext in [".pdf"]:
            doc = fitz.open(filepath)
            pages_info = []
            for i, page in enumerate(doc):
                rect = page.rect
                w_mm = rect.width * 25.4 / 72.0
                h_mm = rect.height * 25.4 / 72.0
                pages_info.append({
                    "page_number": i + 1,
                    "width_pt": rect.width,
                    "height_pt": rect.height,
                    "width_mm": round(w_mm, 2),
                    "height_mm": round(h_mm, 2)
                })
            doc_info = {
                "format": "PDF",
                "page_count": len(doc),
                "pages": pages_info,
                "title": doc.metadata.get("title", "") if doc.metadata else ""
            }
            doc.close()
            return doc_info
        elif ext in [".tif", ".tiff"]:
            img = Image.open(filepath)
            dpi = img.info.get("dpi", (300.0, 300.0))
            dpi_val = dpi[0] if isinstance(dpi, tuple) else (dpi or 300.0)
            w_mm = img.size[0] * 25.4 / dpi_val
            h_mm = img.size[1] * 25.4 / dpi_val
            return {
                "format": "TIFF",
                "page_count": 1,
                "pages": [{
                    "page_number": 1,
                    "width_px": img.size[0],
                    "height_px": img.size[1],
                    "width_mm": round(w_mm, 2),
                    "height_mm": round(h_mm, 2),
                    "detected_dpi": dpi_val,
                    "mode": img.mode
                }]
            }
        else:
            raise ValueError(f"Unsupported file format: {ext}. Please upload a PDF or TIFF file.")

    @staticmethod
    def rasterize_page(filepath: str, page_index: int = 0, dpi: float = 300.0) -> PageRaster:
        """Rasterizes a single page of a PDF or TIFF into 8-bit CMYK bytes."""
        ext = os.path.splitext(filepath)[1].lower()

        if ext == ".pdf":
            doc = fitz.open(filepath)
            if page_index < 0 or page_index >= len(doc):
                doc.close()
                raise IndexError(f"Page index {page_index} out of range (total pages: {len(doc)})")

            page = doc[page_index]
            rect = page.rect
            width_mm = rect.width * 25.4 / 72.0
            height_mm = rect.height * 25.4 / 72.0

            # Render directly to CMYK colorspace at requested DPI
            pix = page.get_pixmap(colorspace=fitz.csCMYK, dpi=int(dpi), alpha=False)
            cmyk_bytes = pix.samples
            width = pix.width
            height = pix.height
            doc.close()

            return PageRaster(
                page_num=page_index + 1,
                width=width,
                height=height,
                dpi=dpi,
                width_mm=round(width_mm, 2),
                height_mm=round(height_mm, 2),
                cmyk_bytes=cmyk_bytes
            )

        elif ext in [".tif", ".tiff"]:
            img = Image.open(filepath)
            img_dpi = img.info.get("dpi", (dpi, dpi))
            dpi_val = img_dpi[0] if isinstance(img_dpi, tuple) else (img_dpi or dpi)

            # Ensure CMYK mode
            if img.mode != "CMYK":
                img = img.convert("CMYK")

            width_mm = img.size[0] * 25.4 / dpi_val
            height_mm = img.size[1] * 25.4 / dpi_val
            cmyk_bytes = img.tobytes()

            return PageRaster(
                page_num=1,
                width=img.size[0],
                height=img.size[1],
                dpi=dpi_val,
                width_mm=round(width_mm, 2),
                height_mm=round(height_mm, 2),
                cmyk_bytes=cmyk_bytes
            )
        else:
            raise ValueError(f"Unsupported file format: {ext}")