"""Layout analysis: page extraction, block classification, column detection."""
from __future__ import annotations

import re
import fitz
from ..types import (
    BBox, Span, TextLine, DocumentBlock, BlockType, PipelineConfig,
)
from ..utils.fonts import FontHierarchy, is_math_font, build_font_hierarchy
from ..utils.geometry import merge_bboxes, find_columns, cluster_1d


def _smart_title_case(text: str) -> str:
    """Title-case ALL-CAPS text while preserving likely acronyms.

    Uses .title() as a base, then fixes:
    - Words/parts containing digits are preserved (GPT-3, 175B)
    - Short (2-3 char) all-uppercase non-article words are kept (XXL, NLP, AI)
    """
    _SMALL_WORDS = frozenset({
        # 1-2 letter articles/prepositions/pronouns
        "a", "an", "as", "at", "be", "by", "do", "go", "he", "if",
        "in", "is", "it", "my", "no", "of", "on", "or", "so", "to",
        "up", "us", "vs", "we",
        # 3-letter common English words (not acronyms)
        "the", "and", "but", "for", "nor", "yet", "all", "any", "are",
        "can", "did", "end", "few", "got", "had", "has", "her", "him",
        "his", "how", "its", "let", "low", "may", "men", "new", "not",
        "now", "old", "one", "our", "out", "own", "per", "put", "ran",
        "run", "saw", "say", "set", "she", "six", "ten", "the", "too",
        "top", "try", "two", "use", "via", "was", "way", "who", "why",
        "won", "yes", "big", "cut", "far", "hot", "key", "net", "red",
    })

    # Start with basic .title() then fix specific patterns
    result_words = []
    for idx, w in enumerate(text.split()):
        # Words with digits: preserve entirely (GPT-3, 175B, R2D2)
        if any(c.isdigit() for c in w):
            result_words.append(w)
            continue
        # Hyphenated words: handle each part
        if "-" in w:
            parts = w.split("-")
            new_parts = []
            for p in parts:
                core = p.rstrip(".,;:'\"")
                if (core.isupper() and core.isalpha()
                        and 2 <= len(core) <= 3
                        and core.lower() not in _SMALL_WORDS):
                    new_parts.append(p)  # Short acronym: keep
                else:
                    new_parts.append(p.title())
            result_words.append("-".join(new_parts))
            continue
        # Short all-uppercase words: keep if likely acronym
        core = w.rstrip(".,;:'\"")
        if (core.isupper() and core.isalpha()
                and 2 <= len(core) <= 3
                and core.lower() not in _SMALL_WORDS):
            result_words.append(w)  # Short acronym
            continue
        # Default: title case
        result_words.append(w.title())
    return " ".join(result_words)


def extract_spans_from_page(page: fitz.Page) -> list[Span]:
    """Extract all text spans from a page."""
    spans = []
    blocks = page.get_text("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)["blocks"]
    for block in blocks:
        if block.get("type") != 0:
            continue
        for line in block.get("lines", []):
            for span in line.get("spans", []):
                text = span.get("text", "")
                if not text:
                    continue
                bbox = BBox.from_tuple(span["bbox"])
                origin = tuple(span.get("origin", (bbox.x0, bbox.y0)))
                spans.append(Span(
                    text=text,
                    font=span.get("font", ""),
                    size=span.get("size", 12.0),
                    flags=span.get("flags", 0),
                    color=span.get("color", 0),
                    bbox=bbox,
                    origin=origin,
                ))
    return spans


_ARXIV_SIDEBAR_RE = re.compile(r"arXiv:\d|^\d{4}\.\d{4,5}")


def extract_lines_from_page(page: fitz.Page) -> list[TextLine]:
    """Extract text lines from a page, preserving span info."""
    lines = []
    blocks = page.get_text("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)["blocks"]
    page_width = page.rect.width
    for block in blocks:
        if block.get("type") != 0:
            continue
        for line_data in block.get("lines", []):
            line_spans = []
            for span in line_data.get("spans", []):
                text = span.get("text", "")
                if not text:
                    continue
                # CMEX fonts encode large delimiters as ASCII chars — strip
                # them to avoid garbled text in tables and body output
                font_name = span.get("font", "")
                if font_name.startswith(("CMEX", "cmex")):
                    text = "".join(
                        c for c in text
                        if c in ' \n\t' or ord(c) > 0x7F
                    )
                    if not text.strip():
                        continue
                bbox = BBox.from_tuple(span["bbox"])
                origin = tuple(span.get("origin", (bbox.x0, bbox.y0)))
                line_spans.append(Span(
                    text=text,
                    font=span.get("font", ""),
                    size=span.get("size", 12.0),
                    flags=span.get("flags", 0),
                    color=span.get("color", 0),
                    bbox=bbox,
                    origin=origin,
                ))
            if line_spans:
                line_bbox = merge_bboxes([s.bbox for s in line_spans])
                line = TextLine(spans=line_spans, bbox=line_bbox)
                # Filter arXiv sidebar stamps: narrow text in the left margin
                # with arXiv ID pattern (rotated 90° text spanning full page height)
                if (line_bbox.x1 < page_width * 0.1
                        and line_bbox.height > 100
                        and _ARXIV_SIDEBAR_RE.search(line.text)):
                    continue
                lines.append(line)
    return _merge_same_y_lines(lines, page_width=page.rect.width)


def _merge_same_y_lines(
    lines: list[TextLine], y_tol: float = 3.0, page_width: float = 612.0,
) -> list[TextLine]:
    """Merge text lines at nearly the same vertical position.

    Handles equation fragments (e.g. fractions where numerator and denominator
    are separate PDF 'line' elements) that should be a single logical line.
    Guards against cross-column merges using page_width.
    """
    if len(lines) <= 1:
        return lines

    sorted_lines = sorted(lines, key=lambda l: (l.bbox.y0, l.bbox.x0))
    merged: list[TextLine] = []
    for line in sorted_lines:
        if merged:
            prev = merged[-1]
            y_close = abs(line.bbox.y0 - prev.bbox.y0) < y_tol
            x_gap = line.bbox.x0 - prev.bbox.x1
            # Allow larger gaps for heading-sized bold text (>12pt fonts
            # have wider word spacing) but NOT for body-sized bold text
            # (which may be separate author name blocks)
            max_size = max(s.size for s in line.spans)
            x_gap_limit = 25 if max_size > 12 else 10
            x_adjacent = -50 < x_gap < x_gap_limit
            # Only merge lines with reasonably sized text (>= 5pt)
            both_reasonable = (max(s.size for s in line.spans) >= 5
                               and max(s.size for s in prev.spans) >= 5)
            # Don't merge lines that would span > 60% of page (cross-column)
            combined_width = max(line.bbox.x1, prev.bbox.x1) - min(line.bbox.x0, prev.bbox.x0)
            not_cross_column = combined_width < page_width * 0.6
            if y_close and x_adjacent and both_reasonable and not_cross_column:
                # Insert a space span between the two lines if there's a gap
                new_spans = list(prev.spans)
                if x_gap > 1:
                    # Create a synthetic space span using the prev line's font
                    last = prev.spans[-1]
                    space_bbox = BBox(last.bbox.x1, last.bbox.y0,
                                     line.spans[0].bbox.x0, last.bbox.y1)
                    new_spans.append(Span(
                        text=" ", font=last.font, size=last.size,
                        flags=last.flags, color=last.color, bbox=space_bbox,
                        origin=(space_bbox.x0, last.origin[1]),
                    ))
                new_spans.extend(line.spans)
                new_spans.sort(key=lambda s: s.bbox.x0)
                new_bbox = merge_bboxes([s.bbox for s in new_spans])
                merged[-1] = TextLine(spans=new_spans, bbox=new_bbox)
                continue
        merged.append(line)
    return merged


def _dehyphenate(text: str) -> str:
    """Rejoin words that were hyphenated across line breaks.

    'fac- tors' → 'factors', but preserve real hyphens like 'self-attention'.
    """
    # Pattern: lowercase letter, hyphen, space(s), lowercase letter
    return re.sub(r'([a-z])-\s+([a-z])', r'\1\2', text)


def _join_body_lines(lines: list[TextLine]) -> str:
    """Join body text lines with dehyphenation."""
    texts = [l.text.strip() for l in lines]
    if not texts:
        return ""
    # Join with spaces, then dehyphenate
    joined = " ".join(texts)
    return _dehyphenate(joined)


def _is_in_any_bbox(line: TextLine, exclude_bboxes: list[BBox]) -> bool:
    """Check if a line falls within any of the exclusion bounding boxes."""
    for bbox in exclude_bboxes:
        if (line.bbox.y0 >= bbox.y0 - 3 and line.bbox.y1 <= bbox.y1 + 3
                and line.bbox.x0 >= bbox.x0 - 10 and line.bbox.x1 <= bbox.x1 + 10):
            return True
    return False


def classify_lines(
    lines: list[TextLine],
    hierarchy: FontHierarchy,
    page_num: int,
    page_height: float,
    config: PipelineConfig,
    exclude_bboxes: list[BBox] | None = None,
) -> list[DocumentBlock]:
    """Classify text lines into document blocks."""
    blocks: list[DocumentBlock] = []
    current_body_lines: list[TextLine] = []
    _exclude = exclude_bboxes or []

    def flush_body():
        nonlocal current_body_lines
        if current_body_lines:
            text = _join_body_lines(current_body_lines)
            bbox = merge_bboxes([l.bbox for l in current_body_lines])
            blocks.append(DocumentBlock(
                block_type=BlockType.BODY,
                text=text,
                bbox=bbox,
                page_num=page_num,
                lines=list(current_body_lines),
            ))
            current_body_lines = []

    i = 0
    while i < len(lines):
        line = lines[i]
        text = line.text.strip()
        if not text:
            i += 1
            continue

        # Skip lines in table/figure regions
        if _exclude and _is_in_any_bbox(line, _exclude):
            i += 1
            continue

        # Skip page numbers (standalone digits near bottom of page)
        if text.isdigit() and line.bbox.y0 > page_height - 80:
            i += 1
            continue

        # Skip headers/footers near page edges — but not section headings
        if line.bbox.y0 < config.page_margin or line.bbox.y0 > page_height - config.page_margin:
            if len(text) < 80:
                # Don't skip bold/large text that looks like a section heading
                line_bold = any(s.is_bold for s in line.spans if s.text.strip())
                line_size = line.dominant_size
                is_heading_like = line_bold and line_size >= hierarchy.body_size
                if not is_heading_like:
                    i += 1
                    continue

        size = line.dominant_size
        font = line.dominant_font
        is_bold = any(s.is_bold for s in line.spans if len(s.text.strip()) > 0)

        # For bold-at-body-size detection, require majority bold (not just any)
        # to avoid classifying "**Datasets.** We use..." as a heading.
        if hierarchy.bold_heading_at_body_size and abs(size - hierarchy.body_size) < 1.0:
            bold_chars = sum(len(s.text) for s in line.spans if s.is_bold)
            total_chars = sum(len(s.text) for s in line.spans)
            if total_chars > 0 and bold_chars / total_chars < 0.7:
                is_bold = False

        level = hierarchy.heading_level(size, is_bold, font)

        # Footnote detection: text in the bottom 20% of the page that is
        # noticeably smaller than body text AND starts with a footnote marker.
        # This catches footnotes like "∗Equal contribution..." at 9pt when
        # body is 10pt, without false-positives on reference text.
        if (size < hierarchy.body_size - 0.5
                and line.bbox.y0 > page_height * 0.75
                and level == 0
                and not _is_allcaps_heading(text)
                and re.match(r"^[∗†‡§¶*]", text)):
            flush_body()
            blocks.append(DocumentBlock(
                block_type=BlockType.FOOTNOTE,
                text=text,
                bbox=line.bbox,
                page_num=page_num,
                lines=[line],
            ))
            i += 1
            continue

        if level > 0:
            # Check if this is really an author line misclassified as heading:
            # bold-at-body-size text with footnote markers (with commas, or short name-like text)
            if (level >= 2 and hierarchy.bold_heading_at_body_size
                    and abs(size - hierarchy.body_size) < 1.0
                    and any(ch in text for ch in '*∗†‡§¶★⋆')
                    and (',' in text or len(text.strip()) < 40)):
                flush_body()
                blocks.append(DocumentBlock(
                    block_type=BlockType.METADATA,
                    text=text,
                    bbox=line.bbox,
                    page_num=page_num,
                    lines=[line],
                ))
                i += 1
                continue
            # Reject numeric-only "headings" — table cell values like "70.4±.1"
            stripped = text.strip()
            alpha_chars = sum(1 for c in stripped if c.isalpha())
            if alpha_chars < 3 and level >= 2:
                current_body_lines.append(line)
                i += 1
                continue
            flush_body()
            bt = {1: BlockType.TITLE, 2: BlockType.HEADING, 3: BlockType.SUBHEADING}.get(level, BlockType.SUBHEADING)
            blocks.append(DocumentBlock(
                block_type=bt,
                text=text,
                bbox=line.bbox,
                page_num=page_num,
                lines=[line],
                level=level,
            ))
        elif is_bold and text.strip().lower() in _KNOWN_BOLD_HEADINGS:
            # Bold standalone heading word (e.g. "Abstract" at 9pt in Mamba)
            # even when smaller than body size
            flush_body()
            blocks.append(DocumentBlock(
                block_type=BlockType.HEADING,
                text=text,
                bbox=line.bbox,
                page_num=page_num,
                lines=[line],
                level=2,
            ))
        elif hierarchy.is_footnote(size) and not _is_allcaps_heading(text):
            flush_body()
            blocks.append(DocumentBlock(
                block_type=BlockType.FOOTNOTE,
                text=text,
                bbox=line.bbox,
                page_num=page_num,
                lines=[line],
            ))
        elif _is_allcaps_heading(text):
            # Check if this line has multiple numeric siblings at the same y (table row).
            # A single short number is likely a section number, not table data.
            numeric_neighbors = 0
            for other in lines:
                if other is line:
                    continue
                if abs(other.bbox.y0 - line.bbox.y0) < 3:
                    ot = other.text.strip()
                    # Count numeric cells (ignore single-digit section numbers)
                    if ot and len(ot) > 1 and (re.match(r'^[\d.,%+\-]+$', ot) or re.match(r'^\d{4}', ot)):
                        numeric_neighbors += 1
            if numeric_neighbors >= 2:
                current_body_lines.append(line)
            else:
                flush_body()
                # Determine level: longer all-caps = H2, shorter = H3
                ac_level = 2 if size >= hierarchy.body_size * 0.9 else 3
                bt = BlockType.HEADING if ac_level == 2 else BlockType.SUBHEADING
                blocks.append(DocumentBlock(
                    block_type=bt,
                    text=_smart_title_case(text),
                    bbox=line.bbox,
                    page_num=page_num,
                    lines=[line],
                    level=ac_level,
                ))
        elif _is_numbered_subsection(text, is_bold):
            flush_body()
            blocks.append(DocumentBlock(
                block_type=BlockType.SUBHEADING,
                text=text,
                bbox=line.bbox,
                page_num=page_num,
                lines=[line],
                level=3,
            ))
        elif is_bold and _BARE_SUBSECTION_RE.match(text) and _valid_subsection_number(text):
            # Bare subsection number (e.g. "3.1") — look ahead for title on same row
            sub_lines = [line]
            sub_text = text
            found_title = False
            if i + 1 < len(lines):
                next_l = lines[i + 1]
                next_bold = any(s.is_bold for s in next_l.spans if s.text.strip())
                next_text = next_l.text.strip()
                # Same row: similar y-position, and next line has alphabetic content
                if (next_bold and abs(next_l.bbox.y0 - line.bbox.y0) < 3
                        and any(c.isalpha() for c in next_text)):
                    sub_text = text + " " + next_text
                    sub_lines.append(next_l)
                    i += 1
                    found_title = True
            # Only classify as heading if we found an alphabetic title
            if found_title or any(c.isalpha() for c in text.split(".", 1)[-1]):
                flush_body()
                blocks.append(DocumentBlock(
                    block_type=BlockType.SUBHEADING,
                    text=sub_text,
                    bbox=merge_bboxes([l.bbox for l in sub_lines]),
                    page_num=page_num,
                    lines=sub_lines,
                    level=3,
                ))
            else:
                # Just a number in a table — treat as body text
                current_body_lines.append(line)
        elif _is_list_item(text):
            flush_body()
            # Collect continuation lines for this list item
            list_lines = [line]
            j = i + 1
            while j < len(lines):
                next_line = lines[j]
                next_text = next_line.text.strip()
                if not next_text:
                    j += 1
                    continue
                # Stop if: new list item, heading, too far vertically, or back at left margin
                if _is_list_item(next_text):
                    break
                next_size = next_line.dominant_size
                next_font = next_line.dominant_font
                next_bold = any(s.is_bold for s in next_line.spans if s.text.strip())
                if hierarchy.heading_level(next_size, next_bold, next_font) > 0:
                    break
                # Check if continuation: indented or close vertically
                gap = next_line.bbox.y0 - list_lines[-1].bbox.y1
                if gap > 20:  # too far
                    break
                # Continuation line should be at same or greater indentation
                if next_line.bbox.x0 >= line.bbox.x0 - 5:
                    list_lines.append(next_line)
                    j += 1
                else:
                    break

            full_text = _join_body_lines(list_lines)
            bbox = merge_bboxes([l.bbox for l in list_lines])
            blocks.append(DocumentBlock(
                block_type=BlockType.LIST_ITEM,
                text=full_text,
                bbox=bbox,
                page_num=page_num,
                lines=list_lines,
            ))
            i = j
            continue
        else:
            current_body_lines.append(line)

        i += 1

    flush_body()
    return blocks


def _is_list_item(text: str) -> bool:
    """Detect list items by bullet markers or numbered patterns."""
    text = text.strip()
    if not text:
        return False
    if text[0] in "\u2022\u2023\u25e6\u2043\u2219\u25cf\u25cb\u2013\u2014":
        return True
    if re.match(r"^\d{1,3}[.)]\s", text):
        return True
    if re.match(r"^\([a-z0-9]{1,3}\)\s", text):
        return True
    return False


_ALLCAPS_HEADING_RE = re.compile(
    r"^(\d+[\s.]*|[IVX]+[\s.]+)?[A-Z][A-Z\d\s&:,\-]{3,}$"
)

# Bold standalone words that should always be classified as headings,
# even when their font size is smaller than body text (e.g. Mamba paper).
_KNOWN_BOLD_HEADINGS = {
    "abstract", "introduction", "conclusion", "conclusions",
    "references", "acknowledgements", "acknowledgments",
    "bibliography", "appendix",
}

_ALLCAPS_EXCLUDE = {
    "GPU", "CPU", "TPU", "RAM", "CNN", "RNN", "LSTM", "GAN", "MLP",
    "BLEU", "ROUGE", "BERT", "GPT", "VIT", "CLIP",
}


def _is_allcaps_heading(text: str) -> bool:
    """Detect all-caps section headings (e.g. ICLR/ICML format).

    Examples: "ABSTRACT", "1 INTRODUCTION", "RELATED WORK", "METHOD"
    """
    t = text.strip()
    if len(t) < 5 or len(t) > 60:
        return False
    letters = [c for c in t if c.isalpha()]
    if len(letters) < 5:
        return False
    upper_ratio = sum(1 for c in letters if c.isupper()) / len(letters)
    if upper_ratio < 0.85:
        return False
    # Skip common acronyms and short dataset/model names
    if t in _ALLCAPS_EXCLUDE:
        return False
    # Headings don't end with commas or semicolons (those are mid-sentence fragments)
    if t.endswith(",") or t.endswith(";"):
        return False
    # Must have at least 2 words OR be a well-known heading
    words = t.split()
    alpha_words = [w for w in words if w[0].isalpha()]
    _KNOWN_HEADINGS = {"ABSTRACT", "INTRODUCTION", "CONCLUSION", "ACKNOWLEDGEMENTS",
                       "ACKNOWLEDGMENTS", "REFERENCES", "BIBLIOGRAPHY", "APPENDIX"}
    if len(alpha_words) == 1 and alpha_words[0] not in _KNOWN_HEADINGS:
        return False
    # Must match heading pattern (optional number prefix + caps words)
    return bool(_ALLCAPS_HEADING_RE.match(t))



_NUMBERED_SUBSECTION_RE = re.compile(
    r"^([A-Z]|\d+)\.\d+(\.\d+)*\.?\s+\S"
)
# Matches bare numbered subsection (number only, e.g. "3.1", "A.2")
_BARE_SUBSECTION_RE = re.compile(
    r"^([A-Z]|\d+)\.\d+(\.\d+)*\.?\s*$"
)


def _valid_subsection_number(text: str) -> bool:
    """Check if subsection number parts are reasonable (≤ 20 each)."""
    num_part = text.strip().split()[0].rstrip(".")
    parts = num_part.split(".")
    for p in parts[1:]:
        try:
            if int(p) > 20:
                return False
        except ValueError:
            pass
    return True


def _is_numbered_subsection(text: str, is_bold: bool) -> bool:
    """Detect bold numbered subsection headings at body text size."""
    if not is_bold:
        return False
    text = text.strip()
    if len(text) >= 100:
        return False
    if not _NUMBERED_SUBSECTION_RE.match(text):
        return False
    # Validate: sub-numbers should be reasonable (≤ 20) — reject decimal
    # values like "3.87" or "4.53" that appear in tables/metrics
    num_part = text.split()[0].rstrip(".")
    parts = num_part.split(".")
    for p in parts[1:]:
        try:
            if int(p) > 20:
                return False
        except ValueError:
            pass
    return True


def detect_page_columns(lines: list[TextLine], page_width: float, config: PipelineConfig) -> list[tuple[float, float]]:
    """Detect N-column layout from line positions.

    Clusters line midpoints to find 1, 2, or 3+ columns.
    Requires ≥3 lines per cluster to avoid false positives.
    Falls back to single column if clustering is ambiguous.
    """
    if not lines:
        return [(0, page_width)]

    # Only use body-width lines (not full-width headings/figures)
    # Adaptive threshold: for N columns, body line width < 1/N + margin
    # Start with a generous threshold (55%) and narrow if needed
    body_lines = [l for l in lines if l.bbox.width < page_width * 0.55]
    if len(body_lines) < 6:
        return [(0, page_width)]

    x_mids = [l.bbox.midx for l in body_lines]
    gap_threshold = config.column_gap_threshold

    # Cluster midpoints using 1D clustering
    clusters = cluster_1d(x_mids, gap_threshold)

    # Filter: require ≥3 lines per cluster
    valid_clusters = [c for c in clusters if len(c) >= 3]

    if len(valid_clusters) <= 1:
        return [(0, page_width)]

    # Absorb small clusters into larger ones when the small cluster's
    # x-range is mostly contained within a larger cluster.  This handles
    # spurious sub-clusters from heading fragments (e.g. a "3.1" label
    # creating a small cluster overlapping the main column) without
    # merging two genuine columns that happen to share a few bridging
    # table/figure lines.
    if len(valid_clusters) >= 2:
        # Compute x-ranges (min x0, max x1) for each cluster
        cluster_ranges = []
        for c in valid_clusters:
            c_lines = [body_lines[i] for i in c]
            cluster_ranges.append((min(l.bbox.x0 for l in c_lines),
                                   max(l.bbox.x1 for l in c_lines)))

        # Iteratively absorb small clusters into overlapping large ones
        changed = True
        while changed:
            changed = False
            # Sort by size descending so large clusters absorb small ones
            order = sorted(range(len(valid_clusters)),
                           key=lambda i: len(valid_clusters[i]), reverse=True)
            absorbed: set[int] = set()
            for i in order:
                if i in absorbed:
                    continue
                for j in order:
                    if j in absorbed or j == i:
                        continue
                    # Check if cluster j is much smaller than i
                    if len(valid_clusters[j]) >= len(valid_clusters[i]):
                        continue
                    # Check if j's x-range is mostly within i's x-range
                    i_x0, i_x1 = cluster_ranges[i]
                    j_x0, j_x1 = cluster_ranges[j]
                    overlap_x0 = max(i_x0, j_x0)
                    overlap_x1 = min(i_x1, j_x1)
                    j_width = j_x1 - j_x0
                    if j_width <= 0:
                        continue
                    overlap_width = max(0, overlap_x1 - overlap_x0)
                    # Absorb if >50% of j's range overlaps with i
                    if overlap_width > j_width * 0.5:
                        valid_clusters[i] = valid_clusters[i] + valid_clusters[j]
                        # Update x-range for merged cluster
                        c_lines = [body_lines[k] for k in valid_clusters[i]]
                        cluster_ranges[i] = (min(l.bbox.x0 for l in c_lines),
                                             max(l.bbox.x1 for l in c_lines))
                        absorbed.add(j)
                        changed = True
            if absorbed:
                valid_clusters = [valid_clusters[i] for i in range(len(valid_clusters))
                                  if i not in absorbed]
                cluster_ranges = [cluster_ranges[i] for i in range(len(cluster_ranges))
                                  if i not in absorbed]
        valid_clusters = [c for c in valid_clusters if len(c) >= 3]

    if len(valid_clusters) <= 1:
        return [(0, page_width)]

    # Sort clusters by centroid position (left to right)
    def _sort_by_centroid(clusters):
        centroids = [sum(x_mids[i] for i in c) / len(c) for c in clusters]
        order = sorted(range(len(clusters)), key=lambda i: centroids[i])
        return [clusters[i] for i in order]

    valid_clusters = _sort_by_centroid(valid_clusters)

    # Compute column boundaries from gaps between clusters.
    # If the gap check fails (clusters overlap), retry with only the
    # two largest clusters — the rest are likely noise from table/figure
    # fragments that bridge both columns.
    def _compute_columns(clusters):
        columns: list[tuple[float, float]] = []
        for ci in range(len(clusters)):
            cluster_lines = [body_lines[i] for i in clusters[ci]]
            col_min_x0 = min(l.bbox.x0 for l in cluster_lines)
            col_max_x1 = max(l.bbox.x1 for l in cluster_lines)

            if ci == 0:
                left_bound = 0.0
            else:
                prev_lines = [body_lines[i] for i in clusters[ci - 1]]
                prev_max_x1 = max(l.bbox.x1 for l in prev_lines)
                if col_min_x0 - prev_max_x1 <= 5:
                    return None  # overlap — signal failure
                left_bound = (prev_max_x1 + col_min_x0) / 2

            if ci == len(clusters) - 1:
                right_bound = page_width
            else:
                next_lines = [body_lines[i] for i in clusters[ci + 1]]
                next_min_x0 = min(l.bbox.x0 for l in next_lines)
                right_bound = (col_max_x1 + next_min_x0) / 2

            columns.append((left_bound, right_bound))
        return columns

    columns = _compute_columns(valid_clusters)
    if columns is not None:
        return columns

    # Gap check failed — retry with only the 2 largest clusters
    if len(valid_clusters) > 2:
        top2 = sorted(valid_clusters, key=len, reverse=True)[:2]
        top2 = _sort_by_centroid(top2)
        columns = _compute_columns(top2)
        if columns is not None:
            return columns

    return [(0, page_width)]


def split_lines_by_column(
    lines: list[TextLine],
    columns: list[tuple[float, float]],
    page_width: float,
) -> tuple[list[list[TextLine]], list[TextLine]]:
    """Split lines into per-column groups plus full-width lines.

    Returns: (column_lines_list, full_width_lines)
    """
    if len(columns) <= 1:
        return [lines], []

    col_lines: list[list[TextLine]] = [[] for _ in columns]
    full_width: list[TextLine] = []

    # Adaptive full-width threshold: for N columns, line is full-width if
    # it spans more than 1/N + 0.1 of page width
    n_cols = len(columns)
    fw_threshold = 1.0 / n_cols + 0.1 if n_cols > 1 else 1.0

    for line in lines:
        # Full-width: spans more than the adaptive threshold
        if line.bbox.width > page_width * fw_threshold:
            full_width.append(line)
            continue

        # Assign to column by midpoint: prefer strict containment, then
        # closest center.  The old tolerance-based first-match approach
        # mis-assigned narrow heading fragments near column boundaries.
        mid = line.bbox.midx
        # 1. Strict containment (no tolerance)
        strict = [ci for ci, (c0, c1) in enumerate(columns) if c0 <= mid <= c1]
        if len(strict) == 1:
            col_lines[strict[0]].append(line)
        elif strict:
            # On boundary — pick closest center
            best = min(strict,
                       key=lambda ci: abs(mid - (columns[ci][0] + columns[ci][1]) / 2))
            col_lines[best].append(line)
        else:
            # Outside all columns — closest center
            closest = min(range(len(columns)),
                         key=lambda ci: abs(mid - (columns[ci][0] + columns[ci][1]) / 2))
            col_lines[closest].append(line)

    return col_lines, full_width


def classify_page_with_columns(
    lines: list[TextLine],
    hierarchy: FontHierarchy,
    page_num: int,
    page_height: float,
    page_width: float,
    config: PipelineConfig,
    exclude_bboxes: list[BBox] | None = None,
) -> list[DocumentBlock]:
    """Classify lines with proper column handling."""
    columns = detect_page_columns(lines, page_width, config)

    if len(columns) <= 1:
        # Single column - simple path
        blocks = classify_lines(lines, hierarchy, page_num, page_height, config, exclude_bboxes)
        return sorted(blocks, key=lambda b: b.bbox.y0)

    # Multi-column: split lines, classify each column, then merge
    col_lines_list, full_width_lines = split_lines_by_column(lines, columns, page_width)

    all_blocks: list[DocumentBlock] = []

    for col_idx, col_lines in enumerate(col_lines_list):
        if not col_lines:
            continue
        blocks = classify_lines(col_lines, hierarchy, page_num, page_height, config, exclude_bboxes)
        # Tag column index for reading order
        for b in blocks:
            b.metadata["column"] = col_idx
        all_blocks.extend(blocks)

    # Full-width lines (spanning figures, equations, etc.)
    if full_width_lines:
        fw_blocks = classify_lines(full_width_lines, hierarchy, page_num, page_height, config, exclude_bboxes)
        for b in fw_blocks:
            b.metadata["column"] = -1  # full width
        all_blocks.extend(fw_blocks)

    # Sort: for each y-position band, left column before right, with full-width interleaved
    return _sort_column_blocks(all_blocks, columns)


def _sort_column_blocks(blocks: list[DocumentBlock], columns: list[tuple[float, float]]) -> list[DocumentBlock]:
    """Sort blocks in reading order for multi-column layout.

    Reading order: process left column top-to-bottom, then right column top-to-bottom.
    Full-width elements interleaved by y-position.
    """
    if not blocks:
        return blocks

    col_blocks: dict[int, list[DocumentBlock]] = {}
    for b in blocks:
        col = b.metadata.get("column", 0)
        col_blocks.setdefault(col, []).append(b)

    # Sort each column by y-position
    for col in col_blocks:
        col_blocks[col].sort(key=lambda b: b.bbox.y0)

    # Interleave: left col, right col, with full-width in between
    result = []
    fw = col_blocks.pop(-1, [])

    # Process columns left to right
    for col_idx in sorted(col_blocks.keys()):
        col_b = col_blocks[col_idx]
        # Insert any full-width blocks that come before this column's first block
        while fw and col_b and fw[0].bbox.y0 < col_b[0].bbox.y0:
            result.append(fw.pop(0))
        result.extend(col_b)

    # Remaining full-width blocks
    result.extend(fw)
    return result


def sort_blocks_reading_order(blocks: list[DocumentBlock], columns: list[tuple[float, float]]) -> list[DocumentBlock]:
    """Sort blocks in reading order: top-to-bottom within each column, left-to-right across columns."""
    if len(columns) <= 1:
        return sorted(blocks, key=lambda b: (b.page_num, b.bbox.y0))

    def column_index(block: DocumentBlock) -> int:
        mid = block.bbox.midx
        for i, (c0, c1) in enumerate(columns):
            if c0 <= mid <= c1:
                return i
        return 0

    return sorted(blocks, key=lambda b: (b.page_num, column_index(b), b.bbox.y0))
