"""
PDF Text Extraction Module

Extracts text tokens with bounding boxes from text-layer PDFs.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Generator, Optional
import fitz  # PyMuPDF

from .detector import is_text_pdf as _is_text_pdf_standalone


@dataclass
class Token:
    """Represents a text token with its bounding box."""
    text: str
    bbox: tuple[float, float, float, float]  # (x0, y0, x1, y1)
    page_no: int

    @property
    def x0(self) -> float:
        return self.bbox[0]

    @property
    def y0(self) -> float:
        return self.bbox[1]

    @property
    def x1(self) -> float:
        return self.bbox[2]

    @property
    def y1(self) -> float:
        return self.bbox[3]

    @property
    def width(self) -> float:
        return self.x1 - self.x0

    @property
    def height(self) -> float:
        return self.y1 - self.y0

    @property
    def center(self) -> tuple[float, float]:
        return ((self.x0 + self.x1) / 2, (self.y0 + self.y1) / 2)


class PDFDocument:
    """
    Context manager for efficient PDF document handling.

    Caches the open document handle to avoid repeated open/close cycles.
    Use this when you need to perform multiple operations on the same PDF.
    """

    def __init__(self, pdf_path: str | Path):
        self.pdf_path = Path(pdf_path)
        self._doc: Optional[fitz.Document] = None
        self._dimensions_cache: dict[int, tuple[float, float]] = {}

    def __enter__(self) -> 'PDFDocument':
        self._doc = fitz.open(self.pdf_path)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._doc:
            self._doc.close()
            self._doc = None

    @property
    def doc(self) -> fitz.Document:
        if self._doc is None:
            raise RuntimeError("PDFDocument must be used within a context manager")
        return self._doc

    @property
    def page_count(self) -> int:
        return len(self.doc)

    def is_text_pdf(self, min_chars: int = 30) -> bool:
        """
        Check if PDF has extractable AND READABLE text layer.

        Uses the improved detection from detector.py that also checks
        for garbled text (custom font encoding issues).
        """
        return _is_text_pdf_standalone(self.pdf_path, min_chars)

    def get_page_dimensions(self, page_no: int = 0) -> tuple[float, float]:
        """Get page dimensions in points (cached)."""
        if page_no not in self._dimensions_cache:
            page = self.doc[page_no]
            rect = page.rect
            self._dimensions_cache[page_no] = (rect.width, rect.height)
        return self._dimensions_cache[page_no]

    def get_render_dimensions(self, page_no: int = 0, dpi: int = 300) -> tuple[int, int]:
        """Get rendered image dimensions in pixels."""
        width, height = self.get_page_dimensions(page_no)
        zoom = dpi / 72
        return int(width * zoom), int(height * zoom)

    def extract_text_tokens(self, page_no: int) -> Generator[Token, None, None]:
        """Extract text tokens from a specific page."""
        page = self.doc[page_no]

        text_dict = page.get_text("dict")
        tokens_found = False

        for block in text_dict.get("blocks", []):
            if block.get("type") != 0:
                continue

            for line in block.get("lines", []):
                for span in line.get("spans", []):
                    text = span.get("text", "").strip()
                    if not text:
                        continue

                    bbox = span.get("bbox")
                    if bbox and all(abs(b) < 1e9 for b in bbox):
                        tokens_found = True
                        yield Token(
                            text=text,
                            bbox=tuple(bbox),
                            page_no=page_no
                        )

        # Fallback: if dict mode failed, use words mode
        if not tokens_found:
            words = page.get_text("words")
            for word_info in words:
                x0, y0, x1, y1, text, *_ = word_info
                text = text.strip()
                if text:
                    yield Token(
                        text=text,
                        bbox=(x0, y0, x1, y1),
                        page_no=page_no
                    )

    def render_page(self, page_no: int, output_path: Path, dpi: int = 300) -> Path:
        """Render a page to an image file."""
        zoom = dpi / 72
        matrix = fitz.Matrix(zoom, zoom)

        page = self.doc[page_no]
        pix = page.get_pixmap(matrix=matrix)

        output_path.parent.mkdir(parents=True, exist_ok=True)
        pix.save(str(output_path))

        return output_path

    def render_all_pages(
        self,
        output_dir: Path,
        dpi: int = 300
    ) -> Generator[tuple[int, Path], None, None]:
        """Render all pages to images."""
        output_dir.mkdir(parents=True, exist_ok=True)
        pdf_name = self.pdf_path.stem

        zoom = dpi / 72
        matrix = fitz.Matrix(zoom, zoom)

        for page_no in range(self.page_count):
            page = self.doc[page_no]
            pix = page.get_pixmap(matrix=matrix)

            image_path = output_dir / f"{pdf_name}_page_{page_no:03d}.png"
            pix.save(str(image_path))

            yield page_no, image_path


def extract_text_tokens(
    pdf_path: str | Path,
    page_no: int | None = None
) -> Generator[Token, None, None]:
    """
    Extract text tokens with bounding boxes from PDF.

    Args:
        pdf_path: Path to the PDF file
        page_no: Specific page to extract (None for all pages)

    Yields:
        Token objects with text and bbox
    """
    doc = fitz.open(pdf_path)

    pages_to_process = [page_no] if page_no is not None else range(len(doc))

    for pg_no in pages_to_process:
        page = doc[pg_no]

        # Get text with position info using "dict" mode
        text_dict = page.get_text("dict")

        tokens_found = False
        for block in text_dict.get("blocks", []):
            if block.get("type") != 0:  # Skip non-text blocks
                continue

            for line in block.get("lines", []):
                for span in line.get("spans", []):
                    text = span.get("text", "").strip()
                    if not text:
                        continue

                    bbox = span.get("bbox")
                    # Check for corrupted bbox (overflow values)
                    if bbox and all(abs(b) < 1e9 for b in bbox):
                        tokens_found = True
                        yield Token(
                            text=text,
                            bbox=tuple(bbox),
                            page_no=pg_no
                        )

        # Fallback: if dict mode failed, use words mode
        if not tokens_found:
            words = page.get_text("words")
            for word_info in words:
                x0, y0, x1, y1, text, *_ = word_info
                text = text.strip()
                if text:
                    yield Token(
                        text=text,
                        bbox=(x0, y0, x1, y1),
                        page_no=pg_no
                    )

    doc.close()


def extract_words(
    pdf_path: str | Path,
    page_no: int | None = None
) -> Generator[Token, None, None]:
    """
    Extract individual words with bounding boxes.

    Uses PyMuPDF's word extraction which splits text into words.
    """
    doc = fitz.open(pdf_path)

    pages_to_process = [page_no] if page_no is not None else range(len(doc))

    for pg_no in pages_to_process:
        page = doc[pg_no]

        # get_text("words") returns list of (x0, y0, x1, y1, word, block_no, line_no, word_no)
        words = page.get_text("words")

        for word_info in words:
            x0, y0, x1, y1, text, *_ = word_info
            text = text.strip()
            if text:
                yield Token(
                    text=text,
                    bbox=(x0, y0, x1, y1),
                    page_no=pg_no
                )

    doc.close()


def extract_lines(
    pdf_path: str | Path,
    page_no: int | None = None
) -> Generator[Token, None, None]:
    """
    Extract text lines with bounding boxes.
    """
    doc = fitz.open(pdf_path)

    pages_to_process = [page_no] if page_no is not None else range(len(doc))

    for pg_no in pages_to_process:
        page = doc[pg_no]
        text_dict = page.get_text("dict")

        for block in text_dict.get("blocks", []):
            if block.get("type") != 0:
                continue

            for line in block.get("lines", []):
                spans = line.get("spans", [])
                if not spans:
                    continue

                # Combine all spans in the line
                line_text = " ".join(s.get("text", "") for s in spans).strip()
                if not line_text:
                    continue

                # Get line bbox from all spans
                x0 = min(s["bbox"][0] for s in spans)
                y0 = min(s["bbox"][1] for s in spans)
                x1 = max(s["bbox"][2] for s in spans)
                y1 = max(s["bbox"][3] for s in spans)

                yield Token(
                    text=line_text,
                    bbox=(x0, y0, x1, y1),
                    page_no=pg_no
                )

    doc.close()


def get_page_dimensions(pdf_path: str | Path, page_no: int = 0) -> tuple[float, float]:
    """Get the dimensions of a PDF page in points."""
    doc = fitz.open(pdf_path)
    page = doc[page_no]
    rect = page.rect
    doc.close()
    return rect.width, rect.height
