"""Table extraction via booktabs horizontal rule detection."""
from __future__ import annotations

import re

import fitz
from ..types import BBox, Table, TextLine
from ..utils.geometry import cluster_1d


def find_horizontal_rules(page: fitz.Page) -> list[dict]:
    """Find horizontal rules from page drawings (booktabs-style tables).

    Returns list of dicts with keys: y, x0, x1, width (line thickness).
    """
    rules = []
    drawings = page.get_drawings()
    for d in drawings:
        for item in d.get("items", []):
            kind = item[0]
            if kind == "l":  # line
                p1, p2 = item[1], item[2]
                # Horizontal line: y values approximately equal
                if abs(p1.y - p2.y) < 2.0 and abs(p1.x - p2.x) > 30:
                    rules.append({
                        "y": (p1.y + p2.y) / 2,
                        "x0": min(p1.x, p2.x),
                        "x1": max(p1.x, p2.x),
                        "width": d.get("width", 1.0),
                    })
            elif kind == "re":  # rectangle (thin rectangles can be rules)
                rect = item[1]
                if isinstance(rect, fitz.Rect) and rect.height < 3.0 and rect.width > 30:
                    rules.append({
                        "y": (rect.y0 + rect.y1) / 2,
                        "x0": rect.x0,
                        "x1": rect.x1,
                        "width": rect.height,
                    })
    return sorted(rules, key=lambda r: r["y"])


def group_rules_into_tables(rules: list[dict], gap: float = 200.0) -> list[list[dict]]:
    """Group horizontal rules into table regions by vertical proximity."""
    if not rules:
        return []
    groups: list[list[dict]] = [[rules[0]]]
    for r in rules[1:]:
        if r["y"] - groups[-1][-1]["y"] < gap:
            groups[-1].append(r)
        else:
            groups.append([r])
    # A booktabs table needs at least 3 rules (top, mid, bottom)
    return [g for g in groups if len(g) >= 3]


def extract_table_from_rules(
    page: fitz.Page,
    rules: list[dict],
    lines: list[TextLine],
) -> Table | None:
    """Extract a table bounded by horizontal rules.

    Strategy:
    - toprule/bottomrule: first and last rules
    - midrules: everything in between separates header from body
    - Column detection: cluster x-positions of text within table region
    """
    if len(rules) < 3:
        return None

    top_y = rules[0]["y"]
    bottom_y = rules[-1]["y"]
    x0 = min(r["x0"] for r in rules)
    x1 = max(r["x1"] for r in rules)

    # Find lines within table region
    table_lines = [
        l for l in lines
        if l.bbox.y0 >= top_y - 5 and l.bbox.y1 <= bottom_y + 5
        and l.bbox.x0 >= x0 - 10 and l.bbox.x1 <= x1 + 10
    ]

    if not table_lines:
        return None

    # Find midrule y-positions (between top and bottom)
    midrule_ys = [top_y] + [r["y"] for r in rules[1:-1]] + [bottom_y]

    # Segment lines into rows by midrule positions
    row_groups: list[list[TextLine]] = []
    for i in range(len(midrule_ys) - 1):
        y_start = midrule_ys[i]
        y_end = midrule_ys[i + 1]
        row_lines = [l for l in table_lines if l.bbox.midy > y_start and l.bbox.midy < y_end]
        if row_lines:
            # Sub-segment by y-position for multi-line rows
            row_lines.sort(key=lambda l: l.bbox.y0)
            sub_rows = _split_into_rows(row_lines)
            row_groups.extend(sub_rows)

    if not row_groups:
        return None

    # Detect columns from x-positions
    all_spans_x = []
    for row in row_groups:
        for line in row:
            for span in line.spans:
                if span.text.strip():
                    all_spans_x.append(span.bbox.x0)

    if not all_spans_x:
        return None

    col_boundaries = _detect_column_boundaries(all_spans_x, x0, x1)

    # Build cell grid
    cells: list[list[str]] = []
    for row in row_groups:
        row_cells = [""] * len(col_boundaries)
        for line in row:
            for span in line.spans:
                text = span.text.strip()
                if not text:
                    continue
                col_idx = _find_column(span.bbox.x0, col_boundaries)
                if row_cells[col_idx]:
                    row_cells[col_idx] += " " + text
                else:
                    row_cells[col_idx] = text
        cells.append(row_cells)

    # Clean up cell text (rejoin spaced numbers, parentheses, etc.)
    cells = [[_clean_table_cell(cell) for cell in row] for row in cells]

    return Table(
        cells=cells,
        bbox=BBox(x0, top_y, x1, bottom_y),
        page_num=0,  # caller sets this
    )


def _clean_table_cell(text: str) -> str:
    """Clean up table cell text: rejoin spaced numbers and operators."""
    # Join digit . digit patterns: "1 . 0" -> "1.0"
    text = re.sub(r'(\d) \. (\d)', r'\1.\2', text)
    # Join spaced parentheses: "O ( n )" -> "O(n)"
    # Process iteratively since patterns can nest
    prev = None
    while prev != text:
        prev = text
        text = re.sub(r'(\w) \( ', r'\1(', text)
        text = re.sub(r' \)', r')', text)
        text = re.sub(r'\( ', r'(', text)
    # Join spaced commas: "0 , 1" -> "0, 1" (remove leading space)
    text = re.sub(r'(\d) , ', r'\1, ', text)
    # Join "· 10" patterns for scientific notation
    text = re.sub(r'· 10 (\d+)', r'· 10^{\1}', text)
    # Join spaced minus in numbers: "- 4" -> "-4" when at start or after operator
    text = re.sub(r'(?:^|\s)- (\d)', r' -\1', text)
    # Join "3 . 0 e - 4" scientific notation
    text = re.sub(r'(\d+\.\d+) e ([-\u2212]) (\d+)', r'\1e\2\3', text)
    # Also handle "3 . 0 e - 4" with already-joined decimal
    text = re.sub(r'(\d) e ([-\u2212]) (\d)', r'\1e\2\3', text)
    return text.strip()


def _split_into_rows(lines: list[TextLine], gap: float = 6.0) -> list[list[TextLine]]:
    """Split lines into rows based on y-position gaps.

    Uses y0-to-y0 distance to detect row boundaries: lines in the same
    table row share similar y0 (within *gap*), while lines in different
    rows have y0 separated by a full line height (~10pt+).
    """
    if not lines:
        return []
    rows: list[list[TextLine]] = [[lines[0]]]
    for line in lines[1:]:
        # Compare y0 of current line to average y0 of current row
        row_y0 = sum(l.bbox.y0 for l in rows[-1]) / len(rows[-1])
        if abs(line.bbox.y0 - row_y0) > gap:
            rows.append([])
        rows[-1].append(line)
    return rows


def _detect_column_boundaries(x_positions: list[float], table_x0: float, table_x1: float) -> list[tuple[float, float]]:
    """Detect column boundaries from span x-positions."""
    if not x_positions:
        return [(table_x0, table_x1)]

    x_sorted = sorted(set(round(x, 1) for x in x_positions))
    if len(x_sorted) <= 1:
        return [(table_x0, table_x1)]

    # Find gaps between x-positions
    gaps = []
    for i in range(1, len(x_sorted)):
        gap = x_sorted[i] - x_sorted[i - 1]
        if gap > 15:  # significant gap
            gaps.append((x_sorted[i - 1], x_sorted[i], gap))

    if not gaps:
        return [(table_x0, table_x1)]

    # Use largest gaps as column separators
    boundaries = [table_x0]
    for g in sorted(gaps, key=lambda g: g[2], reverse=True):
        mid = (g[0] + g[1]) / 2
        boundaries.append(mid)
    boundaries.append(table_x1)
    boundaries.sort()

    return [(boundaries[i], boundaries[i + 1]) for i in range(len(boundaries) - 1)]


def _find_column(x: float, boundaries: list[tuple[float, float]]) -> int:
    """Find which column a given x-position belongs to."""
    for i, (x0, x1) in enumerate(boundaries):
        if x0 - 10 <= x <= x1 + 10:
            return i
    # Fallback: closest column
    return min(range(len(boundaries)), key=lambda i: abs(x - (boundaries[i][0] + boundaries[i][1]) / 2))


def table_to_pipe(table: Table) -> str:
    """Render a Table as a Markdown pipe table."""
    if not table.cells:
        return ""

    # Normalize column count
    num_cols = table.num_cols
    rows = [row + [""] * (num_cols - len(row)) for row in table.cells]

    # Calculate column widths
    widths = [max(len(rows[r][c]) for r in range(len(rows))) for c in range(num_cols)]
    widths = [max(w, 3) for w in widths]  # minimum 3 chars

    def format_row(row: list[str]) -> str:
        cells = [cell.ljust(widths[i]) for i, cell in enumerate(row)]
        return "| " + " | ".join(cells) + " |"

    lines = [format_row(rows[0])]
    lines.append("| " + " | ".join("-" * w for w in widths) + " |")
    for row in rows[1:]:
        lines.append(format_row(row))

    return "\n".join(lines)


def extract_tables(page: fitz.Page, lines: list[TextLine], page_num: int) -> list[Table]:
    """Extract all tables from a page."""
    rules = find_horizontal_rules(page)
    table_groups = group_rules_into_tables(rules)
    tables = []
    for group in table_groups:
        table = extract_table_from_rules(page, group, lines)
        if table:
            table.page_num = page_num
            tables.append(table)
    return tables
