"""
Estratégias de compressão e otimização de PDF.

Usa pypdf + Pillow para re-amostragem de imagens e remoção de
objetos redundantes.  Mantém a estrutura original do PDF (inclusive
PDF/A) quando possível.
"""

import io
from typing import Optional

from PIL import Image
from pypdf import PdfReader, PdfWriter
from pypdf.constants import ImageAttributes as IA

from config import QUALITY_PRESETS


def _downsample_image(
    img_bytes: bytes,
    quality: int = 60,
    max_dim: int = 2000,
    fmt: str = "JPEG",
) -> bytes:
    """
    Redimensiona e re-comprime uma imagem.
    - quality: 1-100 (JPEG)
    - max_dim: lado máximo em px
    - fmt: "JPEG" ou "PNG"
    """
    img = Image.open(io.BytesIO(img_bytes))

    # Converte RGBA/CMYK/P (palette) para RGB para compressão JPEG
    if img.mode not in ("RGB", "L"):
        if fmt == "JPEG":
            img = img.convert("RGB")
        elif img.mode == "P":
            img = img.convert("RGB")

    # Redimensiona proporcionalmente se exceder max_dim
    w, h = img.size
    if max(w, h) > max_dim:
        ratio = max_dim / max(w, h)
        new_size = (int(w * ratio), int(h * ratio))
        # Usa ANTIALIAS/LANCZOS para melhor qualidade na redução
        img = img.resize(new_size, Image.LANCZOS)

    buf = io.BytesIO()
    if fmt == "JPEG":
        img.save(buf, format="JPEG", quality=quality, optimize=True)
    else:
        img.save(buf, format="PNG", optimize=True)

    return buf.getvalue()


def compress_pdf(
    data: bytes,
    preset: str = "ebook",
    remove_unused: bool = True,
) -> bytes:
    """
    Comprime um PDF re-amostrando imagens e otimizando estrutura.

    Args:
        data:       Conteúdo do PDF original.
        preset:     Perfil de qualidade ("screen", "ebook", "printer", "prepress").
        remove_unused: Remove objetos não referenciados.

    Retorna bytes do PDF comprimido.
    """
    opts = QUALITY_PRESETS.get(preset, QUALITY_PRESETS["ebook"])
    quality = opts["quality"]
    max_dim = opts["max_dim"]

    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter()

    # Copia metadados
    writer.metadata = reader.metadata

    # Registra fontes e objetos que usamos (para limpeza)
    used_objects = set()

    for page_num, page in enumerate(reader.pages):
        used_objects.add(id(page))

        # Verifica se há imagens na página
        if "/XObject" in page.get("/Resources", {}):
            resources = page["/Resources"]
            if "/XObject" in resources:
                xobjects = resources["/XObject"]
                for obj_name in list(xobjects.keys()):
                    xobj = xobjects[obj_name]
                    subtype = xobj.get("/Subtype", "")
                    if subtype == "/Image":
                        try:
                            raw = xobj.get_data()
                            if raw:
                                # Determina formato — PNG se tiver transparência
                                fmt = "PNG" if xobj.get("/ColorSpace") == "/DeviceCMYK" else "JPEG"
                                compressed = _downsample_image(raw, quality, max_dim, fmt)
                                xobj._data = compressed
                                xobj[IA.WIDTH], xobj[IA.HEIGHT] = Image.open(
                                    io.BytesIO(compressed)
                                ).size
                        except Exception:
                            # Se falhar na recompressão, mantém original
                            pass

        # Adiciona página modificada
        writer.add_page(page)

    # Escreve com compressão de streams
    out_buf = io.BytesIO()
    writer.write_stream(out_buf)

    return out_buf.getvalue()


def linearize_pdf(data: bytes) -> bytes:
    """
    Re-escreve o PDF de forma linearizada (otimizada para web).
    Usa pypdf para reestruturar.
    """
    reader = PdfReader(io.BytesIO(data))
    writer = PdfWriter()
    writer.metadata = reader.metadata
    for page in reader.pages:
        writer.add_page(page)

    out_buf = io.BytesIO()
    writer.write_stream(out_buf)
    return out_buf.getvalue()