"""Core in-memory PDF manipulation utilities for merging, splitting, rotating,
compressing, format conversion, and OCR processing.
"""

from __future__ import annotations

import io
import os
import re
from typing import Sequence
import numpy as np
from PIL import Image

try:
    import pymupdf as fitz
except ImportError:  # pragma: no cover
    import fitz


class PDFProcessingError(ValueError):
    """Raised when PDF processing fails due to invalid input, encryption, or corruption."""


def _open_pdf(pdf_data: bytes) -> fitz.Document:
    """Safely opens an in-memory PDF, checking for corruption and encryption."""
    if not pdf_data:
        raise PDFProcessingError("The provided PDF data is empty.")
    try:
        doc = fitz.open(stream=pdf_data, filetype="pdf")
    except Exception as exc:
        raise PDFProcessingError(f"Could not open or parse the PDF document: {exc}") from exc

    if doc.needs_pass:
        doc.close()
        raise PDFProcessingError("The PDF document is password-protected. Please upload an unlocked PDF.")
    return doc


def parse_page_ranges(range_str: str, max_pages: int) -> list[int]:
    """
    Parses a page range string (1-indexed) into a list of 0-indexed page indices.
    Supports formats: "1", "1-3", "1, 3, 5", "2-4, 1", "all".
    """
    if not range_str or not range_str.strip():
        raise PDFProcessingError("Page range cannot be empty.")

    clean_str = range_str.strip().lower()
    if clean_str == "all":
        return list(range(max_pages))

    indices: list[int] = []
    tokens = [t.strip() for t in clean_str.split(",") if t.strip()]
    if not tokens:
        raise PDFProcessingError("Invalid page range format.")

    for token in tokens:
        if "-" in token:
            parts = token.split("-")
            if len(parts) != 2:
                raise PDFProcessingError(f"Invalid range format: '{token}'")
            try:
                start = int(parts[0].strip())
                end = int(parts[1].strip())
            except ValueError:
                raise PDFProcessingError(f"Invalid numbers in range: '{token}'")

            if start < 1 or end > max_pages or start > end:
                raise PDFProcessingError(
                    f"Range '{start}-{end}' is invalid. Document has {max_pages} pages (allowed: 1-{max_pages})."
                )
            indices.extend(range(start - 1, end))
        else:
            try:
                page_num = int(token)
            except ValueError:
                raise PDFProcessingError(f"Invalid page number: '{token}'")

            if page_num < 1 or page_num > max_pages:
                raise PDFProcessingError(
                    f"Page number {page_num} is out of bounds (1-{max_pages})."
                )
            indices.append(page_num - 1)

    return indices


def merge_pdfs(pdf_bytes_list: Sequence[bytes]) -> bytes:
    """
    Merges multiple PDF files in sequential order.
    Returns the merged PDF as bytes.
    """
    if not pdf_bytes_list:
        raise PDFProcessingError("At least one PDF file is required to merge.")

    merged_doc = fitz.open()
    try:
        for index, pdf_data in enumerate(pdf_bytes_list):
            sub_doc = _open_pdf(pdf_data)
            try:
                merged_doc.insert_pdf(sub_doc)
            finally:
                sub_doc.close()

        if len(merged_doc) == 0:
            raise PDFProcessingError("The resulting merged PDF has 0 pages.")

        return merged_doc.tobytes(deflate=True, garbage=3)
    finally:
        merged_doc.close()


def split_pdf(pdf_bytes: bytes, page_range_str: str) -> bytes:
    """
    Extracts specified pages from a PDF and returns a new PDF containing only those pages.
    """
    src_doc = _open_pdf(pdf_bytes)
    try:
        target_indices = parse_page_ranges(page_range_str, len(src_doc))
        out_doc = fitz.open()
        try:
            for page_idx in target_indices:
                out_doc.insert_pdf(src_doc, from_page=page_idx, to_page=page_idx)
            return out_doc.tobytes(deflate=True, garbage=3)
        finally:
            out_doc.close()
    finally:
        src_doc.close()


def split_pdf_to_individual_pages(pdf_bytes: bytes) -> list[tuple[str, bytes]]:
    """
    Splits a multi-page PDF into a list of single-page PDFs: [('page_1.pdf', bytes), ...].
    """
    src_doc = _open_pdf(pdf_bytes)
    try:
        pages_list: list[tuple[str, bytes]] = []
        for i in range(len(src_doc)):
            single_doc = fitz.open()
            try:
                single_doc.insert_pdf(src_doc, from_page=i, to_page=i)
                page_bytes = single_doc.tobytes(deflate=True, garbage=3)
                pages_list.append((f"page_{i + 1}.pdf", page_bytes))
            finally:
                single_doc.close()
        return pages_list
    finally:
        src_doc.close()


def rotate_pdf_pages(
    pdf_bytes: bytes,
    rotation_degrees: int,
    page_range_str: str | None = None,
) -> bytes:
    """
    Rotates pages by the specified degrees (90, 180, 270).
    If page_range_str is provided, only specified pages are rotated.
    """
    if rotation_degrees not in (90, 180, 270):
        raise PDFProcessingError(
            f"Invalid rotation degrees: {rotation_degrees}. Allowed values: 90, 180, 270."
        )

    doc = _open_pdf(pdf_bytes)
    try:
        if page_range_str and page_range_str.strip():
            target_indices = set(parse_page_ranges(page_range_str, len(doc)))
        else:
            target_indices = set(range(len(doc)))

        for idx in target_indices:
            page = doc[idx]
            new_rotation = (page.rotation + rotation_degrees) % 360
            page.set_rotation(new_rotation)

        return doc.tobytes(deflate=True, garbage=3)
    finally:
        doc.close()


def compress_pdf(
    pdf_bytes: bytes,
    deflate: bool = True,
    clean: bool = True,
) -> tuple[bytes, int, int]:
    """
    Compresses and optimizes a PDF by deflating streams and cleaning redundant objects.
    Returns: (compressed_bytes, original_size_bytes, new_size_bytes)
    """
    doc = _open_pdf(pdf_bytes)
    try:
        orig_size = len(pdf_bytes)
        compressed_bytes = doc.tobytes(
            garbage=4,
            deflate=deflate,
            clean=clean,
        )
        return compressed_bytes, orig_size, len(compressed_bytes)
    finally:
        doc.close()


def images_to_pdf(images_data: Sequence[tuple[str, bytes]]) -> bytes:
    """
    Combines multiple images into a single PDF document.
    images_data: list of (filename, image_bytes)
    """
    if not images_data:
        raise PDFProcessingError("At least one image is required to build a PDF.")

    pil_images: list[Image.Image] = []
    try:
        for name, img_bytes in images_data:
            if not img_bytes:
                continue
            try:
                img = Image.open(io.BytesIO(img_bytes))
                # Convert RGBA/P/LA to RGB for clean PDF output
                if img.mode in ("RGBA", "P", "LA"):
                    rgb_img = Image.new("RGB", img.size, (255, 255, 255))
                    if img.mode == "RGBA":
                        rgb_img.paste(img, mask=img.split()[3])
                    else:
                        rgb_img.paste(img.convert("RGBA"), mask=img.convert("RGBA").split()[3])
                    pil_images.append(rgb_img)
                elif img.mode != "RGB":
                    pil_images.append(img.convert("RGB"))
                else:
                    pil_images.append(img)
            except Exception as exc:
                raise PDFProcessingError(f"Could not load image '{name}': {exc}") from exc

        if not pil_images:
            raise PDFProcessingError("No valid images could be decoded.")

        out_buf = io.BytesIO()
        pil_images[0].save(
            out_buf,
            format="PDF",
            save_all=True,
            append_images=pil_images[1:] if len(pil_images) > 1 else [],
            resolution=150.0,
        )
        return out_buf.getvalue()
    finally:
        for img in pil_images:
            img.close()


def pdf_to_images(
    pdf_bytes: bytes,
    dpi: int = 200,
    image_format: str = "PNG",
) -> list[tuple[str, bytes]]:
    """
    Renders each page of a PDF as a high-resolution image.
    Returns: [('page_1.png', image_bytes), ...]
    """
    doc = _open_pdf(pdf_bytes)
    ext = "jpg" if image_format.upper() in ("JPG", "JPEG") else "png"
    fmt = "jpeg" if ext == "jpg" else "png"

    try:
        rendered: list[tuple[str, bytes]] = []
        scale = max(dpi, 72) / 72
        matrix = fitz.Matrix(scale, scale)

        for i, page in enumerate(doc):
            pix = page.get_pixmap(matrix=matrix, alpha=False)
            img_bytes = pix.tobytes(output=fmt)
            rendered.append((f"page_{i + 1}.{ext}", img_bytes))

        return rendered
    finally:
        doc.close()


def ocr_scanned_pdf(
    pdf_bytes: bytes,
    ocr_reader,
    dpi: int = 200,
) -> tuple[str, bytes]:
    """
    Runs OCR on each page of a scanned PDF.
    Returns: (extracted_text, searchable_pdf_bytes)
    """
    doc = _open_pdf(pdf_bytes)
    extracted_pages: list[str] = []
    scale = max(dpi, 72) / 72
    matrix = fitz.Matrix(scale, scale)

    try:
        out_doc = fitz.open()
        for i, page in enumerate(doc):
            pix = page.get_pixmap(matrix=matrix, alpha=False)
            image_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(
                pix.height,
                pix.width,
                pix.n,
            )

            # Read text with EasyOCR
            results = ocr_reader.readtext(image_np, detail=0, paragraph=True)
            page_text = "\n".join(str(r).strip() for r in results if str(r).strip())
            extracted_pages.append(f"--- [Page {i + 1} / {len(doc)}] ---\n{page_text}")

            # Recreate page in output PDF with OCR text
            new_page = out_doc.new_page(width=page.rect.width, height=page.rect.height)
            new_page.insert_image(page.rect, stream=pix.tobytes(output="png"))
            if page_text:
                # Insert hidden/overlay text to make it searchable
                new_page.insert_text(
                    (36, 36),
                    page_text,
                    fontsize=8,
                    color=(0, 0, 0),
                    render_mode=3,  # Invisible text render mode in PDF
                )

        full_extracted_text = "\n\n".join(extracted_pages)
        searchable_pdf_bytes = out_doc.tobytes(deflate=True, garbage=3)
        out_doc.close()
        return full_extracted_text, searchable_pdf_bytes
    finally:
        doc.close()