"""PDF extraction via PyMuPDF — text to Markdown + figures to PNG.

Architecture:
    extract_markdown(pdf)  → str (Markdown via pymupdf4llm)
    extract_figures(pdf, output_dir) → list[FigureInfo] (embedded images)
    extract_pdf(pdf, output_dir)     → ExtractionResult (orchestrator)

Note:
    pymupdf is imported lazily (inside functions) to avoid a segfault
    when pytest's assertion rewriter processes the C extension import.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from axm_bib.models import ExtractionResult, FigureInfo

if TYPE_CHECKING:
    import pymupdf

__all__ = ["extract_figures", "extract_markdown", "extract_pdf"]

_MIN_IMAGE_SIZE = 50  # Skip images smaller than 50x50 (icons, bullets)


def _validate_pdf(pdf_path: Path) -> pymupdf.Document:
    """Open and validate a PDF file.

    Args:
        pdf_path: Path to the PDF.

    Returns:
        An open pymupdf.Document.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the file is not a valid PDF or is encrypted.
    """
    import pymupdf

    if not pdf_path.exists():
        msg = f"PDF not found: {pdf_path}"
        raise FileNotFoundError(msg)

    try:
        doc = pymupdf.open(str(pdf_path))
    except Exception as exc:
        msg = f"Not a valid PDF: {pdf_path}"
        raise ValueError(msg) from exc

    if not doc.is_pdf:
        doc.close()
        msg = f"Not a valid PDF: {pdf_path}"
        raise ValueError(msg)

    if doc.is_encrypted:
        doc.close()
        msg = f"PDF is encrypted: {pdf_path}"
        raise ValueError(msg)

    if doc.page_count == 0:
        doc.close()
        msg = f"PDF has no pages: {pdf_path}"
        raise ValueError(msg)

    return doc


def extract_markdown(pdf_path: Path) -> str:
    """Convert a PDF to Markdown via pymupdf4llm.

    Args:
        pdf_path: Path to the PDF file.

    Returns:
        Markdown string with structured text.

    Raises:
        FileNotFoundError: If the PDF does not exist.
        ValueError: If the file is not a valid PDF or is encrypted.
    """
    import pymupdf4llm

    doc = _validate_pdf(pdf_path)
    doc.close()

    md: str = pymupdf4llm.to_markdown(str(pdf_path))
    return md


def extract_figures(pdf_path: Path, output_dir: Path) -> list[FigureInfo]:
    """Extract all embedded images from a PDF as numbered PNGs.

    Deduplicates by xref (an image referenced on multiple pages is
    extracted only once). Filters images smaller than 50x50 pixels.

    Args:
        pdf_path: Path to the PDF file.
        output_dir: Directory to save PNG files into.

    Returns:
        List of FigureInfo with metadata for each extracted figure.

    Raises:
        FileNotFoundError: If the PDF does not exist.
    """
    import pymupdf

    if not pdf_path.exists():
        msg = f"PDF not found: {pdf_path}"
        raise FileNotFoundError(msg)

    output_dir.mkdir(parents=True, exist_ok=True)

    doc = pymupdf.open(str(pdf_path))
    seen_xrefs: set[int] = set()
    figures: list[FigureInfo] = []
    index = 0

    for page_num in range(doc.page_count):
        page = doc.load_page(page_num)
        for img_info in page.get_images():
            xref = img_info[0]

            # Deduplicate
            if xref in seen_xrefs:
                continue
            seen_xrefs.add(xref)

            # Extract image data
            try:
                img_data = doc.extract_image(xref)
            except Exception:  # noqa: S112
                continue

            width = int(img_data["width"])
            height = int(img_data["height"])

            # Filter small images (icons, decorations)
            if width < _MIN_IMAGE_SIZE or height < _MIN_IMAGE_SIZE:
                continue

            index += 1
            filename = f"fig_{index:03d}.png"
            img_bytes: bytes = img_data["image"]

            # Save as PNG (PyMuPDF returns raw image bytes)
            out_path = output_dir / filename
            ext = img_data.get("ext", "png")

            if ext != "png":
                # Convert non-PNG to PNG via pymupdf Pixmap
                pix = pymupdf.Pixmap(img_bytes)
                pix.save(str(out_path))
            else:
                out_path.write_bytes(img_bytes)

            figures.append(
                FigureInfo(
                    index=index,
                    filename=filename,
                    page=page.number,
                    width=width,
                    height=height,
                    xref=xref,
                )
            )

    doc.close()
    return figures


def extract_pdf(
    pdf_path: Path,
    output_dir: Path | None = None,
) -> ExtractionResult:
    """Extract PDF content to a directory: Markdown + figures.

    Creates:
        output_dir/content.md          — structured Markdown text
        output_dir/figures/fig_NNN.png — extracted images

    Args:
        pdf_path: Path to the PDF file.
        output_dir: Output directory (default: {stem}/ next to PDF).

    Returns:
        ExtractionResult with paths and metadata.

    Raises:
        FileNotFoundError: If the PDF does not exist.
        ValueError: If the file is not a valid PDF or is encrypted.
    """
    import pymupdf

    if output_dir is None:
        output_dir = pdf_path.parent / pdf_path.stem

    output_dir.mkdir(parents=True, exist_ok=True)
    figures_dir = output_dir / "figures"
    figures_dir.mkdir(parents=True, exist_ok=True)

    # Extract markdown
    markdown = extract_markdown(pdf_path)

    # Write content.md
    md_path = output_dir / "content.md"
    md_path.write_text(markdown, encoding="utf-8")

    # Extract figures
    figures = extract_figures(pdf_path, figures_dir)

    # Get page count
    doc = pymupdf.open(str(pdf_path))
    page_count = doc.page_count
    doc.close()

    return ExtractionResult(
        source_pdf=str(pdf_path),
        output_dir=str(output_dir),
        markdown_path=str(md_path),
        figures=figures,
        page_count=page_count,
        char_count=len(markdown),
    )
