scripts/extract.py

#!/usr/bin/env python3
"""PDF extraction CLI script for the RAG chatbot build pipeline.

This script processes PDF files from an input directory and extracts their
content to Markdown format in an output directory. It is Step 2 of the offline
build pipeline: PDF -> Markdown extraction using pymupdf4llm.

Features:
    - Incremental extraction: Skip files that already have markdown output
    - Force overwrite: Re-extract all files with --force flag
    - Timestamp-based detection: Re-extract if PDF is newer than markdown
    - Progress reporting: Visual progress bar during batch extraction
    - Verbose mode: Show detailed file names being processed
    - Quiet mode: Suppress all output except errors
    - Statistics summary: Display extraction stats on completion

Exit Codes:
    0: Success - All files processed successfully (or skipped)
    1: Partial failure - Some files failed but some succeeded
    2: Total failure - No files processed or invalid arguments

Example Usage:
    # Basic extraction
    poetry run python scripts/extract.py data/raw/ data/processed/

    # Force overwrite existing files
    poetry run python scripts/extract.py data/raw/ data/processed/ --force

    # Verbose mode (show file names)
    poetry run python scripts/extract.py data/raw/ data/processed/ -v

    # Quiet mode (no output except errors)
    poetry run python scripts/extract.py data/raw/ data/processed/ -q

Note:
----
    This script uses lazy loading for heavy dependencies (PDFExtractor,
    MarkdownConverter) to ensure fast CLI startup times.

"""

from __future__ import annotations

import argparse
import hashlib
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

# =============================================================================
# Environment Variable Loading
# =============================================================================
# Load environment variables from .env file at script startup.
# This ensures configuration is available before it's needed.
# The .env file should be in the project root directory.
# =============================================================================
from dotenv import load_dotenv

# Find the project root (parent of scripts/ directory) and load .env from there
_PROJECT_ROOT = Path(__file__).parent.parent
_ENV_FILE = _PROJECT_ROOT / ".env"

if _ENV_FILE.exists():
    load_dotenv(_ENV_FILE)

# =============================================================================
# Type Checking Imports
# =============================================================================
# Heavy dependencies are imported lazily to ensure fast CLI startup.
# Type checkers still see the types for proper type checking.
# =============================================================================

if TYPE_CHECKING:
    from rag_chatbot.extraction import MarkdownConverter, PDFExtractor
    from rag_chatbot.extraction.models import ExtractedDocument

# =============================================================================
# Module Exports
# =============================================================================
__all__: list[str] = [
    "ExtractionStatistics",
    "parse_args",
    "run_extraction",
    "main",
]

# =============================================================================
# Constants
# =============================================================================

# Exit codes for the CLI script
EXIT_SUCCESS = 0  # All files processed successfully (or skipped)
EXIT_PARTIAL_FAILURE = 1  # Some files failed but some succeeded
EXIT_TOTAL_FAILURE = 2  # No files processed or invalid arguments

# Progress bar update function placeholder for testing
# This allows tests to track progress updates by mocking this function
_update_progress: None = None  # Will be set during extraction if needed


# =============================================================================
# Data Classes
# =============================================================================


@dataclass
class ExtractionStatistics:
    """Statistics from an extraction run.

    This dataclass tracks metrics from a batch PDF extraction operation,
    including counts of processed files and timing information.

    Attributes:
    ----------
        total : int
            Total number of PDF files found in the input directory.
            Must be non-negative.

        extracted : int
            Number of files successfully extracted to markdown.
            Must be non-negative.

        skipped : int
            Number of files skipped due to existing output (incremental mode).
            Must be non-negative.

        failed : int
            Number of files that failed to extract due to errors.
            Must be non-negative.

        total_pages : int
            Total number of pages extracted across all successful files.
            Must be non-negative.

        elapsed_seconds : float
            Total time elapsed during extraction in seconds.
            Must be non-negative.

    Example:
    -------
        >>> stats = ExtractionStatistics(
        ...     total=10,
        ...     extracted=8,
        ...     skipped=1,
        ...     failed=1,
        ...     total_pages=42,
        ...     elapsed_seconds=15.3,
        ... )
        >>> stats.extracted + stats.skipped + stats.failed == stats.total
        True

    """

    total: int
    extracted: int
    skipped: int
    failed: int
    total_pages: int
    elapsed_seconds: float

    def __post_init__(self) -> None:
        """Validate statistics values after initialization.

        Raises
        ------
            ValueError: If any count is negative.
            TypeError: If elapsed_seconds is not a number.

        """
        # Validate all counts are non-negative
        if self.total < 0:
            msg = f"total must be non-negative, got {self.total}"
            raise ValueError(msg)
        if self.extracted < 0:
            msg = f"extracted must be non-negative, got {self.extracted}"
            raise ValueError(msg)
        if self.skipped < 0:
            msg = f"skipped must be non-negative, got {self.skipped}"
            raise ValueError(msg)
        if self.failed < 0:
            msg = f"failed must be non-negative, got {self.failed}"
            raise ValueError(msg)
        if self.total_pages < 0:
            msg = f"total_pages must be non-negative, got {self.total_pages}"
            raise ValueError(msg)
        if self.elapsed_seconds < 0:
            msg = f"elapsed_seconds must be non-negative, got {self.elapsed_seconds}"
            raise ValueError(msg)


# =============================================================================
# Argument Parsing
# =============================================================================


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse command-line arguments for the extraction script.

    This function sets up the argument parser with all supported options
    and returns the parsed arguments. It handles validation of mutually
    exclusive flags (--verbose and --quiet cannot be used together).

    Args:
    ----
        argv : list[str] | None, optional
            Command-line arguments to parse. If None, uses sys.argv[1:].
            This parameter enables testing without modifying sys.argv.

    Returns:
    -------
        argparse.Namespace
            Parsed arguments with the following attributes:
            - input_dir: Path - Directory containing PDF files
            - output_dir: Path - Directory for markdown output
            - force: bool - Whether to overwrite existing files
            - verbose: bool - Whether to show detailed output
            - quiet: bool - Whether to suppress output
            - dump_raw_for: str | None - Optional PDF filename/stem to dump raw output

    Raises:
    ------
        SystemExit
            If required arguments are missing, unknown arguments are provided,
            or --verbose and --quiet are both specified.

    Example:
    -------
        >>> args = parse_args(["data/raw/", "data/processed/", "--force"])
        >>> args.input_dir
        PosixPath('data/raw')
        >>> args.force
        True

    """
    # -------------------------------------------------------------------------
    # Create the argument parser with program description
    # -------------------------------------------------------------------------
    parser = argparse.ArgumentParser(
        prog="extract.py",
        description=(
            "Extract PDF documents to Markdown format for the RAG pipeline. "
            "Processes all PDF files in the input directory and saves "
            "Markdown output to the output directory."
        ),
        epilog=(
            "Examples:\n"
            "  %(prog)s data/raw/ data/processed/           # Basic extraction\n"
            "  %(prog)s data/raw/ data/processed/ --force   # Force re-extract\n"
            "  %(prog)s data/raw/ data/processed/ -v        # Verbose output\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    # -------------------------------------------------------------------------
    # Positional Arguments
    # -------------------------------------------------------------------------
    parser.add_argument(
        "input_dir",
        type=Path,
        help="Directory containing PDF files to extract",
    )

    parser.add_argument(
        "output_dir",
        type=Path,
        help="Directory where Markdown files will be written",
    )

    # -------------------------------------------------------------------------
    # Optional Flags
    # -------------------------------------------------------------------------
    parser.add_argument(
        "--force",
        "-f",
        action="store_true",
        default=False,
        help="Force overwrite of existing Markdown files (default: skip existing)",
    )

    # Create mutually exclusive group for verbose/quiet
    # These flags cannot be used together
    output_group = parser.add_mutually_exclusive_group()

    output_group.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        default=False,
        help="Show detailed output including file names being processed",
    )

    output_group.add_argument(
        "--quiet",
        "-q",
        action="store_true",
        default=False,
        help="Suppress all output except errors (still shows summary)",
    )
    parser.add_argument(
        "--dump-raw-for",
        default=None,
        help=(
            "Dump raw pymupdf4llm markdown (pre-MarkdownConverter) for a "
            "matching PDF filename or stem to <output_dir>/<stem>.raw.md"
        ),
    )

    # -------------------------------------------------------------------------
    # Parse and return arguments
    # -------------------------------------------------------------------------
    return parser.parse_args(argv)


# =============================================================================
# Extraction Logic
# =============================================================================


def _should_extract(
    pdf_path: Path,
    md_path: Path,
    force: bool,
) -> bool:
    """Determine if a PDF file should be extracted.

    This function implements the incremental extraction logic by checking:
    1. If --force is set, always extract
    2. If markdown output doesn't exist, extract
    3. If PDF is newer than markdown, extract (re-extract modified files)
    4. Otherwise, skip (markdown is up-to-date)

    Args:
    ----
        pdf_path : Path
            Path to the PDF file to potentially extract.
        md_path : Path
            Path where the markdown output would be written.
        force : bool
            Whether to force extraction regardless of existing output.

    Returns:
    -------
        bool
            True if the file should be extracted, False if it should be skipped.

    """
    # Force mode: always extract
    if force:
        return True

    # No existing output: extract
    if not md_path.exists():
        return True

    # Compare modification times: extract if PDF is newer
    pdf_mtime = pdf_path.stat().st_mtime
    md_mtime = md_path.stat().st_mtime

    return pdf_mtime > md_mtime


def _get_extractor() -> PDFExtractor:
    """Lazily load and return a PDFExtractor instance.

    This function handles the lazy import of the PDFExtractor class to avoid
    loading heavy dependencies (pymupdf, pymupdf4llm) at module import time.

    Returns
    -------
        PDFExtractor
            A configured PDF extractor instance.

    """
    from rag_chatbot.extraction import PDFExtractor

    return PDFExtractor()


def _get_converter() -> MarkdownConverter:
    """Lazily load and return a MarkdownConverter instance.

    This function handles the lazy import of the MarkdownConverter class.

    Returns
    -------
        MarkdownConverter
            A configured Markdown converter instance.

    """
    from rag_chatbot.extraction import MarkdownConverter

    return MarkdownConverter()


def _print_summary(stats: ExtractionStatistics, quiet: bool) -> None:  # noqa: ARG001
    """Print the extraction summary statistics.

    Displays a formatted summary of the extraction run including counts
    of extracted, skipped, and failed files, plus timing information.

    Args:
    ----
        stats : ExtractionStatistics
            The statistics from the extraction run.
        quiet : bool
            Unused - summary is always printed even in quiet mode.
            Kept for API consistency with run_extraction parameters.

    Note:
    ----
        The quiet parameter is intentionally unused because the summary
        should always be shown to the user, even in quiet mode. Only
        the progress bar is suppressed in quiet mode.

    """
    # Always print summary, even in quiet mode
    # Only progress is suppressed in quiet mode
    print()
    print("Extraction Complete")
    print("=" * 39)
    print(f"Total PDF files:     {stats.total:>4}")
    print(f"Extracted:           {stats.extracted:>4}")
    print(f"Skipped (existing):  {stats.skipped:>4}")
    print(f"Failed:              {stats.failed:>4}")
    print(f"Total pages:         {stats.total_pages:>4}")
    print(f"Elapsed time:      {stats.elapsed_seconds:.2f}s")
    print("=" * 39)


def run_extraction(  # noqa: PLR0912, PLR0915
    input_dir: Path,
    output_dir: Path,
    force: bool,
    verbose: bool,
    quiet: bool,
    dump_raw_for: str | None = None,
) -> ExtractionStatistics:
    """Run the PDF extraction process on all PDF files in the input directory.

    This function is the core extraction logic. It:
    1. Finds all PDF files in the input directory
    2. Determines which files need extraction (incremental or force)
    3. Extracts each PDF using PDFExtractor
    4. Applies MarkdownConverter for normalization
    5. Writes output to the output directory
    6. Tracks and returns extraction statistics

    Args:
    ----
        input_dir : Path
            Directory containing PDF files to extract. Must exist.
        output_dir : Path
            Directory where Markdown files will be written. Created if needed.
        force : bool
            If True, overwrite existing markdown files. If False, skip files
            that already have up-to-date markdown output.
        verbose : bool
            If True, print detailed information including file names.
        quiet : bool
            If True, suppress progress bar (but still print summary).
        dump_raw_for : str | None
            Optional PDF filename or stem to dump raw markdown for.

    Returns:
    -------
        ExtractionStatistics
            Statistics about the extraction run including counts and timing.

    Note:
    ----
        The function handles errors gracefully, continuing to process remaining
        files if one fails. Failed files are logged and counted in statistics.

    """
    # -------------------------------------------------------------------------
    # Start timing
    # -------------------------------------------------------------------------
    start_time = time.perf_counter()

    # -------------------------------------------------------------------------
    # Create output directory if it doesn't exist
    # -------------------------------------------------------------------------
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
    except PermissionError:
        # Cannot create output directory - this is a fatal error
        print(
            f"Error: Permission denied creating output directory: {output_dir}",
            file=sys.stderr,
        )
        elapsed = time.perf_counter() - start_time
        return ExtractionStatistics(
            total=0,
            extracted=0,
            skipped=0,
            failed=0,
            total_pages=0,
            elapsed_seconds=elapsed,
        )

    # -------------------------------------------------------------------------
    # Find all PDF files in input directory (non-recursive)
    # -------------------------------------------------------------------------
    pdf_files = sorted(input_dir.glob("*.pdf"))
    total_files = len(pdf_files)

    # -------------------------------------------------------------------------
    # Initialize statistics
    # -------------------------------------------------------------------------
    extracted_count = 0
    skipped_count = 0
    failed_count = 0
    total_pages = 0

    # -------------------------------------------------------------------------
    # Handle empty directory case
    # -------------------------------------------------------------------------
    if total_files == 0:
        elapsed = time.perf_counter() - start_time
        return ExtractionStatistics(
            total=0,
            extracted=0,
            skipped=0,
            failed=0,
            total_pages=0,
            elapsed_seconds=elapsed,
        )

    # -------------------------------------------------------------------------
    # Lazy load the extractor and converter
    # -------------------------------------------------------------------------
    extractor = _get_extractor()
    converter = _get_converter()

    # -------------------------------------------------------------------------
    # Setup progress bar (unless quiet mode)
    # -------------------------------------------------------------------------
    if not quiet:
        try:
            from tqdm import tqdm  # type: ignore[import-untyped]

            progress_bar = tqdm(
                pdf_files,
                desc="Extracting",
                unit="file",
                disable=False,
            )
        except ImportError:
            # tqdm not available, fall back to simple iteration
            progress_bar = pdf_files
            if not quiet:
                print(f"Processing {total_files} PDF files...")
    else:
        progress_bar = pdf_files

    # -------------------------------------------------------------------------
    # Process each PDF file
    # -------------------------------------------------------------------------
    raw_target = (dump_raw_for or "").strip()
    raw_target_lower = raw_target.lower()
    raw_target_is_filename = raw_target_lower.endswith(".pdf")

    for idx, pdf_path in enumerate(progress_bar):
        # Determine output path
        md_filename = pdf_path.stem + ".md"
        md_path = output_dir / md_filename
        pdf_name_lower = pdf_path.name.lower()
        pdf_stem_lower = pdf_path.stem.lower()
        raw_dump_match = False
        if raw_target_lower:
            if raw_target_is_filename:
                raw_dump_match = pdf_name_lower == raw_target_lower
            else:
                raw_dump_match = raw_target_lower in {pdf_name_lower, pdf_stem_lower}

        # Check if extraction is needed
        if not _should_extract(pdf_path, md_path, force):
            skipped_count += 1
            if verbose:
                print(f"  Skipping (up-to-date): {pdf_path.name}")

            # Call progress update for testing
            if _update_progress is not None:
                _update_progress(idx, total_files)

            continue

        # Attempt extraction
        try:
            if verbose:
                print(f"  Extracting: {pdf_path.name}")

            # Extract PDF content
            document: ExtractedDocument = extractor.extract(pdf_path)

            # Get markdown content and apply converter
            raw_markdown = document.to_markdown()
            if raw_dump_match:
                raw_dump_path = output_dir / f"{pdf_path.stem}.raw.md"
                raw_dump_path.write_text(raw_markdown, encoding="utf-8")
                if verbose:
                    print(f"    Wrote raw markdown: {raw_dump_path}")
            if verbose and raw_dump_match:
                raw_hash = hashlib.sha256(raw_markdown.encode("utf-8")).hexdigest()[:12]
                raw_underscore_count = raw_markdown.count("_")
                print(
                    "    Raw checksum/underscores:",
                    raw_hash,
                    f"underscores={raw_underscore_count}",
                )
            clean_markdown = converter.convert(raw_markdown)
            if verbose and raw_dump_match:
                clean_hash = hashlib.sha256(clean_markdown.encode("utf-8")).hexdigest()[
                    :12
                ]
                clean_underscore_count = clean_markdown.count("_")
                print(
                    "    Clean checksum/underscores:",
                    clean_hash,
                    f"underscores={clean_underscore_count}",
                )

            # Write output file
            md_path.write_text(clean_markdown, encoding="utf-8")

            # Update statistics
            extracted_count += 1
            total_pages += document.page_count

            if verbose:
                print(
                    f"    Extracted {document.page_count} pages, "
                    f"{document.total_tables} tables, "
                    f"{document.total_images} images"
                )

        except PermissionError as e:
            # Handle permission errors
            failed_count += 1
            print(f"Error: Permission denied for {pdf_path.name}: {e}", file=sys.stderr)

        except Exception as e:
            # Handle other errors (corrupted PDF, extraction failure, etc.)
            failed_count += 1
            print(f"Error: Failed to extract {pdf_path.name}: {e}", file=sys.stderr)

        # Call progress update for testing
        if _update_progress is not None:
            _update_progress(idx, total_files)

    # -------------------------------------------------------------------------
    # Calculate elapsed time and create statistics
    # -------------------------------------------------------------------------
    elapsed = time.perf_counter() - start_time

    stats = ExtractionStatistics(
        total=total_files,
        extracted=extracted_count,
        skipped=skipped_count,
        failed=failed_count,
        total_pages=total_pages,
        elapsed_seconds=elapsed,
    )

    # -------------------------------------------------------------------------
    # Print summary (always shown, even in quiet mode)
    # -------------------------------------------------------------------------
    _print_summary(stats, quiet)

    return stats


# =============================================================================
# Main Entry Point
# =============================================================================


def main(argv: list[str] | None = None) -> int:  # noqa: PLR0911
    """Execute the extraction CLI script.

    This function orchestrates the entire extraction process:
    1. Parses command-line arguments
    2. Validates input directory existence
    3. Runs the extraction process
    4. Prints summary statistics
    5. Returns appropriate exit code

    Args:
    ----
        argv : list[str] | None, optional
            Command-line arguments to parse. If None, uses sys.argv[1:].

    Returns:
    -------
        int
            Exit code indicating success or failure:
            - 0: Success (all files processed or skipped)
            - 1: Partial failure (some files failed)
            - 2: Total failure (no files processed or invalid input)

    Example:
    -------
        >>> exit_code = main(["data/raw/", "data/processed/"])
        >>> exit_code
        0

    """
    # -------------------------------------------------------------------------
    # Parse arguments
    # -------------------------------------------------------------------------
    args = parse_args(argv)

    # -------------------------------------------------------------------------
    # Validate input directory
    # -------------------------------------------------------------------------
    if not args.input_dir.exists():
        print(
            f"Error: Input directory does not exist: {args.input_dir}",
            file=sys.stderr,
        )
        return EXIT_TOTAL_FAILURE

    if not args.input_dir.is_dir():
        print(
            f"Error: Input path is not a directory: {args.input_dir}",
            file=sys.stderr,
        )
        return EXIT_TOTAL_FAILURE

    # -------------------------------------------------------------------------
    # Check for PDF files
    # -------------------------------------------------------------------------
    pdf_files = list(args.input_dir.glob("*.pdf"))
    if not pdf_files:
        print(
            f"Error: No PDF files found in {args.input_dir}",
            file=sys.stderr,
        )
        return EXIT_TOTAL_FAILURE

    # -------------------------------------------------------------------------
    # Run extraction (prints summary internally)
    # -------------------------------------------------------------------------
    stats = run_extraction(
        input_dir=args.input_dir,
        output_dir=args.output_dir,
        force=args.force,
        verbose=args.verbose,
        quiet=args.quiet,
        dump_raw_for=args.dump_raw_for,
    )

    # -------------------------------------------------------------------------
    # Determine exit code
    # -------------------------------------------------------------------------
    # Check for early failure (couldn't even start extraction)
    # This happens when e.g. output directory creation fails
    if stats.total == 0 and len(pdf_files) > 0:
        return EXIT_TOTAL_FAILURE

    # Success: all files processed (extracted or skipped), none failed
    if stats.failed == 0:
        return EXIT_SUCCESS

    # Partial failure: some files succeeded, some failed
    if stats.extracted > 0 or stats.skipped > 0:
        return EXIT_PARTIAL_FAILURE

    # Total failure: no files succeeded
    return EXIT_TOTAL_FAILURE


# =============================================================================
# Script Entry Point
# =============================================================================

if __name__ == "__main__":
    sys.exit(main())