#!/usr/bin/env python3
"""
PDF Text Extraction Module

This module extracts text from PDF documents using pypdfium2 first, then falls back to OCR (pytesseract) if needed or if content is sparse. Extraction is limited to the first 15 pages, which can be configured in the constants.py file.

Note: PDF processing libraries are imported lazily inside functions to reduce memory footprint.
"""

from pathlib import Path
from typing import Optional
import logging
import gc
import time

from .constants import (
    OCR_DPI,
    TESSERACT_CONFIG,
    MAX_PDF_PAGES,
    MIN_LINES_PER_PAGE,
)

logger = logging.getLogger(__name__)


def _extract_text_with_pypdfium2(pdf_file_name: str) -> Optional[str]:
    """
    Extract text from PDF using pypdfium2 (Apache 2.0 licensed, fast and memory-efficient).
    Limited to first MAX_PDF_PAGES pages.

    Args:
        pdf_file_name (str): Path to the input PDF file

    Returns:
        Optional[str]: Extracted text if successful, None otherwise
    """
    try:
        # Lazy import to reduce memory footprint
        import pypdfium2 as pdfium
        
        logger.info(f"Attempting to extract text with pypdfium2 from: {pdf_file_name}")

        pdf = pdfium.PdfDocument(pdf_file_name)
        text = ""
        num_pages = min(len(pdf), MAX_PDF_PAGES)
        logger.info(f"Processing {num_pages} pages (max {MAX_PDF_PAGES})")

        for page_num in range(num_pages):
            page = pdf[page_num]
            textpage = page.get_textpage()
            page_text = textpage.get_text_range()
            
            if page_text and page_text.strip():
                text += f"--- Page {page_num + 1} ---\n{page_text}\n\n"
            
            # Close resources immediately
            textpage.close()
            page.close()

        pdf.close()

        if text.strip():
            logger.info(
                f"pypdfium2 extraction successful. Extracted {len(text)} characters."
            )
            return text
        else:
            logger.info("pypdfium2 extraction returned empty text.")
            return None

    except Exception as e:
        logger.warning(f"pypdfium2 extraction failed: {e}")
        return None


def _has_sufficient_text(text: str) -> bool:
    """
    Check if extracted text has sufficient content.
    Allows some sparse pages without triggering OCR.
    For multi-page PDFs, up to 1 sparse page is acceptable per 5 pages.
    Single-page PDFs with sparse content are still considered insufficient.

    Args:
        text (str): Extracted text to evaluate

    Returns:
        bool: True if text has sufficient content, False otherwise
    """
    if not text or not text.strip():
        return False

    # Split text into pages
    pages = text.split("--- Page ")
    if len(pages) <= 1:
        return False

    # Check each page individually (skip first element which is before first page marker)
    sparse_pages = []
    for i, page in enumerate(pages[1:], start=1):
        # Count non-empty lines in this page (skip first line which is the page number like "1 ---")
        page_lines = page.split("\n")
        lines = [
            line.strip() for line in page_lines[1:] if line.strip()
        ]  # Skip first line (page number)
        line_count = len(lines)

        if line_count < MIN_LINES_PER_PAGE:
            sparse_pages.append((i, line_count))

    total_pages = len(pages) - 1
    total_lines = sum(
        len([line.strip() for line in page.split("\n")[1:] if line.strip()])
        for page in pages[1:]
    )

    # Single-page PDFs with sparse content are likely image-based - trigger OCR
    if total_pages == 1 and sparse_pages:
        logger.info(
            f"Text quality check: {total_lines} lines on single page. "
            f"Sparse single-page PDF likely needs OCR."
        )
        return False

    # For multi-page PDFs: allow 1 sparse page per 5 pages (minimum 1 allowed)
    max_sparse_allowed = max(1, total_pages // 5)

    if len(sparse_pages) > max_sparse_allowed:
        logger.info(
            f"Text quality check: {total_lines} lines across {total_pages} pages. "
            f"Found {len(sparse_pages)} sparse page(s) with < {MIN_LINES_PER_PAGE} lines "
            f"(max allowed: {max_sparse_allowed}): "
            f"{', '.join([f'page {p} ({n_lines} lines)' for p, n_lines in sparse_pages[:3]])}"
        )
        return False
    elif sparse_pages:
        logger.info(
            f"Text quality check: {total_lines} lines across {total_pages} pages. "
            f"Found {len(sparse_pages)} sparse page(s) (within allowed limit of {max_sparse_allowed}): "
            f"{', '.join([f'page {p} ({n_lines} lines)' for p, n_lines in sparse_pages[:3]])}"
        )
        return True
    else:
        logger.info(
            f"Text quality check: {total_lines} lines across {total_pages} pages. All pages have sufficient content."
        )
        return True


def extract_pdf_text(pdf_file_name: str, text_file_name: Optional[str] = None) -> bool:
    """
    Extract text from PDF using pypdfium2 first, then fall back to OCR if needed.
    Falls back to OCR if text-based extraction returns less than MIN_LINES_PER_PAGE lines on any page.
    Limited to first MAX_PDF_PAGES pages.

    Args:
        pdf_file_name (str): Path to the input PDF file
        text_file_name (str, optional): Path to the output text file. If None, creates auto-generated name.

    Returns:
        bool: True if extraction was successful, False otherwise
    """
    try:
        # Generate output filename if not provided
        if text_file_name is None:
            base_name = Path(pdf_file_name).stem
            text_file_name = f"{base_name}_extracted.txt"

        logger.info(f"Starting text extraction from: {pdf_file_name}")

        # Try pypdfium2 first (fast and memory-efficient)
        extracted_text = _extract_text_with_pypdfium2(pdf_file_name)

        # Check if text-based extraction has sufficient content
        # If not, it's likely a scanned/image-based PDF, so fall back to OCR
        if extracted_text and not _has_sufficient_text(extracted_text):
            logger.info(
                f"Text extraction returned sparse results (< {MIN_LINES_PER_PAGE} lines/page). "
                "This may be a scanned PDF. Falling back to OCR..."
            )
            extracted_text = _extract_text_with_ocr(pdf_file_name)
        elif not extracted_text:
            # If pypdfium2 failed completely, fall back to OCR
            logger.info(
                "pypdfium2 extraction failed or returned empty text. Falling back to OCR..."
            )
            extracted_text = _extract_text_with_ocr(pdf_file_name)

        # Check if we got meaningful text
        if extracted_text and extracted_text.strip():
            logger.info(
                f"Text extraction successful. Extracted {len(extracted_text)} characters."
            )

            # Save extracted text to file
            try:
                with open(text_file_name, "w", encoding="utf-8") as f:
                    f.write(extracted_text)

                logger.info(f"Text extracted and saved to: {text_file_name}")
                return True

            except Exception as e:
                logger.error(f"Error saving text to file: {e}")
                return False
        else:
            logger.error(
                "All extraction methods (pypdfium2, OCR) returned empty text."
            )
            return False

    except Exception as e:
        logger.error(f"Error in extract_pdf_text: {e}")
        return False


def _extract_text_with_ocr(pdf_file_name: str) -> Optional[str]:
    """
    Extract text from PDF using OCR (pytesseract) with pypdfium2 rendering.
    Limited to first MAX_PDF_PAGES pages.
    Uses temporary files to minimize memory usage - images saved to disk, not RAM.

    Args:
        pdf_file_name (str): Path to the input PDF file

    Returns:
        Optional[str]: Extracted text if successful, None otherwise
    """
    import tempfile
    import os
    
    # Limit Tesseract to single thread to reduce CPU intensity
    os.environ["OMP_THREAD_LIMIT"] = "1"
    
    try:
        # Lazy import to reduce memory footprint (only load when OCR is needed)
        import pypdfium2 as pdfium
        import pytesseract
        
        logger.info(f"Starting OCR text extraction from: {pdf_file_name}")
        extracted_text = ""
        
        pdf = pdfium.PdfDocument(pdf_file_name)
        num_pages = min(len(pdf), MAX_PDF_PAGES)
        logger.info(f"Processing {num_pages} pages with OCR (max {MAX_PDF_PAGES})")

        # Process pages one at a time with temporary file storage
        for page_num in range(num_pages):
            temp_image_path = None
            try:
                logger.info(f"Converting and processing page {page_num + 1}...")
                
                page = pdf[page_num]
                
                # Render page to temporary file for OCR (not memory)
                with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
                    temp_image_path = tmp.name
                
                # Render at lower DPI with grayscale to save memory
                bitmap = page.render(
                    scale=OCR_DPI/72,  # Convert DPI to scale (72 is default)
                    rotation=0,
                    crop=(0, 0, 0, 0),
                    grayscale=True  # 66% less memory vs RGB
                )
                
                # Save bitmap to disk (not kept in memory)
                pil_image = bitmap.to_pil()
                pil_image.save(temp_image_path, 'PNG')
                pil_image.close()
                bitmap.close()
                
                # Run OCR on temp file (not in-memory image)
                text = pytesseract.image_to_string(temp_image_path, config=TESSERACT_CONFIG)
                
                if text.strip():
                    extracted_text += f"--- Page {page_num + 1} ---\n{text}\n\n"
                
                # Immediate cleanup
                if temp_image_path and os.path.exists(temp_image_path):
                    os.unlink(temp_image_path)
                    temp_image_path = None
                
                page.close()
                time.sleep(0.5)
                gc.collect()
                
            except Exception as e:
                logger.warning(f"Error processing page {page_num + 1}: {e}")
                # Clean up temp file on error
                if temp_image_path and os.path.exists(temp_image_path):
                    try:
                        os.unlink(temp_image_path)
                    except OSError as e:
                        logger.debug(f"Failed to delete temp file {temp_image_path}: {e}")
                continue

        pdf.close()

        # Check if we got meaningful text
        if extracted_text.strip():
            logger.info(
                f"OCR extraction successful. Extracted {len(extracted_text)} characters."
            )
            return extracted_text
        else:
            logger.error("OCR extraction returned empty text.")
            return None

    except Exception as e:
        logger.error(f"Error in OCR extraction: {e}")
        return None
    finally:
        # Final memory cleanup
        gc.collect()
