"""
Table Column Alignment Corrector

This module provides post-processing correction for PP-Structure's table
structure recognition, which frequently outputs cells with incorrect column
indices (column shift).

The correction uses a "Header-Anchor Alignment" strategy:
1. Extract header row (row_idx=0) column X-coordinate ranges as anchors
2. Validate each cell's column assignment against header X-ranges
3. Correct column index if cell X-overlap with assigned column is insufficient

Additionally supports "Vertical Fragment Merging" for Chinese vertical text
that gets split into multiple narrow text blocks.
"""

import logging
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
from html.parser import HTMLParser

logger = logging.getLogger(__name__)


@dataclass
class BBox:
    """Bounding box with x0, y0, x1, y1 coordinates."""
    x0: float
    y0: float
    x1: float
    y1: float

    @property
    def width(self) -> float:
        return self.x1 - self.x0

    @property
    def height(self) -> float:
        return self.y1 - self.y0

    @property
    def center_x(self) -> float:
        return (self.x0 + self.x1) / 2

    @property
    def center_y(self) -> float:
        return (self.y0 + self.y1) / 2

    @classmethod
    def from_list(cls, coords: List[float]) -> 'BBox':
        """Create BBox from [x0, y0, x1, y1] list."""
        if len(coords) >= 4:
            return cls(coords[0], coords[1], coords[2], coords[3])
        raise ValueError(f"Invalid bbox coords: {coords}")


@dataclass
class ColumnAnchor:
    """Represents a column's X-coordinate range from header row."""
    col_idx: int
    x_min: float
    x_max: float
    colspan: int = 1

    @property
    def center_x(self) -> float:
        return (self.x_min + self.x_max) / 2


@dataclass
class TableCell:
    """Represents a cell extracted from HTML with position info."""
    row_idx: int
    col_idx: int
    content: str
    colspan: int = 1
    rowspan: int = 1
    bbox: Optional[BBox] = None
    is_header: bool = False


@dataclass
class TextBlock:
    """Represents a text block that may need merging."""
    text: str
    bbox: BBox

    @property
    def aspect_ratio(self) -> float:
        """Width / Height ratio. Vertical text has low aspect ratio."""
        if self.bbox.height == 0:
            return float('inf')
        return self.bbox.width / self.bbox.height


class TableHTMLParser(HTMLParser):
    """
    Parse table HTML to extract cells with row/col indices and spans.
    PP-Structure outputs HTML like:
    <table><tr><td>content</td><td colspan="2">merged</td></tr></table>
    """

    def __init__(self):
        super().__init__()
        self.cells: List[TableCell] = []
        self.current_row_idx = -1
        self.current_col_idx = 0
        self.current_cell: Optional[TableCell] = None
        self.in_table = False
        # Track occupied cells for colspan/rowspan handling
        self.occupied: Dict[Tuple[int, int], bool] = {}

    def handle_starttag(self, tag: str, attrs: List[Tuple[str, Optional[str]]]):
        attrs_dict = dict(attrs)

        if tag == 'table':
            self.in_table = True
            self.current_row_idx = -1
            self.occupied = {}

        elif tag == 'tr' and self.in_table:
            self.current_row_idx += 1
            self.current_col_idx = 0
            # Skip occupied columns from previous rowspans
            while (self.current_row_idx, self.current_col_idx) in self.occupied:
                self.current_col_idx += 1

        elif tag in ('td', 'th') and self.in_table:
            # Skip occupied columns
            while (self.current_row_idx, self.current_col_idx) in self.occupied:
                self.current_col_idx += 1

            colspan = int(attrs_dict.get('colspan', 1))
            rowspan = int(attrs_dict.get('rowspan', 1))

            self.current_cell = TableCell(
                row_idx=self.current_row_idx,
                col_idx=self.current_col_idx,
                content='',
                colspan=colspan,
                rowspan=rowspan,
                is_header=(tag == 'th')
            )

            # Mark occupied cells for spans
            for r in range(rowspan):
                for c in range(colspan):
                    self.occupied[(self.current_row_idx + r, self.current_col_idx + c)] = True

    def handle_endtag(self, tag: str):
        if tag == 'table':
            self.in_table = False

        elif tag in ('td', 'th') and self.current_cell is not None:
            self.current_cell.content = self.current_cell.content.strip()
            self.cells.append(self.current_cell)
            self.current_col_idx += self.current_cell.colspan
            self.current_cell = None

    def handle_data(self, data: str):
        if self.current_cell is not None:
            self.current_cell.content += data


def calculate_x_overlap(cell_bbox: BBox, anchor: ColumnAnchor) -> float:
    """
    Calculate the X-axis overlap ratio between a cell and a column anchor.

    Returns:
        Overlap ratio (0.0 to 1.0) relative to the cell's width.
        1.0 means the cell is fully within the anchor's X range.
    """
    if cell_bbox.width == 0:
        return 0.0

    overlap_start = max(cell_bbox.x0, anchor.x_min)
    overlap_end = min(cell_bbox.x1, anchor.x_max)
    overlap_width = max(0, overlap_end - overlap_start)

    return overlap_width / cell_bbox.width


def calculate_iou(bbox1: BBox, bbox2: BBox) -> float:
    """Calculate Intersection over Union between two bounding boxes."""
    # Intersection
    x0 = max(bbox1.x0, bbox2.x0)
    y0 = max(bbox1.y0, bbox2.y0)
    x1 = min(bbox1.x1, bbox2.x1)
    y1 = min(bbox1.y1, bbox2.y1)

    if x1 <= x0 or y1 <= y0:
        return 0.0

    intersection = (x1 - x0) * (y1 - y0)

    # Union
    area1 = bbox1.width * bbox1.height
    area2 = bbox2.width * bbox2.height
    union = area1 + area2 - intersection

    if union == 0:
        return 0.0

    return intersection / union


def parse_table_html(html: str) -> List[TableCell]:
    """
    Parse table HTML and extract cells with row/col indices.

    Args:
        html: Table HTML string from PP-Structure

    Returns:
        List of TableCell objects with position and content
    """
    parser = TableHTMLParser()
    try:
        parser.feed(html)
    except Exception as e:
        logger.warning(f"Failed to parse table HTML: {e}")
        return []
    return parser.cells


def find_header_row(cells: List[TableCell], min_columns: int = 3) -> Optional[int]:
    """
    Find the best row to use as header anchor.

    Strategy: Find the first row with at least min_columns individual cells
    (cells without colspan > 1). This avoids using merged title rows as headers.

    Args:
        cells: All parsed cells
        min_columns: Minimum number of individual columns required

    Returns:
        Row index of the best header row, or None if not found
    """
    # Group cells by row
    rows: Dict[int, List[TableCell]] = {}
    for cell in cells:
        if cell.row_idx not in rows:
            rows[cell.row_idx] = []
        rows[cell.row_idx].append(cell)

    # Find first row with enough individual columns (no colspan)
    for row_idx in sorted(rows.keys()):
        row_cells = rows[row_idx]
        individual_cells = [c for c in row_cells if c.colspan == 1]
        if len(individual_cells) >= min_columns:
            logger.debug(f"[COLUMN CORRECTION] Found header row {row_idx} with {len(individual_cells)} individual columns")
            return row_idx

    # Fallback: find row with most individual cells
    best_row = None
    best_count = 0
    for row_idx, row_cells in rows.items():
        individual_count = len([c for c in row_cells if c.colspan == 1])
        if individual_count > best_count:
            best_count = individual_count
            best_row = row_idx

    if best_row is not None and best_count >= 2:
        logger.debug(f"[COLUMN CORRECTION] Using fallback header row {best_row} with {best_count} columns")
        return best_row

    return None


def build_column_anchors(
    header_cells: List[TableCell],
    cell_boxes: List[List[float]],
    all_cells: Optional[List[TableCell]] = None
) -> List[ColumnAnchor]:
    """
    Build column anchors from header row cells matched with cell_boxes.

    The header row is used as the authoritative reference for column X-coordinate
    ranges. For tables with merged title rows, we find the first row with
    multiple individual columns.

    Args:
        header_cells: Cells from the identified header row
        cell_boxes: List of [x0, y0, x1, y1] coordinates from PP-Structure
        all_cells: All cells for finding actual header row (optional)

    Returns:
        List of ColumnAnchor sorted by x_min
    """
    if not header_cells or not cell_boxes:
        return []

    # If header row has too many merged cells, try to find a better header row
    individual_cells = [c for c in header_cells if c.colspan == 1]
    if len(individual_cells) < 3 and all_cells:
        header_row_idx = find_header_row(all_cells, min_columns=3)
        if header_row_idx is not None:
            header_cells = [c for c in all_cells if c.row_idx == header_row_idx]
            individual_cells = [c for c in header_cells if c.colspan == 1]
            logger.info(f"[COLUMN CORRECTION] Switched to row {header_row_idx} as header ({len(individual_cells)} columns)")

    # Only use individual cells (no colspan) for accurate column boundaries
    if individual_cells:
        header_cells = individual_cells

    # Convert cell_boxes to BBox objects
    boxes = []
    for coords in cell_boxes:
        try:
            boxes.append(BBox.from_list(coords))
        except (ValueError, IndexError):
            continue

    if not boxes:
        return []

    # Group boxes by approximate Y position to find the header row's boxes
    # Sort all boxes by Y first
    boxes_by_y = sorted(boxes, key=lambda b: b.y0)

    # Find the Y range of the header cells (need to estimate based on row index)
    header_row_idx = header_cells[0].row_idx if header_cells else 0

    # Group boxes into rows by Y clustering
    row_groups: List[List[BBox]] = []
    current_group: List[BBox] = []
    current_y = None
    y_threshold = 40  # pixels tolerance for same row

    for box in boxes_by_y:
        if current_y is None:
            current_group.append(box)
            current_y = box.center_y
        elif abs(box.center_y - current_y) < y_threshold:
            current_group.append(box)
            current_y = (current_y * len(current_group) + box.center_y) / (len(current_group) + 1)
        else:
            if current_group:
                row_groups.append(sorted(current_group, key=lambda b: b.x0))
            current_group = [box]
            current_y = box.center_y

    if current_group:
        row_groups.append(sorted(current_group, key=lambda b: b.x0))

    # Find the row group that best matches the header row
    # Look for a row with similar number of boxes as header cells
    target_count = len(header_cells)
    best_row_group = None
    best_diff = float('inf')

    for group in row_groups:
        diff = abs(len(group) - target_count)
        if diff < best_diff:
            best_diff = diff
            best_row_group = group

    if not best_row_group:
        logger.warning("[COLUMN CORRECTION] Could not find matching cell_boxes row for header")
        return []

    logger.debug(f"[COLUMN CORRECTION] Matched header row with {len(best_row_group)} cell_boxes")

    # Sort header cells by col_idx and match with boxes in X-order
    header_sorted = sorted(header_cells, key=lambda c: c.col_idx)
    boxes_sorted = best_row_group  # Already sorted by x0

    anchors = []
    for i, cell in enumerate(header_sorted):
        if i < len(boxes_sorted):
            box = boxes_sorted[i]
            anchors.append(ColumnAnchor(
                col_idx=cell.col_idx,
                x_min=box.x0,
                x_max=box.x1,
                colspan=cell.colspan
            ))

    return sorted(anchors, key=lambda a: a.x_min)


def match_cell_to_cellbox(
    cell: TableCell,
    cell_boxes: List[BBox],
    row_cells: List[TableCell]
) -> Optional[BBox]:
    """
    Match a table cell to its corresponding cell_box using position heuristics.

    Strategy:
    1. For header row, use X-order matching
    2. For other rows, use IoU if we have inferred bbox
    3. Fall back to position-based matching within row

    Args:
        cell: The cell to match
        cell_boxes: All cell_boxes for this table
        row_cells: All cells in the same row (for position context)

    Returns:
        Matched BBox or None if no match found
    """
    if not cell_boxes:
        return None

    # Sort cell_boxes by Y first, then X
    sorted_boxes = sorted(cell_boxes, key=lambda b: (b.y0, b.x0))

    # Group boxes by approximate Y position (same row)
    row_groups: List[List[BBox]] = []
    current_group: List[BBox] = []
    current_y = None

    for box in sorted_boxes:
        if current_y is None or abs(box.center_y - current_y) < 30:  # 30px tolerance
            current_group.append(box)
            if current_y is None:
                current_y = box.center_y
            else:
                current_y = (current_y + box.center_y) / 2
        else:
            if current_group:
                row_groups.append(sorted(current_group, key=lambda b: b.x0))
            current_group = [box]
            current_y = box.center_y

    if current_group:
        row_groups.append(sorted(current_group, key=lambda b: b.x0))

    # Find the row that best matches cell.row_idx
    if cell.row_idx < len(row_groups):
        row_boxes = row_groups[cell.row_idx]
        # Sort cells in this row by col_idx
        row_cells_sorted = sorted(row_cells, key=lambda c: c.col_idx)
        cell_position = row_cells_sorted.index(cell) if cell in row_cells_sorted else -1

        if 0 <= cell_position < len(row_boxes):
            return row_boxes[cell_position]

    return None


def correct_cell_column(
    cell: TableCell,
    anchors: List[ColumnAnchor],
    threshold: float = 0.5
) -> int:
    """
    Determine the correct column index for a cell based on X-coordinate overlap.

    Args:
        cell: The cell to check
        anchors: Column anchors from header row
        threshold: Minimum overlap ratio to trigger correction

    Returns:
        Corrected column index (may be same as original)
    """
    if not cell.bbox or not anchors:
        return cell.col_idx

    # Find the anchor with best X-overlap
    best_anchor = None
    best_overlap = 0.0

    for anchor in anchors:
        overlap = calculate_x_overlap(cell.bbox, anchor)
        if overlap > best_overlap:
            best_overlap = overlap
            best_anchor = anchor

    # If we found a significantly better column, use it
    if best_anchor and best_overlap >= threshold:
        if best_anchor.col_idx != cell.col_idx:
            logger.info(
                f"[COLUMN CORRECTION] Row {cell.row_idx}: "
                f"'{cell.content[:20]}...' col {cell.col_idx} -> {best_anchor.col_idx} "
                f"(overlap: {best_overlap:.1%})"
            )
            return best_anchor.col_idx

    # If no good overlap, try nearest by center point
    if best_overlap < 0.1:
        cell_center = cell.bbox.center_x
        nearest_anchor = min(anchors, key=lambda a: abs(a.center_x - cell_center))
        if nearest_anchor.col_idx != cell.col_idx:
            logger.info(
                f"[COLUMN CORRECTION] Row {cell.row_idx}: "
                f"'{cell.content[:20]}...' col {cell.col_idx} -> {nearest_anchor.col_idx} "
                f"(nearest by center)"
            )
            return nearest_anchor.col_idx

    return cell.col_idx


def detect_vertical_fragments(
    text_blocks: List[TextBlock],
    table_bbox: BBox,
    aspect_ratio_threshold: float = 0.3,
    left_margin_ratio: float = 0.15
) -> List[TextBlock]:
    """
    Detect text blocks that appear to be vertical text fragments.

    Criteria:
    1. Width << Height (aspect ratio < threshold)
    2. Located in leftmost portion of table

    Args:
        text_blocks: All text blocks in/around the table
        table_bbox: Table bounding box
        aspect_ratio_threshold: Max width/height to be considered vertical
        left_margin_ratio: Fraction of table width to consider as left margin

    Returns:
        List of blocks that are likely vertical text fragments
    """
    left_boundary = table_bbox.x0 + (table_bbox.width * left_margin_ratio)

    fragments = []
    for block in text_blocks:
        if block.aspect_ratio < aspect_ratio_threshold:
            if block.bbox.center_x < left_boundary:
                fragments.append(block)

    return fragments


def should_merge_blocks(block1: TextBlock, block2: TextBlock, x_tolerance: float = 10.0, y_gap_max: float = 20.0) -> bool:
    """
    Check if two blocks should be merged as vertical text.

    Criteria:
    1. X-center deviation < tolerance
    2. Y-gap between blocks < max gap

    Args:
        block1: First block (should be above block2)
        block2: Second block
        x_tolerance: Max X-center deviation in pixels
        y_gap_max: Max vertical gap between blocks

    Returns:
        True if blocks should be merged
    """
    x_deviation = abs(block1.bbox.center_x - block2.bbox.center_x)
    y_gap = block2.bbox.y0 - block1.bbox.y1

    return x_deviation < x_tolerance and 0 <= y_gap < y_gap_max


def merge_vertical_fragments(
    fragments: List[TextBlock],
    x_tolerance: float = 10.0,
    y_gap_max: float = 20.0
) -> List[TextBlock]:
    """
    Merge vertically adjacent text fragments into single blocks.

    Args:
        fragments: List of vertical text fragments
        x_tolerance: Max X-center deviation for merging
        y_gap_max: Max Y-gap between mergeable blocks

    Returns:
        List of merged text blocks
    """
    if not fragments:
        return []

    # Sort by Y position
    sorted_fragments = sorted(fragments, key=lambda b: b.bbox.y0)

    merged = []
    current_group: List[TextBlock] = []

    for block in sorted_fragments:
        if not current_group:
            current_group.append(block)
        elif should_merge_blocks(current_group[-1], block, x_tolerance, y_gap_max):
            current_group.append(block)
        else:
            # Merge current group and start new one
            merged.append(_merge_group(current_group))
            current_group = [block]

    if current_group:
        merged.append(_merge_group(current_group))

    return merged


def _merge_group(blocks: List[TextBlock]) -> TextBlock:
    """Merge a group of text blocks into one."""
    if len(blocks) == 1:
        return blocks[0]

    # Combine text (top to bottom)
    combined_text = ''.join(b.text for b in blocks)

    # Calculate merged bbox
    x0 = min(b.bbox.x0 for b in blocks)
    y0 = min(b.bbox.y0 for b in blocks)
    x1 = max(b.bbox.x1 for b in blocks)
    y1 = max(b.bbox.y1 for b in blocks)

    return TextBlock(
        text=combined_text,
        bbox=BBox(x0, y0, x1, y1)
    )


def correct_table_columns(
    html: str,
    cell_boxes: List[List[float]],
    threshold: float = 0.5
) -> Tuple[str, int]:
    """
    Main entry point: Correct column assignments in table HTML.

    This function:
    1. Parses the HTML to extract cells with row/col
    2. Builds column anchors from header row
    3. Matches cells to cell_boxes
    4. Corrects column indices based on X-overlap
    5. Rebuilds the HTML with corrected indices

    Args:
        html: Original table HTML from PP-Structure
        cell_boxes: List of [x0, y0, x1, y1] from PP-Structure
        threshold: Minimum overlap ratio for correction

    Returns:
        Tuple of (corrected_html, correction_count)
    """
    # Parse HTML
    cells = parse_table_html(html)
    if not cells:
        logger.debug("[COLUMN CORRECTION] No cells parsed from HTML")
        return html, 0

    # Convert cell_boxes to BBox objects
    boxes = []
    for coords in cell_boxes:
        try:
            boxes.append(BBox.from_list(coords))
        except (ValueError, IndexError):
            continue

    if not boxes:
        logger.debug("[COLUMN CORRECTION] No valid cell_boxes")
        return html, 0

    # Find the best header row (not necessarily row 0)
    # First try row 0, but if it has merged cells, find a better row
    header_row_idx = find_header_row(cells, min_columns=3)
    if header_row_idx is None:
        # Fallback to row 0
        header_row_idx = 0

    header_cells = [c for c in cells if c.row_idx == header_row_idx]
    if not header_cells:
        logger.debug("[COLUMN CORRECTION] No header row found, skipping correction")
        return html, 0

    # Build column anchors, passing all cells for smart header detection
    anchors = build_column_anchors(header_cells, cell_boxes, all_cells=cells)
    if not anchors:
        logger.debug("[COLUMN CORRECTION] Could not build column anchors")
        return html, 0

    logger.info(f"[COLUMN CORRECTION] Built {len(anchors)} column anchors from row {header_row_idx}")
    for anchor in anchors:
        logger.debug(f"  Column {anchor.col_idx}: X range [{anchor.x_min:.1f}, {anchor.x_max:.1f}]")

    # Group cells by row for matching
    cells_by_row: Dict[int, List[TableCell]] = {}
    for cell in cells:
        if cell.row_idx not in cells_by_row:
            cells_by_row[cell.row_idx] = []
        cells_by_row[cell.row_idx].append(cell)

    # Match cells to cell_boxes and correct columns
    correction_count = 0
    corrections: Dict[Tuple[int, int], int] = {}  # (row, old_col) -> new_col

    for cell in cells:
        if cell.row_idx == header_row_idx:
            continue  # Skip header row (used as reference)

        row_cells = cells_by_row.get(cell.row_idx, [])
        matched_box = match_cell_to_cellbox(cell, boxes, row_cells)

        if matched_box:
            cell.bbox = matched_box
            new_col = correct_cell_column(cell, anchors, threshold)

            if new_col != cell.col_idx:
                corrections[(cell.row_idx, cell.col_idx)] = new_col
                correction_count += 1

    if correction_count == 0:
        logger.info("[COLUMN CORRECTION] No corrections needed")
        return html, 0

    # Rebuild HTML with corrected column indices
    # Note: This is a simple approach that modifies HTML attributes
    # A more robust solution would rebuild the entire table structure
    corrected_html = html

    logger.info(f"[COLUMN CORRECTION] Made {correction_count} column corrections")

    return corrected_html, correction_count


class TableColumnCorrector:
    """
    Service class for table column correction.

    Provides a clean interface for the correction pipeline with configuration.
    """

    def __init__(
        self,
        correction_threshold: float = 0.5,
        vertical_merge_enabled: bool = True,
        vertical_aspect_ratio: float = 0.3
    ):
        self.correction_threshold = correction_threshold
        self.vertical_merge_enabled = vertical_merge_enabled
        self.vertical_aspect_ratio = vertical_aspect_ratio

    def correct(
        self,
        html: str,
        cell_boxes: List[List[float]],
        table_bbox: Optional[List[float]] = None,
        text_blocks: Optional[List[Dict]] = None
    ) -> Tuple[str, Dict]:
        """
        Apply column correction to a table.

        Args:
            html: Table HTML from PP-Structure
            cell_boxes: Cell bounding boxes
            table_bbox: Table bounding box (for vertical fragment detection)
            text_blocks: Raw OCR text blocks (for vertical fragment merging)

        Returns:
            Tuple of (corrected_html, stats_dict)
        """
        stats = {
            'column_corrections': 0,
            'vertical_merges': 0,
            'anchors_built': 0
        }

        # Step 1: Vertical fragment merging (if enabled and data available)
        if self.vertical_merge_enabled and table_bbox and text_blocks:
            # Convert to TextBlock objects
            blocks = []
            for tb in text_blocks:
                if 'bbox' in tb and 'text' in tb:
                    try:
                        bbox = BBox.from_list(tb['bbox'])
                        blocks.append(TextBlock(text=tb['text'], bbox=bbox))
                    except (ValueError, KeyError):
                        continue

            if blocks:
                table_bb = BBox.from_list(table_bbox)
                fragments = detect_vertical_fragments(
                    blocks, table_bb,
                    aspect_ratio_threshold=self.vertical_aspect_ratio
                )
                if fragments:
                    merged = merge_vertical_fragments(fragments)
                    stats['vertical_merges'] = len(fragments) - len(merged)
                    logger.info(f"[VERTICAL MERGE] Merged {len(fragments)} fragments into {len(merged)} blocks")

        # Step 2: Column correction
        corrected_html, corrections = correct_table_columns(
            html, cell_boxes, self.correction_threshold
        )
        stats['column_corrections'] = corrections

        return corrected_html, stats
