from __future__ import annotations

import io
from pathlib import Path

from django.core.exceptions import ValidationError
from django.core.files.base import ContentFile
from PIL import Image, ImageFile, ImageOps, UnidentifiedImageError

ImageFile.LOAD_TRUNCATED_IMAGES = True

MAX_PDF_BYTES = 2 * 1024 * 1024
IMAGE_DIMENSION_STEPS = (2200, 1800, 1600, 1400, 1200, 1000, 800, 640)
JPEG_QUALITY_STEPS = (82, 74, 66, 58, 50, 42, 34)
DPI_STEPS = (150, 130, 110, 96, 84, 72)


def normalize_uploaded_file_to_pdf(uploaded_file) -> ContentFile:
    original_name = uploaded_file.name or "document"
    extension = Path(original_name).suffix.lower()
    raw_bytes = uploaded_file.read()

    if extension == ".pdf":
        normalized_name = f"{Path(original_name).stem or 'document'}.pdf"
        _validate_pdf_size(raw_bytes, normalized_name)
        return ContentFile(raw_bytes, name=normalized_name)

    try:
        return _convert_image_bytes_to_pdf(raw_bytes, original_name=original_name)
    except UnidentifiedImageError as exc:
        raise ValidationError("Upload a PDF, JPG, JPEG, or PNG file.") from exc


def _convert_image_bytes_to_pdf(raw_bytes: bytes, *, original_name: str) -> ContentFile:
    with Image.open(io.BytesIO(raw_bytes)) as source_image:
        image = _prepare_image(source_image)

    output_name = f"{Path(original_name).stem or 'document'}.pdf"

    for max_dimension in IMAGE_DIMENSION_STEPS:
        resized = _resize_image(image, max_dimension=max_dimension)
        for quality in JPEG_QUALITY_STEPS:
            for dpi in DPI_STEPS:
                pdf_bytes = _render_image_pdf_bytes(resized, quality=quality, dpi=dpi)
                if len(pdf_bytes) <= MAX_PDF_BYTES:
                    return ContentFile(pdf_bytes, name=output_name)

    raise ValidationError(
        f"'{original_name}' could not be reduced below 2 MB. Upload a smaller image or PDF."
    )


def _prepare_image(image: Image.Image) -> Image.Image:
    image = ImageOps.exif_transpose(image)

    if getattr(image, "is_animated", False):
        image.seek(0)

    if image.mode in {"RGBA", "LA"}:
        background = Image.new("RGB", image.size, "white")
        alpha = image.getchannel("A")
        background.paste(image.convert("RGBA"), mask=alpha)
        return background

    if image.mode == "P":
        return image.convert("RGB")

    if image.mode != "RGB":
        return image.convert("RGB")

    return image.copy()


def _resize_image(image: Image.Image, *, max_dimension: int) -> Image.Image:
    current_max = max(image.size)
    if current_max <= max_dimension:
        return image.copy()

    ratio = max_dimension / float(current_max)
    resized_dimensions = (
        max(1, int(image.width * ratio)),
        max(1, int(image.height * ratio)),
    )
    return image.resize(resized_dimensions, Image.Resampling.LANCZOS)


def _render_image_pdf_bytes(image: Image.Image, *, quality: int, dpi: int) -> bytes:
    jpeg_buffer = io.BytesIO()
    image.save(
        jpeg_buffer,
        format="JPEG",
        quality=quality,
        optimize=True,
        progressive=True,
        dpi=(dpi, dpi),
        subsampling=2,
    )
    jpeg_buffer.seek(0)

    with Image.open(jpeg_buffer) as jpeg_image:
        pdf_buffer = io.BytesIO()
        jpeg_image.save(pdf_buffer, format="PDF", resolution=dpi)
        return pdf_buffer.getvalue()


def _validate_pdf_size(pdf_bytes: bytes, filename: str) -> None:
    if len(pdf_bytes) > MAX_PDF_BYTES:
        total_kb = int(round(len(pdf_bytes) / 1024))
        raise ValidationError(
            f"'{filename}' is {total_kb} KB. Upload a PDF or image file that is 2 MB or smaller."
        )