# pyright: ignore
"""
Enhanced Table Parser with Header Consolidation

Parses tables from Docling documents and consolidates consecutive headers
that immediately precede tables, following strict rules to avoid incorrect
consolidation across content boundaries.
"""

from typing import List, Optional, Tuple, Dict, Any
from dataclasses import dataclass

from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption


@dataclass
class TableMetadata:
    """Container for table metadata including consolidated headers"""
    table_index: int
    page_number: int
    bounding_box: Optional[Tuple[float, float, float, float]]
    consolidated_header: Optional[str]
    headers: List[str]  # Individual headers found
    descriptive_text: Optional[str]  # Text that describes the table
    preceding_text: Optional[str]  # Context before table
    following_text: Optional[str]  # Context after table
    rows: int
    cols: int
    first_cell_content: Optional[str]
    table_markdown: Optional[str]  # Full table in markdown


class EnhancedTableParser:
    """
    Enhanced table parser with intelligent header consolidation.

    Consolidates consecutive headers of the same level that immediately
    precede tables, with strict rules to prevent incorrect consolidation.
    """

    def __init__(
        self,
        separator: str = " ^ ",
        max_lookback: int = 10,
        enable_consolidation: bool = True,
        debug: bool = False
    ):
        """
        Initialize the enhanced table parser.

        Args:
            separator: String to join consolidated headers
            max_lookback: Maximum elements to examine backward from table
            enable_consolidation: Enable/disable header consolidation
            debug: Enable debug output
        """
        self.separator = separator
        self.max_lookback = max_lookback
        self.enable_consolidation = enable_consolidation
        self.debug = debug

    def parse_document(self, document_path: str) -> Tuple[Any, List[TableMetadata]]:
        """
        Parse document and extract tables with consolidated headers.

        Args:
            document_path: Path to PDF document

        Returns:
            Tuple of (document object, list of table metadata)
        """
        # Configure pipeline for table extraction
        pipeline_options = PdfPipelineOptions()
        pipeline_options.do_table_structure = True

        converter = DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
            }
        )

        # Convert document
        result = converter.convert(document_path)
        document = result.document

        if self.debug:
            print(f"Document loaded: {len(document.pages)} pages")
            print(f"Total elements: {len(document.body.children)}")
            print(f"Tables found: {len(document.tables)}")

        # Process tables with header consolidation
        table_metadata = []

        # Build element index for efficient lookup
        element_index = self._build_element_index(document)

        for i, table in enumerate(document.tables):
            metadata = self._process_table(document, table, i, element_index)
            table_metadata.append(metadata)

        return document, table_metadata

    def _build_element_index(self, document) -> Dict[Any, int]:
        """
        Build index mapping elements to their position in document.body.children.

        Args:
            document: Docling document object

        Returns:
            Dictionary mapping resolved elements to their index
        """
        element_index = {}
        for idx, child_ref in enumerate(document.body.children):
            resolved = child_ref.resolve(document)
            if resolved:
                # Store by id or object reference
                element_index[id(resolved)] = idx
        return element_index

    def _process_table(
        self,
        document,
        table,
        table_num: int,
        element_index: Dict[Any, int]
    ) -> TableMetadata:
        """
        Process a single table and extract metadata with consolidated headers.

        Args:
            document: Docling document object
            table: Table element
            table_num: Table number (0-based)
            element_index: Element position index

        Returns:
            TableMetadata object
        """
        # Extract basic table information
        page_number = None
        bbox = None
        if hasattr(table, "prov") and table.prov:
            prov = table.prov[0]
            page_number = prov.page_no  # Already 1-based in Docling
            if hasattr(prov, "bbox"):
                bbox = prov.bbox

        # Get table dimensions
        rows = cols = 0
        first_cell = None
        if hasattr(table, "data") and table.data:
            rows = table.data.num_rows
            cols = table.data.num_cols
            try:
                df = table.export_to_dataframe()
                if not df.empty:
                    first_cell = str(df.iloc[0, 0])
            except:  # noqa: E722
                pass

        # Find table's position in document
        table_doc_index = element_index.get(id(table))

        # Extract context and headers
        headers = []
        consolidated_header = None
        descriptive_text = None
        preceding_text = None
        following_text = None

        if table_doc_index is not None:
            if self.enable_consolidation:
                # Get headers and consolidation
                header_result = self._extract_headers_and_context(
                    document, table_doc_index
                )
                headers = header_result.get('headers', [])
                consolidated_header = header_result.get('consolidated', None)
                descriptive_text = header_result.get('descriptive_text', None)

            # Get surrounding context
            preceding_text = self._get_preceding_text(document, table_doc_index, max_chars=500)
            following_text = self._get_following_text(document, table_doc_index, max_chars=200)

        # Get table markdown
        table_markdown = None
        try:
            table_markdown = table.export_to_markdown(doc=document)
        except:  # noqa: E722
            pass

        metadata = TableMetadata(
            table_index=table_num,
            page_number=page_number,
            bounding_box=bbox,
            consolidated_header=consolidated_header,
            headers=headers,
            descriptive_text=descriptive_text,
            preceding_text=preceding_text,
            following_text=following_text,
            rows=rows,
            cols=cols,
            first_cell_content=first_cell,
            table_markdown=table_markdown
        )

        if self.debug:
            print(f"
Table {table_num + 1}:")
            print(f"  Page: {page_number}")
            print(f"  Dimensions: {rows}x{cols}")
            print(f"  Header: {consolidated_header[:100] if consolidated_header else 'None'}")

        return metadata

    def _extract_headers_and_context(
        self,
        document,
        table_index: int
    ) -> Dict[str, Any]:
        """
        Extract headers and descriptive text before a table.

        Returns dict with:
        - headers: List of header texts
        - consolidated: Consolidated header string
        - descriptive_text: Text that describes the table
        """
        headers = []
        descriptive_text = None
        current_index = table_index - 1
        header_level = None

        # First, check for descriptive text immediately before the table
        if current_index >= 0:
            element_ref = document.body.children[current_index]
            element = element_ref.resolve(document)

            if element and self._is_text_element(element) and not self._is_header_element(element):
                # This is non-header text right before the table
                if hasattr(element, 'text') and element.text:
                    text = element.text.strip()
                    if self._is_table_description(text) or len(text) > 10:  # Capture any meaningful text
                        descriptive_text = text
                        if self.debug:
                            print(f"    Found descriptive text: {text[:50]}")
                        current_index -= 1  # Move past this text to look for headers

        # Now look for headers, continuing past any descriptive text
        # We'll look for the nearest header, skipping over tables and other non-text elements
        found_any_header = False

        while current_index >= 0 and (table_index - current_index) <= self.max_lookback:
            if current_index >= len(document.body.children):
                break

            element_ref = document.body.children[current_index]
            element = element_ref.resolve(document)

            if not element:
                current_index -= 1
                continue

            # Check if this is a text element
            if not self._is_text_element(element):
                # Skip non-text elements (like tables, pictures) but keep looking
                if self.debug:
                    print(f"    Skipping non-text element: {type(element).__name__}")
                current_index -= 1
                continue

            # Check if this is a header
            if self._is_header_element(element):
                element_level = self._get_header_level(element)

                # If we haven't found any headers yet, take this one
                if not found_any_header:
                    found_any_header = True
                    header_level = element_level
                    header_text = element.text.strip() if hasattr(element, 'text') else ""
                    if header_text:
                        headers.insert(0, header_text)
                        if self.debug:
                            print(f"    Found first header: {header_text[:50]}")
                # Check if this is the same level as previous headers
                elif header_level == element_level:
                    header_text = element.text.strip() if hasattr(element, 'text') else ""
                    if header_text:
                        headers.insert(0, header_text)
                        if self.debug:
                            print(f"    Found additional header: {header_text[:50]}")
                else:
                    # Different level header, stop here
                    if self.debug:
                        print(f"    Stopping at different header level: {element_level} != {header_level}")
                    break
            else:
                # Non-header text
                if found_any_header:
                    # We've already found headers, this text marks the end of header section
                    break
                else:
                    # No headers found yet, this might be descriptive text
                    if not descriptive_text and hasattr(element, 'text') and element.text:
                        text = element.text.strip()
                        if len(text) > 10:  # Any meaningful text
                            descriptive_text = text
                            if self.debug:
                                print(f"    Found fallback descriptive text: {text[:50]}")

            current_index -= 1

        # Build result
        result = {'headers': headers}

        # Consolidate headers if multiple found
        if len(headers) > 1:
            result['consolidated'] = self.separator.join(headers)
        elif len(headers) == 1:
            result['consolidated'] = headers[0]
        else:
            result['consolidated'] = None

        result['descriptive_text'] = descriptive_text

        return result

    def _is_table_description(self, text: str) -> bool:
        """Check if text likely describes a table"""
        if not text:
            return False

        indicators = [
            "table", "following", "below", "presents", "shows",
            "summary", "details", "breakdown", "comprises"
        ]
        text_lower = text.lower()
        return any(ind in text_lower for ind in indicators)

    def _get_preceding_text(
        self,
        document,
        table_index: int,
        max_chars: int = 500
    ) -> Optional[str]:
        """Get text content before table up to max_chars"""
        texts = []
        total_chars = 0
        current_index = table_index - 1

        while current_index >= 0 and total_chars < max_chars:
            if current_index >= len(document.body.children):
                break

            element_ref = document.body.children[current_index]
            element = element_ref.resolve(document)

            if element and self._is_text_element(element):
                if hasattr(element, 'text') and element.text:
                    text = element.text.strip()
                    if text:
                        texts.insert(0, text)
                        total_chars += len(text)

            current_index -= 1

        return ' '.join(texts)[:max_chars] if texts else None

    def _get_following_text(
        self,
        document,
        table_index: int,
        max_chars: int = 200
    ) -> Optional[str]:
        """Get text content after table up to max_chars"""
        texts = []
        total_chars = 0
        current_index = table_index + 1

        while current_index < len(document.body.children) and total_chars < max_chars:
            element_ref = document.body.children[current_index]
            element = element_ref.resolve(document)

            if element and self._is_text_element(element):
                if hasattr(element, 'text') and element.text:
                    text = element.text.strip()
                    if text:
                        texts.append(text)
                        total_chars += len(text)

            current_index += 1

        return ' '.join(texts)[:max_chars] if texts else None

    def _is_text_element(self, element) -> bool:
        """
        Check if element is a text element (not Picture, Table, etc.).

        Args:
            element: Resolved document element

        Returns:
            True if text element, False otherwise
        """
        # Check class name
        class_name = type(element).__name__.lower()

        # Exclude non-text elements
        if any(x in class_name for x in ['picture', 'image', 'table', 'figure']):
            return False

        # Must have text attribute
        return hasattr(element, 'text')

    def _is_header_element(self, element) -> bool:
        """
        Determine if a text element is a header.

        Args:
            element: Text element to check

        Returns:
            True if element is a header, False otherwise
        """
        if not hasattr(element, 'text') or not element.text:
            return False

        text = element.text.strip()

        # Exclude common non-header patterns
        # Page numbers, timestamps, single numbers, "Document" labels
        if (text.isdigit() or
            'PM' in text or 'AM' in text or
            '/' in text and ':' in text or  # Date/time patterns
            text.lower() == 'document' or
            len(text) < 3):  # Very short text
            return False

        # Check for markdown headers
        if text.startswith('#'):
            return True

        # Check for label/type indicating header
        if hasattr(element, 'label'):
            label = str(element.label).lower()
            if any(h in label for h in ['header', 'heading', 'title', 'section']):
                return True

        # Heuristic: Short text that looks like a header
        # - Not too long (headers are typically short)
        # - Contains meaningful words (not just numbers/symbols)
        # - Title case or all caps
        # - No ending punctuation (except colon)
        if len(text) < 200 and '
' not in text:
            # Must contain at least one letter
            if not any(c.isalpha() for c in text):
                return False

            # Check if it looks like a title
            is_title_case = text.istitle() or text.isupper()
            no_ending_punctuation = not text.rstrip().endswith(('.', '!', '?'))

            if is_title_case and (no_ending_punctuation or text.endswith(':')):
                return True

        return False

    def _get_header_level(self, element) -> int:
        """
        Extract header level from element (H1=1, H2=2, etc.).

        Args:
            element: Header element

        Returns:
            Header level (1-6) or 0 if undetermined
        """
        if hasattr(element, 'text') and element.text:
            text = element.text.strip()

            # Count markdown header marks
            if text.startswith('#'):
                level = 0
                for char in text:
                    if char == '#':
                        level += 1
                    else:
                        break
                return min(level, 6)

        # Check label for level indication
        if hasattr(element, 'label'):
            label = str(element.label).lower()
            for level in range(1, 7):
                if f'h{level}' in label or f'heading{level}' in label:
                    return level

        # Default to level 2 for identified headers without specific level
        return 2

    def export_results(
        self,
        document,
        table_metadata: List[TableMetadata],
        output_path: str
    ):
        """
        Export results with consolidated headers.

        Args:
            document: Docling document object
            table_metadata: List of table metadata
            output_path: Path for output file
        """
        output = []
        output.append("# Table Extraction Report with Consolidated Headers
")

        for meta in table_metadata:
            output.append(f"
## Table {meta.table_index + 1}")
            output.append(f"**Page:** {meta.page_number}")
            output.append(f"**Dimensions:** {meta.rows} rows × {meta.cols} columns")

            if meta.consolidated_header:
                output.append(f"**Consolidated Header:** {meta.consolidated_header}")
            elif meta.headers:
                output.append(f"**Headers Found:** {', '.join(meta.headers)}")
            else:
                output.append("**Header:** No header found")

            if meta.descriptive_text:
                output.append(f"**Descriptive Text:** {meta.descriptive_text[:200]}")

            if meta.first_cell_content:
                output.append(f"**First Cell:** {meta.first_cell_content[:100]}")

            output.append("")

            # Export table content
            try:
                table = document.tables[meta.table_index]
                table_md = table.export_to_markdown(doc=document)
                output.append("**Table Content:**")
                output.append(table_md)
            except Exception as e:
                output.append(f"*Error exporting table: {e}*")

            output.append("
---")

        # Write to file
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('
'.join(output))

        print(f"Results exported to: {output_path}")


def main():
    """
    Example usage of the enhanced table parser.
    """
    import argparse

    parser = argparse.ArgumentParser(
        description="Enhanced table parser with header consolidation"
    )
    parser.add_argument(
        "document_path",
        help="Path to PDF document"
    )
    parser.add_argument(
        "--output",
        default="table_extraction_report.md",
        help="Output file path (default: table_extraction_report.md)"
    )
    parser.add_argument(
        "--separator",
        default=" ^ ",
        help="Separator for consolidated headers (default: ' ^ ')"
    )
    parser.add_argument(
        "--no-consolidation",
        action="store_true",
        help="Disable header consolidation"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug output"
    )

    args = parser.parse_args()

    # Create parser instance
    table_parser = EnhancedTableParser(
        separator=args.separator,
        enable_consolidation=not args.no_consolidation,
        debug=args.debug
    )

    # Parse document
    print(f"Processing: {args.document_path}")
    document, table_metadata = table_parser.parse_document(args.document_path)

    print(f"Found {len(table_metadata)} tables")

    # Export results
    table_parser.export_results(document, table_metadata, args.output)

    # Print summary
    tables_with_headers = sum(1 for m in table_metadata if m.consolidated_header)
    print(f"Tables with consolidated headers: {tables_with_headers}/{len(table_metadata)}")


if __name__ == "__main__":
    main()