"""PDF parser using the Docling library for advanced document processing."""

import asyncio
import logging
import time
from pathlib import Path
from typing import Any, Dict, Optional

from .parser import ParseResult, PDFParser

logger = logging.getLogger(__name__)

# Default configuration following the design specification
DOCLING_DEFAULT_CONFIG = {
    # OCR Configuration
    "ocr_enabled": True,
    "ocr_engine": "easyocr",  # Options: easyocr, tesseract, rapidocr, ocrmac, onnxtr
    "ocr_languages": ["en"],
    "ocr_device": "auto",  # auto, cpu, cuda, mps
    # Table Processing
    "table_processing_mode": "FAST",  # FAST, ACCURATE
    "table_extraction_enabled": True,
    # Enrichment Features
    "formula_enrichment": True,
    "picture_description": False,  # Requires additional models
    "code_understanding": False,  # Requires additional models
    # Performance Settings
    "processing_timeout": 300,  # seconds
    "max_pages": None,  # None for no limit
    "device_selection": "auto",
    # Output Options
    "export_format": "markdown",
    "include_images": False,
    "preserve_layout": True,
    # Security Settings
    "max_file_size": 100 * 1024 * 1024,  # 100MB default
    "allow_external_resources": False,
}

# Supported OCR engines with their import requirements
OCR_ENGINE_DEPENDENCIES = {
    "easyocr": ["easyocr"],
    "tesseract": ["pytesseract"],
    "rapidocr": ["rapidocr_onnxruntime"],
    "ocrmac": [],  # System-dependent (macOS only)
    "onnxtr": ["doctr"],
}


class DoclingParser(PDFParser):
    """PDF parser using the Docling library for advanced document processing."""

    def __init__(self, config: Optional[Dict[str, Any]] = None, cache_dir: Path = None):
        """Initialize the Docling parser with comprehensive configuration support.

        Args:
            config: Configuration options for docling processing.
            cache_dir: Directory to cache parsed markdown files.
        """
        super().__init__(cache_dir)

        # Merge user config with defaults and validate
        self.config = self._validate_config({**DOCLING_DEFAULT_CONFIG, **(config or {})})

        # Determine available features
        self.available_features = self._configure_available_features()

        logger.debug(f"DoclingParser initialized with config: {self.config}")
        logger.debug(f"Available features: {self.available_features}")

    async def parse(self, file_path: Path) -> ParseResult:
        """Parse a PDF file using Docling library with advanced processing.

        Args:
            file_path: Path to the PDF file.

        Returns:
            ParseResult with markdown content and metadata.
        """
        start_time = time.time()

        try:
            # Check cache first
            cache_path = None
            if self.cache_dir:
                cache_path = self._get_cache_path(file_path)
                if self._is_cache_valid(file_path, cache_path):
                    logger.debug(f"Loading parsed content from cache: {cache_path}")
                    markdown_content = self._load_from_cache(cache_path)
                    metadata = self._load_metadata_from_cache(cache_path)
                    if markdown_content is not None and metadata:
                        return ParseResult(markdown_content=markdown_content, metadata=metadata)

            # Validate input file for security
            self._validate_input_file(file_path)

            logger.info(f"Starting Docling parsing: {file_path}")
            logger.debug(f"Docling configuration: {self.config}")

            # Lazy imports with comprehensive error handling
            try:
                from docling.datamodel.base_models import InputFormat
                from docling.datamodel.pipeline_options import PdfPipelineOptions
                from docling.document_converter import DocumentConverter, PdfFormatOption
            except ImportError as e:
                raise ImportError(
                    "Docling library not available. Install with: pip install docling[complete] "
                    "or pip install docling for basic functionality"
                ) from e

            # Create pipeline options and converter with new v2 API
            pipeline_options = self._build_pipeline_options(PdfPipelineOptions)
            converter = self._create_safe_converter(DocumentConverter, PdfFormatOption, InputFormat, pipeline_options)

            # Apply timeout wrapper for processing
            conversion_result = await self._convert_with_timeout(converter, file_path)

            # Handle any conversion errors or warnings
            self._handle_conversion_errors(conversion_result)

            # Extract markdown content
            markdown_content = conversion_result.document.export_to_markdown()

            if not markdown_content or not markdown_content.strip():
                raise RuntimeError("Docling produced empty markdown content")

            # Extract comprehensive metadata
            metadata = self._extract_metadata_from_result(conversion_result, file_path)

            # Add processing information
            processing_time = time.time() - start_time
            metadata.update(
                {
                    "processing_timestamp": "N/A",  # Will be set by PDFProcessor
                    "processor_version": "docling",
                    "source_filename": file_path.name,
                    "source_directory": str(file_path.parent),
                    "docling_processing_time": processing_time,
                    "docling_features_used": self._get_features_used(),
                }
            )

            # Save to cache if enabled
            if cache_path:
                logger.debug(f"Saving parsed content to cache: {cache_path}")
                self._save_to_cache(cache_path, markdown_content)
                self._save_metadata_to_cache(cache_path, metadata)

            # Log processing statistics
            logger.info(f"Docling parsing completed in {processing_time:.2f}s: {file_path}")
            logger.debug(f"Pages processed: {metadata.get('page_count', 'unknown')}")
            logger.debug(f"OCR enabled: {self.available_features['ocr']}")
            logger.debug(f"Tables extracted: {metadata.get('table_count', 0)}")

            return ParseResult(markdown_content=markdown_content, metadata=metadata)

        except ImportError:
            # Re-raise ImportError as-is for parser fallback logic
            raise
        except Exception as e:
            processing_time = time.time() - start_time
            logger.error(f"Docling parsing failed after {processing_time:.2f}s: {file_path} - {e}")
            raise RuntimeError(f"Failed to parse PDF with Docling: {e}") from e

    def _validate_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Validate and normalize docling configuration.

        Args:
            config: Configuration dictionary to validate.

        Returns:
            Validated and normalized configuration.

        Raises:
            ValueError: If configuration values are invalid.
        """
        validated = config.copy()

        # Validate OCR engine
        if validated["ocr_engine"] not in OCR_ENGINE_DEPENDENCIES:
            supported_engines = list(OCR_ENGINE_DEPENDENCIES.keys())
            raise ValueError(
                f"Unsupported OCR engine: {validated['ocr_engine']}. " f"Supported engines: {supported_engines}"
            )

        # Validate table processing mode
        if validated["table_processing_mode"] not in ["FAST", "ACCURATE"]:
            raise ValueError("table_processing_mode must be FAST or ACCURATE")

        # Validate device selection
        if validated["device_selection"] not in ["auto", "cpu", "cuda", "mps"]:
            raise ValueError("device_selection must be auto, cpu, cuda, or mps")

        # Validate timeout
        if validated["processing_timeout"] <= 0:
            raise ValueError("processing_timeout must be positive")

        # Validate file size limit
        if validated["max_file_size"] <= 0:
            raise ValueError("max_file_size must be positive")

        # Validate OCR languages (must be list)
        if not isinstance(validated["ocr_languages"], list):
            if isinstance(validated["ocr_languages"], str):
                validated["ocr_languages"] = [validated["ocr_languages"]]
            else:
                raise ValueError("ocr_languages must be a list of language codes")

        return validated

    def _configure_available_features(self) -> Dict[str, bool]:
        """Determine which docling features are available based on dependencies.

        Returns:
            Dictionary mapping feature names to availability status.
        """
        features = {
            "ocr": False,
            "table_extraction": True,  # Always available with docling
            "formula_enrichment": True,  # Always available with docling
            "picture_description": False,
            "code_understanding": False,
        }

        # Check OCR availability
        if self.config["ocr_enabled"]:
            selected_engine = self._get_available_ocr_engine(self.config["ocr_engine"])
            if selected_engine:
                features["ocr"] = True
                # Update config with available engine if different from requested
                if selected_engine != self.config["ocr_engine"]:
                    self.config["ocr_engine"] = selected_engine

        # Check enrichment model availability (simplified check)
        try:
            # These features require additional model downloads
            # For now, we'll conservatively set them to False unless explicitly tested
            features["picture_description"] = self.config.get("picture_description", False)
            features["code_understanding"] = self.config.get("code_understanding", False)
        except Exception:
            logger.debug("Advanced enrichment features not available")

        return features

    def _get_available_ocr_engine(self, preferred_engine: str) -> Optional[str]:
        """Select best available OCR engine with fallback logic.

        Args:
            preferred_engine: The preferred OCR engine to use.

        Returns:
            Available OCR engine name or None if no engines available.
        """
        # Check if preferred engine is available
        if self._check_ocr_engine_available(preferred_engine):
            return preferred_engine

        # Fallback order: easyocr -> tesseract -> rapidocr -> ocrmac -> onnxtr
        fallback_order = ["easyocr", "tesseract", "rapidocr", "ocrmac", "onnxtr"]

        for engine in fallback_order:
            if engine != preferred_engine and self._check_ocr_engine_available(engine):
                logger.warning(f"Preferred OCR engine '{preferred_engine}' not available, using '{engine}'")
                return engine

        # No OCR engines available
        logger.warning("No OCR engines available, proceeding without OCR")
        return None

    def _check_ocr_engine_available(self, engine: str) -> bool:
        """Check if a specific OCR engine is available.

        Args:
            engine: OCR engine name to check.

        Returns:
            True if engine is available, False otherwise.
        """
        if engine == "ocrmac":
            import platform

            return platform.system() == "Darwin"

        dependencies = OCR_ENGINE_DEPENDENCIES.get(engine, [])
        for dep in dependencies:
            try:
                __import__(dep)
            except ImportError:
                logger.debug(f"OCR engine '{engine}' dependency '{dep}' not available")
                return False

        return True

    def _build_pipeline_options(self, PdfPipelineOptions):
        """Build PdfPipelineOptions from configuration using Docling v2 API.

        Args:
            PdfPipelineOptions: The PdfPipelineOptions class from docling.

        Returns:
            Configured PdfPipelineOptions instance.
        """
        try:
            pipeline_options = PdfPipelineOptions()

            # Configure OCR using new v2 API
            if self.available_features["ocr"] and self.config["ocr_enabled"]:
                pipeline_options.do_ocr = True

                # Set OCR engine-specific options
                ocr_engine = self.config["ocr_engine"]
                if ocr_engine == "rapidocr":
                    from docling.datamodel.pipeline_options import RapidOcrOptions

                    pipeline_options.ocr_options = RapidOcrOptions()
                elif ocr_engine == "easyocr":
                    from docling.datamodel.pipeline_options import EasyOcrOptions

                    pipeline_options.ocr_options = EasyOcrOptions()
                    if hasattr(pipeline_options.ocr_options, "lang"):
                        pipeline_options.ocr_options.lang = self.config["ocr_languages"]
                elif ocr_engine == "tesseract":
                    from docling.datamodel.pipeline_options import TesseractOcrOptions

                    pipeline_options.ocr_options = TesseractOcrOptions()
                elif ocr_engine == "ocrmac":
                    from docling.datamodel.pipeline_options import OcrMacOptions

                    pipeline_options.ocr_options = OcrMacOptions()
                else:
                    # Fallback to base OcrOptions
                    from docling.datamodel.pipeline_options import OcrOptions

                    pipeline_options.ocr_options = OcrOptions()
            else:
                pipeline_options.do_ocr = False

            # Configure table processing using new v2 API
            pipeline_options.do_table_structure = self.config["table_extraction_enabled"]
            if pipeline_options.do_table_structure:
                # Set table mode based on config
                if self.config["table_processing_mode"].upper() == "ACCURATE":
                    from docling.datamodel.pipeline_options import TableFormerMode

                    pipeline_options.table_structure_options.mode = TableFormerMode.ACCURATE

            # Configure enrichment features using new v2 API
            pipeline_options.do_formula_enrichment = (
                self.config["formula_enrichment"] and self.available_features["formula_enrichment"]
            )
            pipeline_options.do_picture_description = (
                self.config["picture_description"] and self.available_features["picture_description"]
            )
            pipeline_options.do_code_enrichment = (
                self.config["code_understanding"] and self.available_features["code_understanding"]
            )

            # Apply resource limits and timeouts
            pipeline_options.document_timeout = self.config["processing_timeout"]

            # Apply additional resource limits
            pipeline_options = self._apply_resource_limits(pipeline_options)

            return pipeline_options

        except Exception as e:
            logger.warning(f"Failed to configure pipeline options: {e}, using defaults")
            return PdfPipelineOptions()

    def _apply_resource_limits(self, pipeline_options):
        """Apply resource limits to prevent excessive resource usage.

        Args:
            pipeline_options: PdfPipelineOptions instance to modify.

        Returns:
            Modified pipeline_options with resource limits applied.
        """
        # Limit maximum pages processed
        if self.config.get("max_pages"):
            if hasattr(pipeline_options, "page_range"):
                pipeline_options.page_range = (1, self.config["max_pages"])

        # Configure device usage
        device = self.config.get("device_selection", "auto")
        if device != "auto" and hasattr(pipeline_options, "device"):
            pipeline_options.device = device

        # Memory limits for table processing
        if (
            self.config["table_processing_mode"] == "FAST"
            and hasattr(pipeline_options, "table_options")
            and hasattr(pipeline_options.table_options, "max_table_size")
        ):
            pipeline_options.table_options.max_table_size = 1000000  # 1MB limit

        return pipeline_options

    def _create_safe_converter(self, DocumentConverter, PdfFormatOption, InputFormat, pipeline_options):
        """Create DocumentConverter with security restrictions using Docling v2 API.

        Args:
            DocumentConverter: The DocumentConverter class from docling.
            PdfFormatOption: The PdfFormatOption class from docling.
            InputFormat: The InputFormat enum from docling.
            pipeline_options: Configured PdfPipelineOptions instance.

        Returns:
            Configured DocumentConverter instance.
        """
        try:
            # Create converter with format options (new v2 API)
            converter = DocumentConverter(
                format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}
            )

            return converter

        except Exception as e:
            logger.warning(f"Failed to configure converter: {e}, using defaults")
            # Fallback to basic converter
            return DocumentConverter()

    def _optimize_memory_usage(self, converter):
        """Configure converter for optimal memory usage.

        Args:
            converter: DocumentConverter instance to optimize.

        Returns:
            Optimized converter instance.
        """
        try:
            # Configure batch processing for large documents
            if self.config.get("max_pages") and hasattr(converter, "config"):
                converter.config.max_pages = self.config["max_pages"]

            # Configure memory-efficient table processing
            if self.config["table_processing_mode"] == "FAST" and hasattr(converter, "config"):
                # Use faster, less memory-intensive processing
                pass

        except Exception as e:
            logger.debug(f"Memory optimization failed: {e}")

        return converter

    async def _convert_with_timeout(self, converter, file_path: Path):
        """Convert document with timeout handling using Docling v2 API.

        Args:
            converter: DocumentConverter instance.
            file_path: Path to the PDF file.

        Returns:
            ConversionResult from docling.

        Raises:
            RuntimeError: If conversion times out or fails.
        """
        timeout = self.config["processing_timeout"]

        try:
            # Create a coroutine wrapper for the synchronous conversion
            async def convert_async():
                # Use new v2 API - no pipeline_options parameter
                return converter.convert(str(file_path))

            # Apply timeout
            return await asyncio.wait_for(convert_async(), timeout=timeout)

        except asyncio.TimeoutError:
            raise RuntimeError(f"Docling processing timed out after {timeout} seconds")
        except Exception as e:
            raise RuntimeError(f"Docling conversion failed: {e}") from e

    def _handle_conversion_errors(self, conversion_result):
        """Handle partial conversion failures and log detailed error information.

        Args:
            conversion_result: ConversionResult from docling.

        Raises:
            RuntimeError: If conversion failed completely.
        """
        try:
            # Import ConversionStatus enum
            from docling.datamodel.base_models import ConversionStatus

            if conversion_result.status == ConversionStatus.FAILURE:
                error_details = []

                # Collect error information from result
                if hasattr(conversion_result, "errors") and conversion_result.errors:
                    error_details.extend([str(error) for error in conversion_result.errors])

                error_msg = (
                    f"Docling conversion failed: {'; '.join(error_details) if error_details else 'Unknown error'}"
                )
                raise RuntimeError(error_msg)

            elif conversion_result.status == ConversionStatus.PARTIAL_SUCCESS:
                warning_details = []
                if hasattr(conversion_result, "warnings") and conversion_result.warnings:
                    warning_details.extend([str(warning) for warning in conversion_result.warnings])

                logger.warning(f"Docling conversion partially successful: {'; '.join(warning_details)}")

        except ImportError:
            # If we can't import ConversionStatus, do basic checks
            if not conversion_result or not hasattr(conversion_result, "document"):
                raise RuntimeError("Docling conversion failed: Invalid result")

            if not conversion_result.document:
                raise RuntimeError("Docling conversion failed: No document produced")

    def _extract_metadata_from_result(self, conversion_result, file_path: Path) -> Dict[str, Any]:
        """Extract comprehensive metadata from Docling's ConversionResult.

        Args:
            conversion_result: ConversionResult from docling.
            file_path: Path to the original PDF file.

        Returns:
            Dictionary of extracted metadata.
        """
        metadata = {}

        try:
            document = conversion_result.document

            # Basic document information
            metadata["page_count"] = len(document.pages) if document.pages else 1

            # Count different element types
            element_counts = {}
            table_count = 0
            image_count = 0
            formula_count = 0

            for page in document.pages:
                for element in page.elements:
                    element_type = type(element).__name__
                    element_counts[element_type] = element_counts.get(element_type, 0) + 1

                    # Count specific elements of interest
                    if "table" in element_type.lower():
                        table_count += 1
                    elif "image" in element_type.lower() or "picture" in element_type.lower():
                        image_count += 1
                    elif "formula" in element_type.lower() or "equation" in element_type.lower():
                        formula_count += 1

            metadata.update(
                {
                    "element_types": element_counts,
                    "total_elements": sum(element_counts.values()),
                    "table_count": table_count,
                    "image_count": image_count,
                    "formula_count": formula_count,
                }
            )

            # Document-level metadata if available
            if hasattr(document, "metadata") and document.metadata:
                doc_metadata = document.metadata
                for key, value in doc_metadata.items():
                    if key not in metadata:  # Don't override our counts
                        metadata[f"doc_{key}"] = value

            # Processing statistics
            if hasattr(conversion_result, "processing_stats"):
                metadata["processing_stats"] = conversion_result.processing_stats

        except Exception as e:
            logger.warning(f"Failed to extract detailed metadata from Docling result: {e}")
            # Provide minimal fallback metadata
            metadata.update(
                {
                    "page_count": 1,
                    "element_types": {},
                    "total_elements": 0,
                    "table_count": 0,
                    "image_count": 0,
                    "formula_count": 0,
                }
            )

        return metadata

    def _get_features_used(self) -> Dict[str, bool]:
        """Get a summary of which features were actually used during processing.

        Returns:
            Dictionary mapping feature names to usage status.
        """
        return {
            "ocr_enabled": self.available_features["ocr"] and self.config["ocr_enabled"],
            "table_extraction": self.config["table_extraction_enabled"],
            "formula_enrichment": (self.config["formula_enrichment"] and self.available_features["formula_enrichment"]),
            "picture_description": (
                self.config["picture_description"] and self.available_features["picture_description"]
            ),
            "code_understanding": (self.config["code_understanding"] and self.available_features["code_understanding"]),
            "ocr_engine": self.config["ocr_engine"] if self.available_features["ocr"] else None,
            "table_mode": self.config["table_processing_mode"],
        }

    def _validate_input_file(self, file_path: Path) -> None:
        """Validate input file for security and basic requirements.

        Args:
            file_path: Path to the file to validate.

        Raises:
            RuntimeError: If file validation fails.
        """
        # Check file exists
        if not file_path.exists():
            raise RuntimeError(f"File not found: {file_path}")

        # Check file size limits
        file_size = file_path.stat().st_size
        max_size = self.config["max_file_size"]
        if file_size > max_size:
            raise RuntimeError(f"File too large: {file_size} bytes (max: {max_size})")

        if file_size == 0:
            raise RuntimeError("File is empty")

        # Validate file extension
        if file_path.suffix.lower() != ".pdf":
            raise RuntimeError(f"Invalid file type: {file_path.suffix}")

        # Basic PDF header validation
        try:
            with open(file_path, "rb") as f:
                header = f.read(8)
                if not header.startswith(b"%PDF-"):
                    # Allow non-PDF files for testing, but log a warning
                    logger.warning(f"File does not have PDF header: {file_path}")
        except Exception as e:
            raise RuntimeError(f"Cannot read file: {e}")
