"""
HTML Table Extractor for parsing Docling-generated HTML.

This module extracts table structures from HTML, handling merged cells (rowspan/colspan)
and extracting the first 10 rows for LLM header detection.
"""

from typing import List, Optional
from dataclasses import dataclass

import structlog
from bs4 import BeautifulSoup

logger = structlog.get_logger(__name__)


@dataclass
class HtmlTable:
    """Represents an extracted HTML table."""
    html: str  # Raw HTML of the table
    rows: List[List[str]]  # 2D array of cell text content
    sheet_name: Optional[str] = None  # Sheet name if available
    table_index: int = 0  # Index of this table in the document
    row_count: int = 0  # Total number of rows
    col_count: int = 0  # Maximum number of columns


class HtmlTableExtractor:
    """Extracts tables from HTML structure generated by Docling."""

    # Table filtering configuration
    MIN_TABLE_COLS = 2  # Minimum columns for valid table (filters single-column text)

    def __init__(self, min_cols: int | None = None):
        """
        Initialize the HTML table extractor.

        Args:
            min_cols: Override minimum column threshold (default: 2)
        """
        # Allow threshold override for testing/tuning
        if min_cols is not None:
            self.MIN_TABLE_COLS = min_cols

        logger.debug(
            f"HtmlTableExtractor initialized with column threshold: "
            f"cols≥{self.MIN_TABLE_COLS}"
        )

    def extract_tables(self, html: str) -> List[HtmlTable]:
        """
        Parse HTML and extract all <table> elements with filtering.

        Filters out single-column tables (likely text blocks, not data).

        Args:
            html: HTML string from Docling conversion

        Returns:
            List of HtmlTable objects with structured table data (filtered)
        """
        valid_tables, _ = self.extract_all_tables(html)
        return valid_tables

    def extract_all_tables(self, html: str) -> tuple[List[HtmlTable], List[HtmlTable]]:
        """
        Parse HTML and extract all <table> elements, classifying them.

        Separates tables into:
        - Valid data tables (>= MIN_TABLE_COLS columns)
        - Raw text tables (single-column, containing text content)

        Args:
            html: HTML string from Docling conversion

        Returns:
            Tuple of (valid_tables, raw_text_tables)
            - valid_tables: Tables with >= MIN_TABLE_COLS columns (for LLM header detection)
            - raw_text_tables: Single-column tables (raw text, skip LLM)
        """
        try:
            logger.info("Parsing HTML to extract and classify tables...")
            soup = BeautifulSoup(html, 'lxml')

            # Find all table elements
            table_elements = soup.find_all('table')
            logger.info(f"Found {len(table_elements)} table(s) in HTML")

            extracted_tables = []
            for idx, table_elem in enumerate(table_elements):
                try:
                    table = self._extract_table(table_elem, idx)
                    extracted_tables.append(table)

                    logger.debug(
                        f"Extracted table {idx + 1}: {table.row_count} rows x "
                        f"{table.col_count} cols"
                    )
                except Exception as e:
                    logger.warning(f"Failed to extract table {idx + 1}: {e}")
                    continue

            # Classify tables into valid data tables and raw text tables
            valid_tables = []
            raw_text_tables = []

            for table in extracted_tables:
                if self._is_valid_table(table):
                    valid_tables.append(table)
                else:
                    # Single-column tables are raw text content
                    raw_text_tables.append(table)
                    logger.debug(
                        f"Classified table {table.table_index} as raw text: "
                        f"{table.row_count} rows × {table.col_count} col(s)"
                    )

            logger.info(
                f"Table classification complete: {len(table_elements)} found, "
                f"{len(valid_tables)} valid data tables, "
                f"{len(raw_text_tables)} raw text tables"
            )

            return valid_tables, raw_text_tables

        except Exception as e:
            logger.error(f"Error parsing HTML: {e}")
            return [], []

    def _is_valid_table(self, table: HtmlTable) -> bool:
        """
        Determine if HTML table is a legitimate data table.

        Filters out:
        - Single-column tables (likely text blocks or lists, not data)

        Args:
            table: HtmlTable to validate

        Returns:
            True if table has at least 2 columns, False otherwise
        """
        # Check column threshold only
        if table.col_count < self.MIN_TABLE_COLS:
            logger.debug(
                f"Table {table.table_index} invalid: "
                f"single-column table ({table.col_count} < {self.MIN_TABLE_COLS}) "
                f"- likely text, not data"
            )
            return False

        return True

    def _extract_table(self, table_elem, table_index: int) -> HtmlTable:
        """
        Extract a single table element into structured data.

        Handles:
        - Standard <tr> and <td>/<th> structure
        - Merged cells (rowspan, colspan)
        - Mixed <th> and <td> elements

        Args:
            table_elem: BeautifulSoup table element
            table_index: Index of this table in the document

        Returns:
            HtmlTable object with extracted data
        """
        # Extract all rows
        rows = self._extract_rows_with_spans(table_elem)

        # Calculate dimensions
        row_count = len(rows)
        col_count = max(len(row) for row in rows) if rows else 0

        # Try to extract sheet name from nearby context
        sheet_name = self._extract_sheet_name(table_elem)

        return HtmlTable(
            html=str(table_elem),
            rows=rows,
            sheet_name=sheet_name,
            table_index=table_index,
            row_count=row_count,
            col_count=col_count
        )

    def _extract_rows_with_spans(self, table_elem) -> List[List[str]]:
        """
        Extract table rows, handling rowspan and colspan attributes.

        This creates a normalized 2D grid where merged cells are properly expanded:
        - Colspan: Value stored in first column, rest are empty strings
        - Rowspan: Value repeated in first column of each spanned row, rest empty

        This normalization ensures:
        1. Each row has consistent column count
        2. Merged cell values appear in predictable positions
        3. Empty strings mark spanned positions for clarity

        Args:
            table_elem: BeautifulSoup table element

        Returns:
            2D list of cell text content (normalized grid)
        """
        # First pass: collect all cells with their span info
        raw_rows = []
        for tr in table_elem.find_all('tr'):
            cells = []
            for cell in tr.find_all(['td', 'th']):
                # Get cell text (strip whitespace)
                text = cell.get_text(strip=True)

                # Get span attributes
                rowspan = int(cell.get('rowspan', 1))
                colspan = int(cell.get('colspan', 1))

                cells.append({
                    'text': text,
                    'rowspan': rowspan,
                    'colspan': colspan
                })

            if cells:  # Only add non-empty rows
                raw_rows.append(cells)

        if not raw_rows:
            return []

        # Second pass: expand spans into a proper 2D grid
        # Pre-calculate max columns needed
        max_cols = 0
        for row_cells in raw_rows:
            cols_in_row = sum(cell['colspan'] for cell in row_cells)
            max_cols = max(max_cols, cols_in_row)

        # Create empty grid
        grid = []

        # Track which cells are already filled by rowspan from previous rows
        filled_positions = {}  # {(row_idx, col_idx): text}

        for row_idx, row_cells in enumerate(raw_rows):
            current_row = []
            col_idx = 0
            cell_idx = 0

            while col_idx < max_cols:
                # Check if this position is already filled by a rowspan
                if (row_idx, col_idx) in filled_positions:
                    current_row.append(filled_positions[(row_idx, col_idx)])
                    col_idx += 1
                    continue

                # No more cells in this row
                if cell_idx >= len(row_cells):
                    current_row.append('')
                    col_idx += 1
                    continue

                # Get current cell
                cell = row_cells[cell_idx]
                text = cell['text']
                rowspan = cell['rowspan']
                colspan = cell['colspan']

                # NORMALIZATION: Fill colspan with value in first column only
                # This ensures merged cells don't duplicate data
                for c in range(colspan):
                    if col_idx + c < max_cols:
                        # Only first column gets the actual value
                        current_row.append(text if c == 0 else '')

                # NORMALIZATION: Handle rowspan by repeating value in first column
                # This makes it clear which rows the merged cell spans
                if rowspan > 1:
                    for r in range(1, rowspan):
                        for c in range(colspan):
                            # Only first column of each row gets the value
                            filled_positions[(row_idx + r, col_idx + c)] = text if c == 0 else ''

                col_idx += colspan
                cell_idx += 1

            # Ensure all rows have same length
            while len(current_row) < max_cols:
                current_row.append('')

            grid.append(current_row)

        return grid

    def _extract_sheet_name(self, table_elem) -> Optional[str]:
        """
        Try to extract sheet name from HTML context near the table.

        Docling may include sheet names as headings or data attributes.

        Args:
            table_elem: BeautifulSoup table element

        Returns:
            Sheet name if found, None otherwise
        """
        # Strategy 1: Look for data-sheet attribute
        sheet_attr = table_elem.get('data-sheet')
        if sheet_attr:
            return sheet_attr

        # Strategy 2: Look for preceding heading with sheet name pattern
        prev_sibling = table_elem.find_previous_sibling()
        if prev_sibling and prev_sibling.name in ['h1', 'h2', 'h3', 'h4']:
            heading_text = prev_sibling.get_text(strip=True)
            if heading_text:
                return heading_text

        # Strategy 3: Look for parent div with sheet identifier
        parent = table_elem.find_parent('div', {'class': 'sheet'})
        if parent:
            sheet_name = parent.get('data-name') or parent.get('id')
            if sheet_name:
                return sheet_name

        # No sheet name found
        return None

    def get_first_n_rows(self, table: HtmlTable, n: int = 10) -> List[List[str]]:
        """
        Get the first N rows of a table for header detection.

        Args:
            table: HtmlTable object
            n: Number of rows to return (default: 10)

        Returns:
            First N rows (or all rows if table has fewer than N rows)
        """
        return table.rows[:n]


# Convenience function for easy imports
def extract_tables_from_html(html: str) -> List[HtmlTable]:
    """
    Convenience function to extract tables from HTML.

    Args:
        html: HTML string

    Returns:
        List of HtmlTable objects (valid data tables only)
    """
    extractor = HtmlTableExtractor()
    return extractor.extract_tables(html)


def extract_all_tables_from_html(html: str) -> tuple[list[HtmlTable], list[HtmlTable]]:
    """
    Convenience function to extract and classify all tables from HTML.

    Args:
        html: HTML string

    Returns:
        Tuple of (valid_tables, raw_text_tables)
    """
    extractor = HtmlTableExtractor()
    return extractor.extract_all_tables(html)