"""
PDF Document Extraction

Extracts text, tables, and structure from PDF documents.
Uses pymupdf (fitz) for text extraction and pdfplumber for table extraction.
"""

import logging
import re
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

logger = logging.getLogger(__name__)


class PDFType(Enum):
    """Type of PDF based on content extraction method."""

    TEXT_BASED = "text_based"  # Native text layer
    SCANNED = "scanned"  # Image-only, requires OCR
    HYBRID = "hybrid"  # Mix of text and scanned pages


@dataclass
class ExtractedTable:
    """Represents a table extracted from a PDF."""

    page_number: int
    table_index: int  # Index on the page
    headers: list[str]
    rows: list[list[str]]
    bbox: tuple[float, float, float, float] | None = None  # Bounding box
    caption: str | None = None

    @property
    def num_rows(self) -> int:
        return len(self.rows)

    @property
    def num_cols(self) -> int:
        return len(self.headers) if self.headers else (len(self.rows[0]) if self.rows else 0)

    def to_markdown(self) -> str:
        """Convert table to markdown format."""
        if not self.headers and not self.rows:
            return ""

        lines = []

        # Caption if present
        if self.caption:
            lines.append(f"**{self.caption}**\n")

        # Headers
        headers = self.headers or ([""] * self.num_cols)
        lines.append("| " + " | ".join(headers) + " |")
        lines.append("| " + " | ".join(["---"] * len(headers)) + " |")

        # Rows
        for row in self.rows:
            # Ensure row has correct number of columns
            padded_row = row + [""] * (len(headers) - len(row))
            lines.append("| " + " | ".join(padded_row[:len(headers)]) + " |")

        return "\n".join(lines)


@dataclass
class ExtractedFigure:
    """Represents a figure extracted from a PDF."""

    page_number: int
    figure_index: int
    bbox: tuple[float, float, float, float]  # Bounding box
    caption: str | None = None
    image_data: bytes | None = None
    image_format: str = "png"


@dataclass
class ExtractedSection:
    """Represents a document section with hierarchy."""

    title: str
    level: int  # Heading level (1=top, 2=subsection, etc.)
    page_number: int
    content: str = ""
    subsections: list["ExtractedSection"] = field(default_factory=list)


@dataclass
class PDFExtractionResult:
    """Result of PDF extraction."""

    file_path: Path
    pdf_type: PDFType
    total_pages: int
    text_content: str
    sections: list[ExtractedSection] = field(default_factory=list)
    tables: list[ExtractedTable] = field(default_factory=list)
    figures: list[ExtractedFigure] = field(default_factory=list)
    metadata: dict = field(default_factory=dict)
    extraction_warnings: list[str] = field(default_factory=list)
    pages_requiring_ocr: list[int] = field(default_factory=list)

    @property
    def has_tables(self) -> bool:
        return len(self.tables) > 0

    @property
    def has_figures(self) -> bool:
        return len(self.figures) > 0


class DocumentTypeDetector:
    """Detects if a PDF is text-based or scanned."""

    # Minimum characters per page to consider it text-based
    MIN_CHARS_PER_PAGE = 100

    def detect(self, file_path: str | Path) -> PDFType:
        """
        Detect if PDF is text-based or scanned.

        Args:
            file_path: Path to PDF file

        Returns:
            PDFType indicating the document type
        """
        try:
            import fitz  # pymupdf
        except ImportError:
            logger.error("pymupdf not installed. Run: pip install pymupdf")
            raise

        file_path = Path(file_path)

        doc = fitz.open(file_path)
        text_pages = 0
        scanned_pages = 0

        try:
            for page in doc:
                text = page.get_text()
                char_count = len(text.strip())

                if char_count >= self.MIN_CHARS_PER_PAGE:
                    text_pages += 1
                else:
                    scanned_pages += 1
        finally:
            doc.close()

        if scanned_pages == 0:
            return PDFType.TEXT_BASED
        elif text_pages == 0:
            return PDFType.SCANNED
        else:
            return PDFType.HYBRID


class PDFExtractor:
    """
    Extracts content from PDF documents.

    Features:
    - Text extraction with structure preservation
    - Table extraction with header detection
    - Figure/image extraction with captions
    - Multi-column layout handling
    - Automatic OCR fallback for scanned pages

    Example:
        extractor = PDFExtractor()
        result = extractor.extract("document.pdf")
        print(result.text_content)
        for table in result.tables:
            print(table.to_markdown())
    """

    def __init__(
        self,
        extract_tables: bool = True,
        extract_figures: bool = True,
        extract_structure: bool = True,
        ocr_fallback: bool = True,
    ):
        """
        Initialize PDF extractor.

        Args:
            extract_tables: Whether to extract tables
            extract_figures: Whether to extract figures
            extract_structure: Whether to extract document structure
            ocr_fallback: Whether to use OCR for scanned pages
        """
        self.extract_tables = extract_tables
        self.extract_figures = extract_figures
        self.extract_structure = extract_structure
        self.ocr_fallback = ocr_fallback
        self._type_detector = DocumentTypeDetector()

    def extract(self, file_path: str | Path) -> PDFExtractionResult:
        """
        Extract content from a PDF file.

        Args:
            file_path: Path to PDF file

        Returns:
            PDFExtractionResult containing extracted content
        """
        try:
            import fitz  # pymupdf
        except ImportError:
            logger.error("pymupdf not installed. Run: pip install pymupdf")
            raise

        file_path = Path(file_path)

        if not file_path.exists():
            raise FileNotFoundError(f"PDF file not found: {file_path}")

        # Detect document type
        pdf_type = self._type_detector.detect(file_path)
        logger.info(f"Detected PDF type: {pdf_type.value}")

        doc = fitz.open(file_path)
        warnings: list[str] = []
        pages_requiring_ocr: list[int] = []

        try:
            # Extract metadata
            metadata = self._extract_metadata(doc)

            # Extract text content
            text_content, pages_requiring_ocr = self._extract_text(doc, pdf_type)

            # Extract tables
            tables = []
            if self.extract_tables:
                tables = self._extract_tables(file_path)

            # Extract figures
            figures = []
            if self.extract_figures:
                figures = self._extract_figures(doc)

            # Extract document structure
            sections = []
            if self.extract_structure:
                sections = self._extract_structure(doc)

            return PDFExtractionResult(
                file_path=file_path,
                pdf_type=pdf_type,
                total_pages=len(doc),
                text_content=text_content,
                sections=sections,
                tables=tables,
                figures=figures,
                metadata=metadata,
                extraction_warnings=warnings,
                pages_requiring_ocr=pages_requiring_ocr,
            )

        finally:
            doc.close()

    def _extract_metadata(self, doc) -> dict:
        """Extract PDF metadata."""
        metadata = doc.metadata
        return {
            "title": metadata.get("title", ""),
            "author": metadata.get("author", ""),
            "subject": metadata.get("subject", ""),
            "keywords": metadata.get("keywords", ""),
            "creator": metadata.get("creator", ""),
            "producer": metadata.get("producer", ""),
            "creation_date": metadata.get("creationDate", ""),
            "modification_date": metadata.get("modDate", ""),
        }

    def _extract_text(self, doc, pdf_type: PDFType) -> tuple[str, list[int]]:
        """Extract text from all pages."""
        all_text = []
        pages_requiring_ocr = []

        for page_num, page in enumerate(doc):
            text = page.get_text("text")

            # Check if page needs OCR
            if len(text.strip()) < DocumentTypeDetector.MIN_CHARS_PER_PAGE:
                pages_requiring_ocr.append(page_num + 1)

                if self.ocr_fallback:
                    # OCR will be handled by the OCR module
                    logger.debug(f"Page {page_num + 1} requires OCR")
                    text = f"\n[Page {page_num + 1} requires OCR]\n"

            # Add page marker
            all_text.append(f"\n--- Page {page_num + 1} ---\n")
            all_text.append(text)

        return "\n".join(all_text), pages_requiring_ocr

    def _extract_tables(self, file_path: Path) -> list[ExtractedTable]:
        """Extract tables using pdfplumber."""
        try:
            import pdfplumber
        except ImportError:
            logger.warning("pdfplumber not installed. Table extraction disabled.")
            return []

        tables = []

        with pdfplumber.open(file_path) as pdf:
            for page_num, page in enumerate(pdf.pages):
                page_tables = page.extract_tables()

                for table_idx, table in enumerate(page_tables):
                    if not table or len(table) < 2:
                        continue

                    # First row as headers
                    headers = [str(cell or "") for cell in table[0]]

                    # Remaining rows as data
                    rows = []
                    for row in table[1:]:
                        rows.append([str(cell or "") for cell in row])

                    # Try to find caption (text above table)
                    caption = self._find_table_caption(page, table_idx)

                    tables.append(ExtractedTable(
                        page_number=page_num + 1,
                        table_index=table_idx,
                        headers=headers,
                        rows=rows,
                        caption=caption,
                    ))

        logger.info(f"Extracted {len(tables)} tables")
        return tables

    def _find_table_caption(self, page, table_idx: int) -> str | None:
        """Try to find table caption text."""
        # Look for "Table X" pattern in page text
        text = page.extract_text() or ""
        pattern = rf"Table\s+{table_idx + 1}[.:]\s*([^\n]+)"
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return f"Table {table_idx + 1}: {match.group(1).strip()}"
        return None

    def _extract_figures(self, doc) -> list[ExtractedFigure]:
        """Extract figures/images from PDF."""
        figures = []

        for page_num, page in enumerate(doc):
            # Use full=True to get complete image info needed for get_image_bbox()
            image_list = page.get_images(full=True)

            for img_idx, img in enumerate(image_list):
                xref = img[0]

                try:
                    # Get image data
                    base_image = doc.extract_image(xref)
                    image_data = base_image["image"]
                    image_format = base_image["ext"]

                    # Get bounding box - may fail for inline images or complex cases
                    try:
                        bbox = page.get_image_bbox(img)
                        bbox_tuple = (bbox.x0, bbox.y0, bbox.x1, bbox.y1)
                    except (ValueError, RuntimeError):
                        # Some images (e.g., inline images, smasks) don't have extractable bboxes
                        # Use a default bbox covering the page
                        rect = page.rect
                        bbox_tuple = (rect.x0, rect.y0, rect.x1, rect.y1)

                    # Try to find caption
                    caption = self._find_figure_caption(page, img_idx)

                    figures.append(ExtractedFigure(
                        page_number=page_num + 1,
                        figure_index=img_idx,
                        bbox=bbox_tuple,
                        caption=caption,
                        image_data=image_data,
                        image_format=image_format,
                    ))
                except Exception as e:
                    # Only log actual extraction failures, not expected limitations
                    error_msg = str(e).lower()
                    if "smask" not in error_msg and "mask" not in error_msg:
                        logger.warning(f"Failed to extract image {img_idx} on page {page_num + 1}: {e}")

        logger.info(f"Extracted {len(figures)} figures")
        return figures

    def _find_figure_caption(self, page, fig_idx: int) -> str | None:
        """Try to find figure caption text."""
        text = page.get_text()
        pattern = rf"Figure\s+{fig_idx + 1}[.:]\s*([^\n]+)"
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return f"Figure {fig_idx + 1}: {match.group(1).strip()}"
        return None

    def _extract_structure(self, doc) -> list[ExtractedSection]:
        """Extract document structure from TOC or headings."""
        sections = []

        # Try to get table of contents
        toc = doc.get_toc()

        if toc:
            for level, title, page_num in toc:
                sections.append(ExtractedSection(
                    title=title,
                    level=level,
                    page_number=page_num,
                ))
        else:
            # Fall back to detecting headings from text
            sections = self._detect_headings(doc)

        return sections

    def _detect_headings(self, doc) -> list[ExtractedSection]:
        """Detect headings from text formatting."""
        sections = []

        # Common heading patterns
        heading_patterns = [
            (1, r"^CHAPTER\s+\d+"),
            (1, r"^SECTION\s+\d+"),
            (1, r"^\d+\.\s+[A-Z]"),
            (2, r"^\d+\.\d+\s+[A-Z]"),
            (3, r"^\d+\.\d+\.\d+\s+[A-Z]"),
        ]

        for page_num, page in enumerate(doc):
            text = page.get_text()
            lines = text.split("\n")

            for line in lines:
                line = line.strip()
                if not line:
                    continue

                for level, pattern in heading_patterns:
                    if re.match(pattern, line):
                        sections.append(ExtractedSection(
                            title=line,
                            level=level,
                            page_number=page_num + 1,
                        ))
                        break

        return sections


def extract_pdf(
    file_path: str | Path,
    **kwargs,
) -> PDFExtractionResult:
    """
    Convenience function to extract content from a PDF.

    Args:
        file_path: Path to PDF file
        **kwargs: Arguments passed to PDFExtractor

    Returns:
        PDFExtractionResult
    """
    extractor = PDFExtractor(**kwargs)
    return extractor.extract(file_path)
