"""
Large Table Extractor - Handles Excel tables exceeding Docling's processing limits.

This module provides an alternative extraction strategy for very large Excel tables
(>1000 rows, >100 columns) that would crash or timeout with standard Docling processing.

Strategy:
1. Border detection to identify table boundaries without loading full content
2. Header extraction to temporary file for normalization
3. Multi-level header parsing using HeaderExtractor
4. Chunked data processing using iter_rows(values_only=True) for ~100x speedup
5. Structured output compatible with V2 pipeline format

Performance:
- 619K cells processed in ~15 seconds
- Memory efficient with read_only mode
- Automatic detection and routing of large tables
"""

from datetime import datetime
from typing import Any
from uuid import UUID
from dataclasses import dataclass

import structlog
from openpyxl.utils import get_column_letter, column_index_from_string

from src.extraction_v2.detect_border import BorderTableDetector
from src.extraction_v2.excel_loader import load_excel_file
from src.extraction_v2.llm_header_detector import LlmHeaderDetector, HeaderStructure

logger = structlog.get_logger(__name__)


@dataclass
class LargeTableConfig:
    """Configuration for large table extraction."""
    max_rows_for_docling: int = 1000  # Tables larger than this use chunked processing
    header_sample_rows: int = 10      # Number of rows to extract for header detection
    chunk_size: int = 100             # Number of data rows to process at once
    max_rows_to_process: int = -1     # -1 means all rows


@dataclass
class LargeTableResult:
    """Result from large table extraction."""
    success: bool
    sheet_name: str
    table_range: str
    total_rows: int
    rows_processed: int
    column_count: int
    header_depth: int
    headers: dict[str, str]  # column_letter -> header_name
    records: list[dict[str, Any]]
    error: str | None = None
    processing_time: float = 0.0


class LargeTableExtractor:
    """
    Extracts data from very large Excel tables using chunked processing.

    Use this extractor when:
    - Table has >1,000 rows
    - File has >100,000 cells per sheet
    - Standard Docling extraction fails or times out
    - Memory is constrained
    """

    def __init__(self, config: LargeTableConfig | None = None):
        """
        Initialize the large table extractor.

        Args:
            config: Optional configuration, uses defaults if not provided
        """
        self.config = config or LargeTableConfig()
        self.border_detector = BorderTableDetector()
        self.header_detector = LlmHeaderDetector()

    def is_large_table(self, table_info: dict[str, Any]) -> bool:
        """
        Check if a table should use large table processing.

        Args:
            table_info: Table information from border detection

        Returns:
            True if table exceeds max_rows_for_docling threshold
        """
        num_rows = table_info['end_row'] - table_info['start_row'] + 1
        return num_rows > self.config.max_rows_for_docling

    def detect_tables(self, file_path: str) -> list[dict[str, Any]]:
        """
        Detect all tables in Excel file using border detection.

        Args:
            file_path: Path to Excel file

        Returns:
            List of table information dictionaries with size classification
        """
        logger.info("Detecting tables using border detection", file_path=file_path)

        tables = self.border_detector.detect_tables_by_border(file_path)

        # Add size classification to each table
        for table in tables:
            table['is_large'] = self.is_large_table(table)
            table['row_count'] = table['end_row'] - table['start_row'] + 1

        large_count = sum(1 for t in tables if t['is_large'])
        logger.info(
            f"Detected {len(tables)} table(s), {large_count} large",
            total=len(tables),
            large=large_count
        )

        return tables

    def extract_table(
        self,
        file_path: str,
        table_info: dict[str, Any],
        document_id: UUID | None = None
    ) -> LargeTableResult:
        """
        Extract data from a large table using chunked processing.

        Args:
            file_path: Path to Excel file
            table_info: Table information from detect_tables()
            document_id: Optional document ID for tracking

        Returns:
            LargeTableResult with extracted records
        """
        start_time = datetime.now()

        logger.info(
            "Extracting large table",
            sheet=table_info['sheet'],
            range=table_info['range'],
            cells=table_info['cell_count']
        )

        try:
            # Step 1: Extract header rows as HTML for LLM
            header_html = self._extract_header_rows_as_html(file_path, table_info)

            # Step 2: Detect headers using LLM
            headers = self._extract_normalized_headers(header_html, table_info)

            # Step 3: Process data rows in chunks
            records = self._process_data_rows(file_path, table_info, headers)

            processing_time = (datetime.now() - start_time).total_seconds()

            logger.info(
                "Large table extraction complete",
                sheet=table_info['sheet'],
                rows_processed=len(records),
                time=f"{processing_time:.2f}s"
            )

            return LargeTableResult(
                success=True,
                sheet_name=table_info['sheet'],
                table_range=table_info['range'],
                total_rows=table_info['row_count'],
                rows_processed=len(records),
                column_count=len(headers['column_mapping']),
                header_depth=headers.get('header_depth', self.config.header_sample_rows),
                headers=headers['column_mapping'],
                records=records,
                processing_time=processing_time
            )

        except Exception as e:
            logger.error(f"Large table extraction failed: {e}", exc_info=True)
            return LargeTableResult(
                success=False,
                sheet_name=table_info.get('sheet', 'Unknown'),
                table_range=table_info.get('range', 'Unknown'),
                total_rows=table_info.get('row_count', 0),
                rows_processed=0,
                column_count=0,
                header_depth=0,
                headers={},
                records=[],
                error=str(e)
            )

    def _extract_header_rows_as_html(
        self,
        file_path: str,
        table_info: dict[str, Any]
    ) -> str:
        """
        Extract first N rows of table as HTML for LLM header detection.

        Args:
            file_path: Path to source Excel file
            table_info: Table information dictionary

        Returns:
            HTML string with table rows
        """
        logger.debug(
            "Extracting header rows as HTML",
            rows=self.config.header_sample_rows,
            sheet=table_info['sheet']
        )

        # Load source workbook
        source_wb = load_excel_file(file_path, data_only=False)
        source_sheet = source_wb[table_info['sheet']]

        # Get column range
        start_col_idx = column_index_from_string(table_info['start_col'])
        end_col_idx = column_index_from_string(table_info['end_col'])

        # Calculate how many rows to extract
        num_rows = min(
            self.config.header_sample_rows,
            table_info['end_row'] - table_info['start_row'] + 1
        )

        # Build merged cell map for this region
        merged_cells_map = {}  # (row, col) -> {'value': str, 'rowspan': int, 'colspan': int}
        skip_cells = set()  # Cells that are part of a merged range but not the anchor

        for merged_range in source_sheet.merged_cells.ranges:
            if (merged_range.min_row >= table_info['start_row'] and
                merged_range.min_row < table_info['start_row'] + num_rows and
                merged_range.min_col >= start_col_idx and
                merged_range.max_col <= end_col_idx):

                # Get the anchor cell value
                anchor_cell = source_sheet.cell(
                    row=merged_range.min_row, column=merged_range.min_col
                )
                value = anchor_cell.value if anchor_cell.value is not None else ""

                # Calculate rowspan and colspan (relative to our extracted region)
                rowspan = min(
                    merged_range.max_row - merged_range.min_row + 1,
                    table_info['start_row'] + num_rows - merged_range.min_row
                )
                colspan = merged_range.max_col - merged_range.min_col + 1

                # Store anchor info (in relative coordinates)
                rel_row = merged_range.min_row - table_info['start_row']
                rel_col = merged_range.min_col - start_col_idx
                merged_cells_map[(rel_row, rel_col)] = {
                    'value': str(value),
                    'rowspan': rowspan,
                    'colspan': colspan
                }

                # Mark cells to skip (non-anchor cells in merged range)
                for r in range(merged_range.min_row, merged_range.max_row + 1):
                    for c in range(merged_range.min_col, merged_range.max_col + 1):
                        if r >= table_info['start_row'] + num_rows:
                            continue
                        rel_r = r - table_info['start_row']
                        rel_c = c - start_col_idx
                        if (rel_r, rel_c) != (rel_row, rel_col):
                            skip_cells.add((rel_r, rel_c))

        # Build HTML table
        html_rows = []
        for row_offset in range(num_rows):
            source_row = table_info['start_row'] + row_offset
            cells_html = []

            for col_offset in range(end_col_idx - start_col_idx + 1):
                rel_pos = (row_offset, col_offset)

                # Skip cells that are part of a merged range but not the anchor
                if rel_pos in skip_cells:
                    continue

                source_col = start_col_idx + col_offset
                cell = source_sheet.cell(row=source_row, column=source_col)

                # Check if this is a merged cell anchor
                if rel_pos in merged_cells_map:
                    merge_info = merged_cells_map[rel_pos]
                    value = merge_info['value']
                    attrs = []
                    if merge_info['rowspan'] > 1:
                        attrs.append(f'rowspan="{merge_info["rowspan"]}"')
                    if merge_info['colspan'] > 1:
                        attrs.append(f'colspan="{merge_info["colspan"]}"')
                    attr_str = ' ' + ' '.join(attrs) if attrs else ''
                    cells_html.append(f"<th{attr_str}>{value}</th>")
                else:
                    # Regular cell
                    value = str(cell.value) if cell.value is not None else ""
                    # Use <th> for first row, <td> for others
                    tag = "th" if row_offset == 0 else "td"
                    cells_html.append(f"<{tag}>{value}</{tag}>")

            html_rows.append(f"<tr>{''.join(cells_html)}</tr>")

        source_wb.close()

        html = f"<table>{''.join(html_rows)}</table>"
        logger.debug(f"Generated HTML for header detection: {len(html)} chars")

        return html

    def _extract_normalized_headers(
        self,
        header_html: str,
        table_info: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Extract and normalize multi-level headers using LLM.

        Args:
            header_html: HTML string with header rows
            table_info: Table information dictionary

        Returns:
            Dictionary with column_mapping, header_structure, header_depth
        """
        logger.debug("Extracting headers using LLM")

        try:
            # Use LLM to detect headers from HTML
            header_structure: HeaderStructure = self.header_detector.detect_headers(
                header_html
            )

            # Determine header depth from detected header rows
            header_depth = len(header_structure.header_rows) if header_structure.header_rows else 1

            logger.debug(
                "LLM header detection complete",
                header_rows=header_structure.header_rows,
                columns=len(header_structure.headers),
                depth=header_depth
            )

            # Get column range for mapping
            start_col_idx = column_index_from_string(table_info['start_col'])

            # Create column mapping from LLM results
            column_mapping = {}
            for header_info in header_structure.headers:
                col_idx = header_info.get('column', 0)
                hierarchy = header_info.get('hierarchy', [])

                # Convert column index to letter (relative to table start)
                col_letter = get_column_letter(start_col_idx + col_idx)

                # Build header name from hierarchy
                if hierarchy:
                    # Remove consecutive duplicates
                    deduplicated = []
                    for level in hierarchy:
                        if not deduplicated or level != deduplicated[-1]:
                            deduplicated.append(level)

                    # Remove first level (table/sheet name) if multiple levels
                    if len(deduplicated) > 1:
                        deduplicated = deduplicated[1:]

                    header_name = " -> ".join(deduplicated) if deduplicated else f"Column {col_idx}"
                else:
                    header_name = f"Column {col_idx}"

                column_mapping[col_letter] = header_name

            return {
                'column_mapping': column_mapping,
                'header_structure': header_structure,
                'header_depth': header_depth
            }

        except Exception as e:
            logger.warning(f"LLM header detection failed, using column letters: {e}")

            # Fallback: use column letters as headers
            start_col_idx = column_index_from_string(table_info['start_col'])
            end_col_idx = column_index_from_string(table_info['end_col'])
            column_mapping = {}
            for col_idx in range(start_col_idx, end_col_idx + 1):
                col_letter = get_column_letter(col_idx)
                column_mapping[col_letter] = col_letter

            return {
                'column_mapping': column_mapping,
                'header_structure': None,
                'header_depth': self.config.header_sample_rows
            }

    def _process_data_rows(
        self,
        file_path: str,
        table_info: dict[str, Any],
        headers: dict[str, Any]
    ) -> list[dict[str, Any]]:
        """
        Process data rows in chunks using iter_rows for performance.

        Args:
            file_path: Path to source Excel file
            table_info: Table information dictionary
            headers: Normalized header information

        Returns:
            List of row dictionaries
        """
        workbook = load_excel_file(file_path, read_only=True, data_only=True)
        sheet = workbook[table_info['sheet']]

        # Calculate data row range using actual header depth
        actual_header_depth = headers.get('header_depth', self.config.header_sample_rows)
        data_start_row = table_info['start_row'] + actual_header_depth
        data_end_row = table_info['end_row']

        # Apply row limit if configured
        if self.config.max_rows_to_process > 0:
            data_end_row = min(
                data_end_row,
                data_start_row + self.config.max_rows_to_process - 1
            )

        total_data_rows = data_end_row - data_start_row + 1

        logger.debug(
            f"Processing data rows",
            start=data_start_row,
            end=data_end_row,
            total=total_data_rows
        )

        # Get column range
        start_col_idx = column_index_from_string(table_info['start_col'])
        end_col_idx = column_index_from_string(table_info['end_col'])

        rows_data = []
        column_mapping = headers.get('column_mapping', {})

        # Process in chunks using iter_rows (FAST!)
        for chunk_start in range(data_start_row, data_end_row + 1, self.config.chunk_size):
            chunk_end = min(chunk_start + self.config.chunk_size - 1, data_end_row)

            for row in sheet.iter_rows(
                min_row=chunk_start,
                max_row=chunk_end,
                min_col=start_col_idx,
                max_col=end_col_idx,
                values_only=True
            ):
                # Skip empty rows
                if all(
                    value is None or (isinstance(value, str) and not value.strip())
                    for value in row
                ):
                    continue

                row_data = {}
                for col_offset, value in enumerate(row):
                    col_idx = start_col_idx + col_offset
                    col_letter = get_column_letter(col_idx)
                    header_name = column_mapping.get(col_letter, col_letter)
                    row_data[header_name] = value

                rows_data.append(row_data)

        workbook.close()

        logger.debug(f"Processed {len(rows_data)} data rows")

        return rows_data

    def extract_all_large_tables(
        self,
        file_path: str,
        document_id: UUID | None = None
    ) -> list[LargeTableResult]:
        """
        Detect and extract all large tables from an Excel file.

        Args:
            file_path: Path to Excel file
            document_id: Optional document ID for tracking

        Returns:
            List of LargeTableResult for each large table found
        """
        tables = self.detect_tables(file_path)
        large_tables = [t for t in tables if t['is_large']]

        if not large_tables:
            logger.info("No large tables found in file")
            return []

        results = []
        for table in large_tables:
            result = self.extract_table(file_path, table, document_id)
            results.append(result)

        return results

    def to_pipeline_records(
        self,
        result: LargeTableResult,
        document_id: UUID
    ) -> list[dict[str, Any]]:
        """
        Convert LargeTableResult to V2 pipeline record format.

        Args:
            result: Extraction result
            document_id: Document UUID

        Returns:
            List of records in V2 pipeline format
        """
        pipeline_records = []

        for idx, row_data in enumerate(result.records):
            record = {
                'document_id': str(document_id),
                'sheet_name': result.sheet_name,
                'table_range': result.table_range,
                'row_index': idx,
                'data': row_data,
                'extraction_method': 'large_table_chunked',
                'header_depth': result.header_depth
            }
            pipeline_records.append(record)

        return pipeline_records


# Convenience function
def extract_large_tables(
    file_path: str,
    document_id: UUID | None = None,
    config: LargeTableConfig | None = None
) -> list[LargeTableResult]:
    """
    Convenience function to extract all large tables from an Excel file.

    Args:
        file_path: Path to Excel file
        document_id: Optional document ID for tracking
        config: Optional configuration

    Returns:
        List of LargeTableResult for each large table
    """
    extractor = LargeTableExtractor(config)
    return extractor.extract_all_large_tables(file_path, document_id)
