"""Equation extraction: CM font detection, math run grouping, LaTeX reconstruction."""
from __future__ import annotations

import re as _re

from ..types import Span, TextLine, BBox
from ..utils.fonts import is_math_font, char_to_latex, CM_TO_LATEX


# Characters that are acceptable as "glue" between adjacent inline math
# fragments — they do not break a mathematical expression.
_INLINE_MATH_GLUE = _re.compile(r"^[\s\+\-\*/=,;\.\!\?\(\)\[\]\{\}\\|^_~&%<>]+$")

def is_math_span(span: Span, body_font: str = "") -> bool:
    """Check if a span is math content based on font or character content."""
    if is_math_font(span.font, body_font):
        return True
    text = span.text
    # Single special math characters
    for ch in text:
        if ch in CM_TO_LATEX:
            return True
    return False


def detect_superscript(span: Span, body_size: float) -> bool:
    """Detect if a span is a superscript based on size and vertical position."""
    return span.size < body_size * 0.8


def detect_subscript(span: Span, prev_span: Span | None, body_size: float) -> bool:
    """Detect if a span is a subscript."""
    if span.size >= body_size * 0.8:
        return False
    if prev_span and span.bbox.y1 > prev_span.bbox.y1 + 1:
        return True
    return False


def _equation_complexity(spans: list[Span], body_size: float) -> int:
    """Score the structural complexity of a math expression.

    Higher score = more complex. Scores ≥3 indicate fractions, matrices,
    or deeply nested sub/superscripts that may benefit from image fallback.
    """
    if not spans:
        return 0

    score = 0
    size_levels = set()
    for s in spans:
        rounded = round(s.size * 2) / 2
        size_levels.add(rounded)

    # Multiple size levels → nested sub/superscripts or fractions
    if len(size_levels) >= 3:
        score += 2
    elif len(size_levels) >= 2:
        score += 1

    # Vertical stacking (fraction-like): spans at same x but different y
    x_groups: dict[int, list[float]] = {}
    for s in spans:
        x_bucket = int(s.bbox.midx / 5)
        x_groups.setdefault(x_bucket, []).append(s.origin[1])
    for ys in x_groups.values():
        if len(ys) >= 2:
            y_range = max(ys) - min(ys)
            if y_range > body_size * 0.8:
                score += 2  # likely a fraction or large structure
                break

    # Sum/product/integral with limits
    text = "".join(s.text for s in spans)
    if any(ch in text for ch in "\u2211\u220f\u222b"):
        score += 1

    # Square root
    if "\u221a" in text:
        score += 1

    # Many spans → complex layout
    if len(spans) > 15:
        score += 1

    # Unbalanced delimiters → likely incomplete/complex
    opens = sum(1 for ch in text if ch in '([{')
    closes = sum(1 for ch in text if ch in ')]}')
    if opens != closes:
        score += 1

    return score


def _detect_fraction_pattern(spans: list[Span], body_size: float) -> list[tuple[int, int, int]]:
    """Detect fraction patterns: numerator over denominator at same x.

    Returns list of (numerator_idx, bar_idx, denominator_idx) or empty list.
    Fractions appear as vertically stacked spans at similar x positions
    with a size smaller than body text.
    """
    # Simple heuristic: look for groups of spans at the same horizontal
    # position with different vertical positions
    fractions: list[tuple[int, int, int]] = []

    for i, s in enumerate(spans):
        text = s.text.strip()
        # Look for fraction bars (horizontal rule characters)
        if text in ("\u2015", "\u2014", "\u2013", "\u2500", "\u2012"):
            # Check for spans above and below
            bar_x = s.bbox.midx
            above = None
            below = None
            for j, other in enumerate(spans):
                if j == i:
                    continue
                if abs(other.bbox.midx - bar_x) < 20:
                    if other.origin[1] < s.origin[1] - 2:
                        above = j
                    elif other.origin[1] > s.origin[1] + 2:
                        below = j
            if above is not None and below is not None:
                fractions.append((above, i, below))

    return fractions


def process_math_spans(spans: list[Span], body_size: float) -> str:
    """Convert a sequence of math-font spans into LaTeX.

    Handles:
    - Greek letters, operators, relations via CM_TO_LATEX
    - Superscripts/subscripts via font size and position
    - Italic variables (CMMI = math italic)
    - Fraction patterns (\frac{}{})
    - Sum/product/integral with limits
    """
    if not spans:
        return ""

    # Detect fraction patterns
    fractions = _detect_fraction_pattern(spans, body_size)
    frac_indices: set[int] = set()
    frac_replacements: dict[int, str] = {}
    for num_idx, bar_idx, den_idx in fractions:
        frac_indices.update({num_idx, bar_idx, den_idx})
        num_text = _convert_math_text(spans[num_idx].text, spans[num_idx].font)
        den_text = _convert_math_text(spans[den_idx].text, spans[den_idx].font)
        # Place the \frac at the bar position
        frac_replacements[bar_idx] = f"\\frac{{{num_text}}}{{{den_text}}}"

    parts: list[str] = []
    prev_span = None
    # Track the current sub/superscript run to merge consecutive ones
    # e.g. _{t}_{-}_{1} → _{t-1}
    pending_script: str | None = None  # "sub" or "super"
    pending_script_parts: list[str] = []

    def _flush_script():
        nonlocal pending_script, pending_script_parts
        if pending_script and pending_script_parts:
            content = "".join(pending_script_parts)
            prefix = "^" if pending_script == "super" else "_"
            parts.append(f"{prefix}{{{content}}}")
        pending_script = None
        pending_script_parts = []

    for i, span in enumerate(spans):
        # Skip spans consumed by fraction detection
        if i in frac_indices and i not in frac_replacements:
            prev_span = span
            continue
        if i in frac_replacements:
            _flush_script()
            parts.append(frac_replacements[i])
            prev_span = span
            continue

        text = span.text
        is_small = span.size < body_size * 0.8

        if is_small and prev_span:
            # Determine super vs subscript from vertical position
            if span.origin[1] < prev_span.origin[1] - 1:
                script_type = "super"
            else:
                script_type = "sub"
            converted = _convert_math_text(text, span.font)
            # Merge consecutive scripts of the same type
            if pending_script == script_type:
                pending_script_parts.append(converted)
            else:
                _flush_script()
                pending_script = script_type
                pending_script_parts = [converted]
        else:
            _flush_script()
            converted = _convert_math_text(text, span.font)
            # Enhance sum/product/integral with limits notation
            if "\u2211" in text:
                converted = converted.replace("\\sum ", "\\sum")
            elif "\u220f" in text:
                converted = converted.replace("\\prod ", "\\prod")
            elif "\u222b" in text:
                converted = converted.replace("\\int ", "\\int")
            parts.append(converted)

        prev_span = span

    _flush_script()
    result = "".join(parts)
    # Merge adjacent subscripts: _{a}_{b} → _{ab} (fragmented subscript arguments)
    prev = None
    while prev != result:
        prev = result
        result = _re.sub(r'_\{([^}]*)\}_\{([^}]*)\}', r'_{\1\2}', result)
    # Merge adjacent superscripts: ^{a}^{b} → ^{ab}
    prev = None
    while prev != result:
        prev = result
        result = _re.sub(r'\^\{([^}]*)\}\^\{([^}]*)\}', r'^{\1\2}', result)
    # Merge superscript minus + subscript digit: ^{-}_{9} → ^{-9}
    result = _re.sub(r'\^\{-\}_\{(\d+[^}]*)\}', r'^{-\1}', result)
    # Strip spaces before sub/superscript markers: \beta _{1} → \beta_{1}
    result = _re.sub(r'\s+([_^])', r'\1', result)
    # Fix \sqrt not braced: \sqrtx → \sqrt{x}, \sqrtd_{k} → \sqrt{d_{k}}
    result = _re.sub(
        r'\\sqrt([a-zA-Z0-9](?:_\{[^}]*\}|\^?\{[^}]*\}|[a-zA-Z0-9])*)',
        r'\\sqrt{\1}',
        result,
    )
    # Fix \sqrt_{x} → \sqrt{x} (subscript after radical should be argument)
    result = _re.sub(r'\\sqrt_\{([^}]*)\}', r'\\sqrt{\1}', result)
    # Fix parenthesized superscripts split across sub/super:
    # ^{(}_{k)} → ^{(k)}, ^{(}_{k}  → ^{(k)}, _{(}^{k)} → _{(k)}
    result = _re.sub(r'\^\{\(\}_\{([^}]*)\)}', r'^{(\1)}', result)
    result = _re.sub(r'\^\{\(}_\{([^}]*)\}', r'^{(\1)}', result)
    result = _re.sub(r'_\{\(\}\^\{([^}]*)\)}', r'_{(\1)}', result)
    result = _re.sub(r'_\{\(}\^\{([^}]*)\}', r'_{(\1)}', result)
    # Also handle ^{(}k) or ^{(}_{k)} with literal close paren
    result = _re.sub(r'\^\{\(\}([a-zA-Z0-9])\)', r'^{(\1)}', result)
    return result


def _convert_math_text(text: str, font: str) -> str:
    """Convert text to LaTeX based on font type."""
    result = []
    for ch in text:
        latex = char_to_latex(ch, font)
        result.append(latex)
    return "".join(result)


def extract_inline_math(line: TextLine, body_size: float, body_font: str = "") -> str:
    """Process a text line, wrapping math-font runs in $...$.

    Adjacent math runs that are separated only by whitespace or
    operator-like glue characters are merged into a single $...$ block
    rather than producing many tiny fragments like $x$$F$$ $$-!$$T$.
    """
    if not line.spans:
        return ""

    # Build a flat list of segments: (is_math: bool, spans: list[Span])
    segments: list[tuple[bool, list]] = []
    for span in line.spans:
        math = is_math_span(span, body_font)
        if segments and segments[-1][0] == math:
            segments[-1][1].append(span)
        else:
            segments.append((math, [span]))

    # Merge non-math "glue" segments that sit between two math runs, and
    # pull the right math run into the same group so it all becomes one
    # $...$ block.  A glue segment contains only whitespace or simple
    # punctuation/operators — no alphabetic content.
    merged: list[tuple[bool, list]] = []
    i = 0
    while i < len(segments):
        is_math, content = segments[i]
        if (not is_math
                and merged and merged[-1][0]           # left neighbour is math
                and i + 1 < len(segments)              # right neighbour exists
                and segments[i + 1][0]):               # right neighbour is math
            glue_text = "".join(s.text for s in content)
            if _INLINE_MATH_GLUE.match(glue_text):
                # Absorb both the glue spans AND the next math run into
                # the current (left) math group, creating one merged run.
                merged[-1][1].extend(content)          # glue
                merged[-1][1].extend(segments[i + 1][1])  # right math
                i += 2  # skip both glue and right math segments
                continue
        merged.append((is_math, list(content)))
        i += 1

    # Pass 2: Absorb subscript-sized non-math text into adjacent math runs.
    # E.g. "model" at 7pt after math "d" at 10pt → becomes \text{model} subscript.
    absorbed: list[tuple[bool, list]] = []
    for is_math, content in merged:
        if (not is_math
                and absorbed and absorbed[-1][0]  # preceded by math
                and all(s.size < body_size * 0.8 for s in content)):
            text = "".join(s.text for s in content).strip()
            if text and len(text) < 25 and not any(ch in text for ch in '.;:!?'):
                # Absorb into preceding math group
                absorbed[-1][1].extend(content)
                continue
        absorbed.append((is_math, list(content)))
    merged = absorbed

    # Build final output, processing spans in document order.
    parts: list[str] = []
    for is_math, content in merged:
        if is_math:
            # Process spans in order: math spans become LaTeX, glue spans
            # are inserted verbatim so operator positions are preserved.
            tokens: list[str] = []
            math_buf: list[Span] = []
            for span in content:
                if is_math_span(span, body_font):
                    math_buf.append(span)
                else:
                    if math_buf:
                        tokens.append(process_math_spans(math_buf, body_size))
                        math_buf = []
                    # Non-math text at subscript size inside a math run
                    # → wrap in \text{} for proper rendering
                    text = span.text.strip()
                    if text and span.size < body_size * 0.8:
                        tokens.append(f"\\text{{{text}}}")
                    else:
                        tokens.append(span.text)
            if math_buf:
                tokens.append(process_math_spans(math_buf, body_size))

            full = "".join(tokens).strip()
            if full and not full.isspace():
                parts.append(f"${full}$")
        else:
            parts.append("".join(s.text for s in content))

    return "".join(parts)


def is_display_equation(line: TextLine, page_width: float, body_size: float,
                        body_font: str = "") -> bool:
    """Check if a line is a display (block) equation.

    Criteria: mostly math fonts, centered on page, possibly with equation number.
    """
    if not line.spans:
        return False

    math_chars = sum(len(s.text) for s in line.spans if is_math_span(s, body_font))
    total_chars = sum(len(s.text) for s in line.spans)
    if total_chars == 0:
        return False

    math_ratio = math_chars / total_chars

    # Check centering
    line_center = line.bbox.midx
    page_center = page_width / 2
    is_centered = abs(line_center - page_center) < page_width * 0.15

    # High math content + centered = display equation
    if math_ratio > 0.6 and is_centered:
        return True
    # All math content even if not perfectly centered
    if math_ratio > 0.9:
        return True

    return False


def _is_trivial_math(latex: str) -> bool:
    """Return True if the latex string is too trivial to be a display equation.

    Filters cases like: single numbers (200), single variables (x, T),
    bare operator fragments (-!, =), or any content shorter than 5 chars.
    A display equation must contain meaningful mathematical structure.
    """
    s = latex.strip()
    if not s:
        return True
    # Strip common LaTeX command prefixes/suffixes to measure raw content
    # length after removing whitespace-only padding.
    content_chars = len(s.replace(" ", ""))
    if content_chars < 5:
        return True
    # Single token: just a number, single letter, or simple operator
    # e.g. "200", "x", "T", "-!", "= "
    if _re.match(r"^[0-9]+$", s):          # bare number
        return True
    if _re.match(r"^[A-Za-z]$", s):        # single letter variable
        return True
    # Only operators/punctuation/whitespace — no variables or numbers
    # e.g. "-! · · · -!", "= ", "/ "
    stripped = _re.sub(r'\\[a-zA-Z]+\s*', '', s)  # remove LaTeX commands
    if not any(c.isalnum() for c in stripped):
        return True
    # Looks like figure axis data: mostly digits and commas
    digits = sum(1 for c in s if c.isdigit())
    if digits > len(s.replace(" ", "")) * 0.5 and digits > 5:
        return True
    # Single LaTeX command with no arguments (e.g., \sqrt, \sum)
    if _re.match(r'^\\[a-zA-Z]+$', s):
        return True
    return False


def extract_display_equation(line: TextLine, body_size: float,
                             body_font: str = "") -> str:
    """Convert a display equation line to a single $$...$$ LaTeX block.

    Returns an empty string if the equation is trivially short (single
    numbers, single variables, operator fragments) — these are better
    rendered as inline math or plain text.
    """
    latex = process_math_spans(
        [s for s in line.spans if is_math_span(s, body_font)],
        body_size,
    )
    # Check for equation number (rightmost non-math span)
    eq_num = ""
    for span in reversed(line.spans):
        text = span.text.strip()
        if text and not is_math_span(span, body_font):
            if _re.match(r"^\([\d.]+\)$", text):
                eq_num = f"  {text}"
            break

    latex = latex.strip()
    if latex and not _is_trivial_math(latex):
        return f"$${latex}$${eq_num}"
    return ""
