"""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

import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING

from axm_bib.exceptions import ExtractionError
from axm_bib.models import ExtractionResult, FigureInfo

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import pymupdf

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

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

# Compiled regexes for markdown stripping (module-level for performance)
_RE_HEADERS = re.compile(r"^#{1,6}\s+", re.MULTILINE)
_RE_BOLD_ITALIC = re.compile(r"\*{1,3}(.*?)\*{1,3}")
_RE_LINKS = re.compile(r"\[([^\]]+)\]\([^)]+\)")
_RE_HRULES = re.compile(r"^[-*_]{3,}\s*$", re.MULTILINE)
_RE_TABLE_PIPES = re.compile(r"\|")
_RE_TABLE_SEP = re.compile(r"^[|\s:_-]+$", re.MULTILINE)
_RE_BLANK_LINES = re.compile(r"\n{3,}")


def strip_markdown(text: str) -> str:
    """Remove markdown formatting for accurate text counting.

    Strips headers, bold/italic, links, horizontal rules, and table
    formatting so that word and character counts reflect actual content.

    Args:
        text: Raw markdown string.

    Returns:
        Plain text with markdown formatting removed.
    """
    text = _RE_HEADERS.sub("", text)
    text = _RE_BOLD_ITALIC.sub(r"\1", text)
    text = _RE_LINKS.sub(r"\1", text)
    text = _RE_HRULES.sub("", text)
    text = _RE_TABLE_SEP.sub("", text)
    text = _RE_TABLE_PIPES.sub(" ", text)
    text = _RE_BLANK_LINES.sub("\n\n", text)
    return text.strip()


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.
        ExtractionError: 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 ExtractionError(msg) from exc

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

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

    if doc.page_count == 0:
        doc.close()
        msg = f"PDF has no pages: {pdf_path}"
        raise ExtractionError(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 _save_image_as_png(
    img_bytes: bytes,
    ext: str,
    out_path: Path,
) -> bool:
    """Save raw image bytes as PNG, converting if necessary.

    Handles color space conversion (CMYK → RGB) and alpha stripping
    for non-PNG source formats via pymupdf Pixmap.

    Args:
        img_bytes: Raw image bytes from PDF extraction.
        ext: Original image format extension (e.g. ``"png"``, ``"jpeg"``).
        out_path: Destination file path.

    Returns:
        True if the image was saved successfully, False on error.
    """
    import pymupdf

    if ext == "png":
        out_path.write_bytes(img_bytes)
        return True

    # Convert non-PNG to PNG via pymupdf Pixmap
    pix = pymupdf.Pixmap(img_bytes)
    # PNG only supports RGB/Gray — convert CMYK/ICC
    if pix.n - pix.alpha not in (1, 3):
        pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
    # Strip alpha if present (for clean PNG output)
    if pix.alpha:
        pix = pymupdf.Pixmap(pymupdf.csRGB, pix, 0)
    pix.save(str(out_path))
    return True


def _extract_figures_from_doc(
    doc: pymupdf.Document,
    output_dir: Path,
) -> list[FigureInfo]:
    """Extract figures from an already-open pymupdf Document.

    Internal helper — callers are responsible for opening/closing the doc.

    Args:
        doc: An open pymupdf.Document.
        output_dir: Directory to save PNG files into.

    Returns:
        List of FigureInfo for extracted figures.
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    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: BLE001
                logger.debug(
                    "Failed to extract image xref=%d on page %d",
                    xref,
                    page_num,
                )
                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"

            out_path = output_dir / filename
            ext = img_data.get("ext", "png")

            try:
                _save_image_as_png(img_data["image"], ext, out_path)
            except Exception:  # noqa: BLE001
                logger.debug("Failed to save image xref=%d as PNG", xref)
                index -= 1  # revert counter for skipped image
                continue

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

    return figures


def extract_figures(pdf_path: Path, output_dir: Path) -> tuple[list[FigureInfo], int]:
    """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:
        Tuple of (figures list, page_count).  The page count is returned
        so that callers do not need to re-open the document.

    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)

    doc = pymupdf.open(str(pdf_path))
    try:
        figures = _extract_figures_from_doc(doc, output_dir)
        page_count = doc.page_count
    finally:
        doc.close()

    return figures, page_count


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

    Opens the PDF once and passes the document handle to both
    the markdown converter and the figure extractor.

    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 pymupdf4llm

    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)

    # Open PDF once — pass handle to both extractors
    doc = _validate_pdf(pdf_path)
    try:
        # Extract markdown (pymupdf4llm accepts an open doc)
        markdown: str = pymupdf4llm.to_markdown(doc)

        # Extract figures from the same doc handle
        figures = _extract_figures_from_doc(doc, figures_dir)
        page_count = doc.page_count
    finally:
        doc.close()

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

    stripped = strip_markdown(markdown)

    return ExtractionResult(
        source_pdf=str(pdf_path),
        output_dir=str(output_dir),
        markdown_path=str(md_path),
        figures=figures,
        page_count=page_count,
        word_count=len(stripped.split()),
        char_count=len(stripped),
    )
