# pyright: ignore
#!/usr/bin/env python3
"""
ProcessPoolExecutor version of page-by-page PDF converter.
Tests if separate process GPU contexts improve utilization.
"""

import argparse
import sys
import tempfile
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import fitz  # PyMuPDF
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, TableFormerMode, VlmPipelineOptions
from docling.datamodel.vlm_model_specs import GRANITEDOCLING_TRANSFORMERS
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline
from docling.pipeline.vlm_pipeline import VlmPipeline
from loguru import logger

# Configure loguru
logger.remove()  # Remove default handler
logger.add(
    sys.stderr, format="<green>{time:HH:mm:ss}</green> | <level>{level}</level> | {message}", level="INFO"
)

# Suppress verbose library logging
import logging

logging.getLogger("docling").setLevel(logging.WARNING)
logging.getLogger("torch").setLevel(logging.WARNING)
logging.getLogger("transformers").setLevel(logging.WARNING)
logging.getLogger("accelerate").setLevel(logging.WARNING)


def init_converters():
    """Initialize converters - called once per process"""
    # Regular pipeline configuration (from user's working script)
    regular_pipeline_options = PdfPipelineOptions()
    regular_pipeline_options.do_table_structure = True
    regular_pipeline_options.table_structure_options.do_cell_matching = True
    regular_pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE
    regular_pipeline_options.do_ocr = True
    regular_pipeline_options.images_scale = 2.0
    regular_pipeline_options.generate_page_images = True

    regular_format_options = PdfFormatOption(
        pipeline_cls=StandardPdfPipeline, pipeline_options=regular_pipeline_options
    )

    regular_converter = DocumentConverter(format_options={InputFormat.PDF: regular_format_options})

    # VLM pipeline configuration
    vlm_pipeline_options = VlmPipelineOptions(vlm_options=GRANITEDOCLING_TRANSFORMERS)

    vlm_format_options = PdfFormatOption(pipeline_cls=VlmPipeline, pipeline_options=vlm_pipeline_options)

    vlm_converter = DocumentConverter(format_options={InputFormat.PDF: vlm_format_options})

    return regular_converter, vlm_converter


def process_page_worker(args: Tuple[str, int, int, str, str]) -> Optional[Dict[str, Any]]:
    """
    Process a single page - runs in separate process.
    Each process initializes its own converters for separate GPU contexts.
    """
    page_file_str, page_num, total_pages, output_dir_str, original_pdf_name = args
    page_file = Path(page_file_str)
    output_dir = Path(output_dir_str)

    try:
        # Initialize converters in this process (separate GPU context)
        regular_converter, vlm_converter = init_converters()

        # Phase 1: Regular processing for table detection
        regular_result = regular_converter.convert(str(page_file))

        if regular_result is None or not regular_result.document.pages:
            logger.error(f"Page {page_num}/{total_pages}: Phase 1 failed")
            return None

        # Analyze for tables
        document = regular_result.document
        page_tables = [
            table
            for table in document.tables
            if table.prov and len(table.prov) > 0 and table.prov[0].page_no == 1
        ]

        has_tables = len(page_tables) > 0
        table_count = len(page_tables)

        # Log once after Phase 1
        if has_tables:
            logger.info(f"Page {page_num}/{total_pages}: Processing")
        else:
            logger.info(f"Page {page_num}/{total_pages}: Processing (no tables)")

        # Phase 2: Conditional VLM processing
        if has_tables:
            vlm_result = vlm_converter.convert(str(page_file))
            if vlm_result is not None:
                final_result = vlm_result
            else:
                logger.debug(f"VLM failed for page {page_num}, using regular result")
                final_result = regular_result
        else:
            final_result = regular_result

        # Generate markdown
        markdown_content = final_result.document.export_to_markdown()

        # Save output
        base_name = Path(original_pdf_name).stem
        output_filename = f"{base_name}_page_{page_num:03d}.md"
        output_path = output_dir / output_filename

        with open(output_path, "w", encoding="utf-8") as f:
            f.write(markdown_content)

        return {
            "page_num": page_num,
            "output_path": str(output_path),
            "has_tables": has_tables,
            "table_count": table_count,
        }

    except Exception as e:
        logger.error(f"Page {page_num}/{total_pages}: Failed - {str(e)}")
        return None


def extract_pages_to_temp(pdf_path: Path, temp_dir: Path) -> List[Path]:
    """Extract individual pages from PDF to temporary files"""
    doc = fitz.open(str(pdf_path))
    total_pages = len(doc)
    page_files = []

    logger.debug(f"Extracting {total_pages} pages from {pdf_path.name}")

    for page_num in range(total_pages):
        new_doc = fitz.open()
        new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
        page_filename = f"page_{page_num + 1:03d}.pdf"
        page_path = temp_dir / page_filename
        new_doc.save(str(page_path))
        new_doc.close()
        page_files.append(page_path)

    doc.close()
    return page_files


def convert_document(pdf_path: Path, output_dir: Path, max_workers: int = 2) -> dict:
    """Convert PDF using ProcessPoolExecutor for parallel processing"""
    logger.info(f"Starting conversion of {pdf_path.name}")

    if not pdf_path.exists():
        logger.error(f"File not found: {pdf_path}")
        return {"success": 0, "failed": 0, "files": []}

    # Ensure output directory exists
    output_dir.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as temp_dir:
        temp_path = Path(temp_dir)

        try:
            # Extract pages
            page_files = extract_pages_to_temp(pdf_path, temp_path)
            total_pages = len(page_files)

            logger.info(f"Processing {total_pages} pages with {max_workers} process workers")

            # Prepare arguments for each page
            page_args = [
                (str(page_file), i + 1, total_pages, str(output_dir), pdf_path.name)
                for i, page_file in enumerate(page_files)
            ]

            # Process pages using ProcessPoolExecutor
            successful_count = 0
            failed_count = 0

            with ProcessPoolExecutor(max_workers=max_workers) as executor:
                # Submit all tasks
                future_to_args = {executor.submit(process_page_worker, args): args for args in page_args}

                # Process as completed
                for future in as_completed(future_to_args):
                    result = future.result()
                    if result:
                        successful_count += 1
                    else:
                        failed_count += 1

            logger.info(f"Conversion complete: {successful_count} successful, {failed_count} failed")

            return {"success": successful_count, "failed": failed_count, "total_pages": total_pages}

        except Exception as e:
            logger.error(f"Critical error during conversion: {str(e)}")
            return {"success": 0, "failed": 0, "total_pages": 0}


def main():
    parser = argparse.ArgumentParser(
        description="ProcessPoolExecutor PDF converter for GPU utilization testing"
    )
    parser.add_argument("pdf_path", help="Path to the PDF file to convert")
    parser.add_argument(
        "--output-dir", "-o", default="output", help="Output directory for markdown files (default: output)"
    )
    parser.add_argument("--workers", "-w", type=int, default=2, help="Number of process workers (default: 2)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose logging")

    args = parser.parse_args()

    if args.verbose:
        logger.remove()
        logger.add(
            sys.stderr,
            format="<green>{time:HH:mm:ss}</green> | <level>{level}</level> | {message}",
            level="DEBUG",
        )

    pdf_path = Path(args.pdf_path).resolve()
    output_dir = Path(args.output_dir).resolve()

    logger.info("Initialized ProcessPool PDF converter")
    logger.info(f"Workers: {args.workers}")

    results = convert_document(pdf_path, output_dir, max_workers=args.workers)

    logger.success(f"Successfully converted {results['success']} pages")
    if results["failed"] > 0:
        logger.warning(f"{results['failed']} pages failed to convert")


if __name__ == "__main__":
    main()