"""
Styling post-processor for OMML to MathML conversion.

This module extracts styling information (bold, fonts, colors) from OMML XML
and applies it to MathML output as mathvariant and other attributes.

The Microsoft OMML2MML.xsl stylesheet does not preserve styling information,
so this post-processor fills that gap by:
1. Extracting m:rPr (run properties) from OMML
2. Building a styling map (position, text) -> StyleInfo
3. Enhancing MathML <mi> elements with mathvariant and other attributes
"""

from dataclasses import dataclass
from typing import Dict, Tuple, Optional, List
from lxml import etree as ET


# Namespace constants
OMML_NS = 'http://schemas.openxmlformats.org/officeDocument/2006/math'
WORD_NS = 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
MATHML_NS = 'http://www.w3.org/1998/Math/MathML'

# OMML namespaces for parsing
OMML_NAMESPACES = {
    'm': OMML_NS,
    'w': WORD_NS,
}


@dataclass
class StyleInfo:
    """Styling information extracted from OMML run properties."""

    # Text content for matching
    text: str

    # Position in equation (for disambiguation)
    position: int

    # Style attributes from m:rPr
    mathvariant: Optional[str] = None  # bold, italic, bold-italic, normal, double-struck

    # Font family from w:rPr
    font_family: Optional[str] = None

    # Text color from w:rPr
    color: Optional[str] = None

    # Background/highlight color from w:rPr (w:highlight element)
    background_color: Optional[str] = None

    # Font size from w:rPr (in points)
    font_size: Optional[float] = None

    # Combined flags for easier processing
    is_bold: bool = False
    is_italic: bool = False
    is_plain: bool = False  # m:sty val="p" - upright/roman
    is_text_mode: bool = False  # m:nor - Word "Text Mode" toggle (formatting deferred to w:rPr)


@dataclass
class StructuralStyleInfo:
    """Styling information for structural elements (fractions, radicals, etc.)."""

    # Element type from OMML: 'f' (fraction), 'rad' (radical), 'func' (function), 'd' (delimiter)
    element_type: str

    # Position in equation (for matching with MathML elements)
    position: int

    # Color from structural properties (m:fPr/w:rPr/w:color or m:radPr/w:rPr/w:color)
    color: Optional[str] = None


class StylingPostProcessor:
    """
    Post-processes MathML to add styling attributes from OMML.

    This bridges the gap between Microsoft's OMML2MML.xsl (which doesn't preserve
    styling) and our LaTeX converter (which does handle styling).

    Usage:
        processor = StylingPostProcessor()
        styling_map = processor.extract_styling_from_omml(omml_xml)
        enhanced_mathml = processor.apply_styling_to_mathml(mathml_xml, styling_map)
    """

    def __init__(self):
        """Initialize the post-processor."""
        pass

    def extract_styling_from_omml(self, omml_xml: str) -> Dict[Tuple[int, str], StyleInfo]:
        """
        Extract styling information from OMML XML.

        Parses OMML to find all <m:r> (run) elements with <m:rPr> (run properties)
        and builds a map of (position, text) -> StyleInfo.

        Args:
            omml_xml: OMML XML string

        Returns:
            Dict mapping (position, text) to StyleInfo objects

        Examples:
            >>> processor = StylingPostProcessor()
            >>> omml = '<m:oMath><m:r><m:rPr><m:sty m:val="b"/></m:rPr><m:t>Z</m:t></m:r></m:oMath>'
            >>> styling = processor.extract_styling_from_omml(omml)
            >>> styling[(0, 'Z')].is_bold
            True
        """
        try:
            root = ET.fromstring(omml_xml.encode('utf-8'))
        except ET.XMLSyntaxError:
            # If parsing fails, return empty map (no styling)
            return {}

        styling_map = {}
        position = 0

        # Find all <m:r> (run) elements
        for r_elem in root.iter(f'{{{OMML_NS}}}r'):
            # Get text content
            t_elem = r_elem.find(f'{{{OMML_NS}}}t', OMML_NAMESPACES)
            if t_elem is None or not t_elem.text:
                continue

            text = t_elem.text.strip()
            if not text:
                continue

            # Check for run properties
            m_rpr = r_elem.find(f'{{{OMML_NS}}}rPr', OMML_NAMESPACES)
            w_rpr = r_elem.find(f'{{{WORD_NS}}}rPr', OMML_NAMESPACES)

            # Skip if no styling information
            if m_rpr is None and w_rpr is None:
                position += 1
                continue

            # Extract styling information
            style_info = StyleInfo(text=text, position=position)

            # Process OMML math run properties (m:rPr)
            if m_rpr is not None:
                self._extract_math_properties(m_rpr, style_info)

            # Process Word run properties (w:rPr)
            if w_rpr is not None:
                self._extract_word_properties(w_rpr, style_info)

            # Post-check: if m:nor (text mode) is active but no explicit style was set
            # (neither m:sty nor w:i/w:b), default to normal (upright)
            if style_info.is_text_mode and style_info.mathvariant is None:
                style_info.mathvariant = 'normal'

            # Only store if there's meaningful styling (not just default font)
            # Skip if only font_family is set and it's Cambria Math (default)
            has_meaningful_styling = (
                style_info.mathvariant is not None or
                style_info.color is not None or
                style_info.background_color is not None or
                style_info.font_size is not None or
                (style_info.font_family and style_info.font_family != 'Cambria Math')
            )

            if has_meaningful_styling:
                styling_map[(position, text)] = style_info

                # CRITICAL FIX: When styled text contains mixed content (letters + operators/numbers),
                # XSLT will split it into separate elements. We need to track character-level styling
                # so each split element can still find its style.
                #
                # Example: "O=2MnO" with mathvariant="normal" gets split by XSLT into:
                #   <mi>O</mi><mo>=</mo><mn>2</mn><mi>M</mi><mi>n</mi><mi>O</mi>
                # We need entries for 'O', '=', '2', 'MnO' (letters merge, operators/numbers split)
                if len(text) > 1 and self._text_will_be_split(text):
                    self._add_character_level_styling(styling_map, text, style_info, position)

            position += 1

        return styling_map

    def _text_will_be_split(self, text: str) -> bool:
        """
        Check if text will be split by XSLT into multiple elements.

        XSLT splits text when it contains:
        - Operators: =, +, -, *, /, <, >, etc.
        - Transitions between letters and digits

        Args:
            text: Text to check

        Returns:
            True if XSLT will split this text
        """
        operators = set('=+-*/:<>|()[]{},;')

        has_letters = any(c.isalpha() for c in text)
        has_operators = any(c in operators for c in text)
        has_digits = any(c.isdigit() for c in text)

        # Split happens when we have letters + operators, or letters + digits + operators
        return has_letters and (has_operators or (has_digits and has_letters))

    def _add_character_level_styling(
        self,
        styling_map: Dict[Tuple[int, str], 'StyleInfo'],
        text: str,
        style_info: 'StyleInfo',
        base_position: int
    ) -> None:
        """
        Add character/token-level styling entries for text that will be split by XSLT.

        XSLT tokenizes text into:
        - Consecutive letters → single <mi> (but may be further split)
        - Consecutive digits → single <mn>
        - Operators → individual <mo>

        We need to create styling entries for each potential token.

        Args:
            styling_map: Map to add entries to
            text: The original styled text
            style_info: The style to apply
            base_position: Base position in the OMML
        """
        operators = set('=+-*/:<>|()[]{},;')

        # Tokenize like XSLT would
        tokens = []
        current_token = ''
        current_type = None  # 'letter', 'digit', 'op'

        for char in text:
            if char.isalpha():
                char_type = 'letter'
            elif char.isdigit():
                char_type = 'digit'
            elif char in operators:
                char_type = 'op'
            else:
                char_type = 'other'

            if current_type is None:
                current_type = char_type
                current_token = char
            elif char_type == current_type and char_type != 'op':
                # Same type and not operator - accumulate
                current_token += char
            else:
                # Type changed or operator - flush and start new
                if current_token:
                    tokens.append((current_token, current_type))
                current_token = char
                current_type = char_type

        # Flush final token
        if current_token:
            tokens.append((current_token, current_type))

        # Add styling entries for each token
        # Use a sub-position to differentiate tokens from the same OMML run
        for sub_pos, (token, token_type) in enumerate(tokens):
            if token_type == 'letter':
                # Create a copy of style_info with updated text
                token_style = StyleInfo(
                    text=token,
                    position=base_position,
                    mathvariant=style_info.mathvariant,
                    font_family=style_info.font_family,
                    color=style_info.color,
                    background_color=style_info.background_color,
                    font_size=style_info.font_size,
                    is_bold=style_info.is_bold,
                    is_italic=style_info.is_italic,
                    is_plain=style_info.is_plain,
                    is_text_mode=style_info.is_text_mode,
                )
                # Use token text as key for text-only lookup
                # Multiple entries for same position with different texts
                styling_map[(base_position, token)] = token_style

                # CRITICAL: XSLT may also split letter groups into individual letters
                # (e.g., 'MnO' → <mi>M</mi><mi>n</mi><mi>O</mi>)
                # Add entries for each individual letter as well
                if len(token) > 1:
                    for char in token:
                        char_style = StyleInfo(
                            text=char,
                            position=base_position,
                            mathvariant=style_info.mathvariant,
                            font_family=style_info.font_family,
                            color=style_info.color,
                            background_color=style_info.background_color,
                            font_size=style_info.font_size,
                            is_bold=style_info.is_bold,
                            is_italic=style_info.is_italic,
                            is_plain=style_info.is_plain,
                            is_text_mode=style_info.is_text_mode,
                        )
                        styling_map[(base_position, char)] = char_style
            # For operators and digits, we typically don't need mathvariant styling
            # as they're already upright by default in LaTeX

    def extract_structural_styling_from_omml(
        self,
        omml_xml: str
    ) -> Dict[Tuple[str, int], StructuralStyleInfo]:
        """
        Extract styling from OMML structural elements (fractions, radicals, etc.).

        This method identifies structural elements like <m:f> (fractions) and <m:rad> (radicals)
        that have color properties in their properties elements (m:fPr, m:radPr).

        Args:
            omml_xml: The OMML XML string

        Returns:
            Dict mapping (element_type, position) -> StructuralStyleInfo
            Example: ('f', 0) -> StructuralStyleInfo(element_type='f', position=0, color='FFFF00')
        """
        try:
            root = ET.fromstring(omml_xml.encode('utf-8'))
        except ET.XMLSyntaxError:
            return {}

        structural_map = {}

        # Track positions for each element type separately
        positions = {'f': 0, 'rad': 0, 'func': 0, 'd': 0}

        # Process fractions (m:f)
        for f_elem in root.iter(f'{{{OMML_NS}}}f'):
            fpr = f_elem.find(f'{{{OMML_NS}}}fPr', OMML_NAMESPACES)
            if fpr is not None:
                color = self._extract_color_from_props(fpr)
                if color:
                    info = StructuralStyleInfo(
                        element_type='f',
                        position=positions['f'],
                        color=color
                    )
                    structural_map[('f', positions['f'])] = info
            positions['f'] += 1

        # Process radicals (m:rad) - square roots and nth roots
        for rad_elem in root.iter(f'{{{OMML_NS}}}rad'):
            radpr = rad_elem.find(f'{{{OMML_NS}}}radPr', OMML_NAMESPACES)
            if radpr is not None:
                color = self._extract_color_from_props(radpr)
                if color:
                    info = StructuralStyleInfo(
                        element_type='rad',
                        position=positions['rad'],
                        color=color
                    )
                    structural_map[('rad', positions['rad'])] = info
            positions['rad'] += 1

        # Future: Add support for m:func (functions), m:d (delimiters), etc.

        return structural_map

    def _extract_color_from_props(self, props_elem: ET.Element) -> Optional[str]:
        """
        Extract w:color from structural property element (m:fPr, m:radPr, etc.).

        Args:
            props_elem: The properties element (m:fPr or m:radPr)

        Returns:
            Color value (hex string without #) or None if no color found

        Note:
            Word may nest the color in two ways:
            1. Direct: m:fPr/w:rPr/w:color
            2. Via control properties: m:fPr/m:ctrlPr/w:rPr/w:color
        """
        # Try direct path first: m:fPr/w:rPr/w:color
        w_rpr = props_elem.find(f'{{{WORD_NS}}}rPr', OMML_NAMESPACES)

        # If not found, try via m:ctrlPr (control properties)
        if w_rpr is None:
            ctrl_pr = props_elem.find(f'{{{OMML_NS}}}ctrlPr', OMML_NAMESPACES)
            if ctrl_pr is not None:
                w_rpr = ctrl_pr.find(f'{{{WORD_NS}}}rPr', OMML_NAMESPACES)

        # Extract color from w:rPr if found
        if w_rpr is not None:
            color_elem = w_rpr.find(f'{{{WORD_NS}}}color', OMML_NAMESPACES)
            if color_elem is not None:
                # Try with namespace prefix first, then without
                color_val = color_elem.get(f'{{{WORD_NS}}}val')
                if color_val is None:
                    color_val = color_elem.get('val')

                # Skip "auto" color (means default)
                if color_val and color_val != 'auto':
                    return color_val

        return None

    def _extract_math_properties(self, m_rpr: ET.Element, style_info: StyleInfo):
        """
        Extract styling from OMML math run properties (m:rPr).

        Handles:
        - m:sty (style): b=bold, p=plain, i=italic, bi=bold-italic
        - m:scr (script): double-struck, script, fraktur, sans-serif, monospace
        """
        # Check for m:nor - indicates "text mode" in Word equation editor
        # m:nor switches to Word text formatting: w:i/w:b take effect.
        # Without explicit m:sty, default is upright (normal) - but w:i can override to italic.
        nor_elem = m_rpr.find(f'{{{OMML_NS}}}nor', OMML_NAMESPACES)
        if nor_elem is not None:
            # <m:nor/> activates text mode - defer mathvariant to w:rPr
            # If w:i is present, text becomes italic; otherwise stays normal
            style_info.is_plain = True
            style_info.is_text_mode = True
            # Do NOT set mathvariant here - let _extract_word_properties handle w:i

        # Check for m:sty (style) element
        sty_elem = m_rpr.find(f'{{{OMML_NS}}}sty', OMML_NAMESPACES)
        if sty_elem is not None:
            sty_val = sty_elem.get(f'{{{OMML_NS}}}val')
            if sty_val is None:
                sty_val = sty_elem.get('val')  # Try without namespace

            if sty_val == 'b':
                style_info.is_bold = True
                style_info.mathvariant = 'bold'
            elif sty_val == 'p':
                # m:sty val="p" indicates plain/upright text (not italic)
                style_info.is_plain = True
                if style_info.text and any(c.isalpha() for c in style_info.text):
                    style_info.mathvariant = 'normal'
            elif sty_val == 'i':
                style_info.is_italic = True
                style_info.mathvariant = 'italic'
            elif sty_val == 'bi':
                style_info.is_bold = True
                style_info.is_italic = True
                style_info.mathvariant = 'bold-italic'

        # Check for m:scr (script/font variant) element
        scr_elem = m_rpr.find(f'{{{OMML_NS}}}scr', OMML_NAMESPACES)
        if scr_elem is not None:
            scr_val = scr_elem.get(f'{{{OMML_NS}}}val')
            if scr_val is None:
                scr_val = scr_elem.get('val')

            # m:scr overrides m:sty for mathvariant
            if scr_val == 'double-struck':
                style_info.mathvariant = 'double-struck'
            elif scr_val == 'script':
                style_info.mathvariant = 'script'
            elif scr_val == 'fraktur':
                style_info.mathvariant = 'fraktur'
            elif scr_val == 'sans-serif':
                style_info.mathvariant = 'sans-serif'
            elif scr_val == 'monospace':
                style_info.mathvariant = 'monospace'

    def _extract_word_properties(self, w_rpr: ET.Element, style_info: StyleInfo):
        """
        Extract styling from Word run properties (w:rPr).

        Handles:
        - w:rFonts (font family)
        - w:color (text color)
        - w:sz (font size)
        - w:b (bold flag)
        - w:i (italic flag)
        """
        # Font family
        rfonts_elem = w_rpr.find(f'{{{WORD_NS}}}rFonts', OMML_NAMESPACES)
        if rfonts_elem is not None:
            # Try different attributes for font name
            font = (rfonts_elem.get(f'{{{WORD_NS}}}ascii') or
                   rfonts_elem.get(f'{{{WORD_NS}}}hAnsi') or
                   rfonts_elem.get(f'{{{WORD_NS}}}cs'))
            if font:
                style_info.font_family = font

        # Text color
        color_elem = w_rpr.find(f'{{{WORD_NS}}}color', OMML_NAMESPACES)
        if color_elem is not None:
            color_val = color_elem.get(f'{{{WORD_NS}}}val')
            if color_val is None:
                color_val = color_elem.get('val')
            if color_val and color_val != 'auto':
                style_info.color = color_val

        # Font size (in half-points)
        sz_elem = w_rpr.find(f'{{{WORD_NS}}}sz', OMML_NAMESPACES)
        if sz_elem is not None:
            sz_val = sz_elem.get(f'{{{WORD_NS}}}val')
            if sz_val is None:
                sz_val = sz_elem.get('val')
            if sz_val:
                try:
                    # Convert half-points to points
                    style_info.font_size = float(sz_val) / 2
                except ValueError:
                    pass

        # Bold flag (from Word, if not already set by m:sty)
        # IMPORTANT: <w:b/> (no val) means bold=true, but <w:b val="0"/> means bold=false
        if not style_info.is_bold:
            b_elem = w_rpr.find(f'{{{WORD_NS}}}b', OMML_NAMESPACES)
            if b_elem is not None:
                # Check the val attribute - "0" or "false" means NOT bold
                b_val = b_elem.get(f'{{{WORD_NS}}}val')
                if b_val is None:
                    b_val = b_elem.get('val')  # Try without namespace
                # Only set bold=True if val is not "0" or "false"
                # Presence without val attribute means true
                if b_val not in ('0', 'false'):
                    style_info.is_bold = True
                    if not style_info.mathvariant:
                        style_info.mathvariant = 'bold'

        # Italic flag (from Word, if not already set by m:sty)
        # IMPORTANT: <w:i/> (no val) means italic=true, but <w:i val="0"/> means italic=false
        if not style_info.is_italic:
            i_elem = w_rpr.find(f'{{{WORD_NS}}}i', OMML_NAMESPACES)
            if i_elem is not None:
                # Check the val attribute - "0" or "false" means NOT italic
                i_val = i_elem.get(f'{{{WORD_NS}}}val')
                if i_val is None:
                    i_val = i_elem.get('val')  # Try without namespace
                # Only set italic=True if val is not "0" or "false"
                # Presence without val attribute means true
                if i_val not in ('0', 'false'):
                    style_info.is_italic = True
                    if not style_info.mathvariant:
                        style_info.mathvariant = 'italic'
                    elif style_info.mathvariant == 'bold':
                        style_info.mathvariant = 'bold-italic'

        # Background/highlight color (w:highlight element)
        highlight_elem = w_rpr.find(f'{{{WORD_NS}}}highlight', OMML_NAMESPACES)
        if highlight_elem is not None:
            highlight_val = highlight_elem.get(f'{{{WORD_NS}}}val')
            if highlight_val is None:
                highlight_val = highlight_elem.get('val')
            if highlight_val:
                style_info.background_color = highlight_val

    def _apply_structural_styling(
        self,
        mathml_root: ET.Element,
        structural_styling: Dict[Tuple[str, int], StructuralStyleInfo]
    ) -> ET.Element:
        """
        Wrap structural MathML elements in mstyle with mathcolor.

        This method finds structural elements like mfrac, msqrt, mroot and wraps
        them in <mstyle mathcolor="#color"> elements based on the styling extracted
        from OMML.

        Args:
            mathml_root: The MathML root element
            structural_styling: Dict mapping (element_type, position) -> StructuralStyleInfo

        Returns:
            Modified MathML root element

        Note:
            Maps OMML types to MathML tags:
            - 'f' -> 'mfrac'
            - 'rad' -> 'msqrt' or 'mroot'
        """
        # Map OMML element types to MathML tag names
        element_map = {
            'f': ['mfrac'],
            'rad': ['msqrt', 'mroot'],
            # Future: 'func': [...], 'd': [...]
        }

        # Track positions for each MathML tag type
        positions = {'mfrac': 0, 'msqrt': 0, 'mroot': 0}

        # Collect elements to wrap (can't modify tree during iteration)
        elements_to_wrap = []

        for elem in mathml_root.iter():
            local_tag = elem.tag.split('}')[-1] if '}' in elem.tag else elem.tag

            # Check if this is a structural element we're tracking
            for omml_type, mathml_tags in element_map.items():
                if local_tag in mathml_tags:
                    # Check if we have styling for this position
                    key = (omml_type, positions.get(local_tag, 0))
                    if key in structural_styling:
                        style_info = structural_styling[key]
                        if style_info.color:
                            elements_to_wrap.append((elem, style_info.color))

                    # Increment position counter for this tag type
                    positions[local_tag] = positions.get(local_tag, 0) + 1

        # Wrap elements in mstyle
        for elem, color in elements_to_wrap:
            self._wrap_element_in_mstyle(elem, color)

        return mathml_root

    def _wrap_element_in_mstyle(self, elem: ET.Element, color: str):
        """
        Wrap element in mstyle with mathcolor attribute.

        Args:
            elem: The MathML element to wrap
            color: Color value (hex string without #)

        Note:
            Modifies the tree in place by:
            1. Creating <mstyle mathcolor="#color">
            2. Replacing elem with mstyle in parent
            3. Adding elem as child of mstyle
        """
        parent = elem.getparent()
        if parent is None:
            # Can't wrap root element
            return

        # Get namespace from element
        namespace = elem.tag.split('}')[0][1:] if '}' in elem.tag else None

        # Create mstyle element with same namespace
        if namespace:
            mstyle_tag = f'{{{namespace}}}mstyle'
        else:
            mstyle_tag = 'mstyle'

        mstyle = ET.Element(mstyle_tag)
        mstyle.set('mathcolor', f'#{color}')

        # Replace elem with mstyle in parent
        index = list(parent).index(elem)
        parent.remove(elem)
        parent.insert(index, mstyle)

        # Add elem as child of mstyle
        mstyle.append(elem)

    def apply_styling_to_mathml(
        self,
        mathml_xml: str,
        styling_map: Dict[Tuple[int, str], StyleInfo],
        structural_styling: Dict[Tuple[str, int], StructuralStyleInfo] = None
    ) -> str:
        """
        Apply styling information to MathML elements.

        Adds mathvariant, mathcolor, and other attributes to <mi>, <mo>, <mn>
        elements based on the styling map extracted from OMML.

        Args:
            mathml_xml: MathML XML string (output from XSLT)
            styling_map: Styling map from extract_styling_from_omml()

        Returns:
            Enhanced MathML XML string with styling attributes

        Examples:
            >>> processor = StylingPostProcessor()
            >>> mathml = '<math><mi>Z</mi></math>'
            >>> styling = {(0, 'Z'): StyleInfo(text='Z', position=0, mathvariant='bold')}
            >>> enhanced = processor.apply_styling_to_mathml(mathml, styling)
            >>> 'mathvariant="bold"' in enhanced
            True
        """
        if not styling_map and not structural_styling:
            return mathml_xml

        try:
            root = ET.fromstring(mathml_xml.encode('utf-8'))
        except ET.XMLSyntaxError:
            # If parsing fails, return original
            return mathml_xml

        # FIRST: Apply structural styling (wraps elements in mstyle)
        if structural_styling:
            root = self._apply_structural_styling(root, structural_styling)

        # THEN: Apply leaf styling (add attributes to mi, mo, mn)
        # Track position for matching
        position = 0

        # Find all <mi>, <mo>, <mn> elements (identifier, operator, number)
        for elem in root.iter():
            local_tag = elem.tag.split('}')[-1] if '}' in elem.tag else elem.tag

            if local_tag not in ('mi', 'mo', 'mn'):
                continue

            text = elem.text
            if not text or not text.strip():
                continue

            text = text.strip()

            # Look up styling
            style_info = styling_map.get((position, text))
            if style_info is None:
                # Try without position (fallback for simple cases)
                for key, info in styling_map.items():
                    if info.text == text:
                        style_info = info
                        break

            if style_info is None:
                position += 1
                continue

            # Apply mathvariant ONLY to identifiers (<mi>), not operators (<mo>) or numbers (<mn>)
            # Operators should never have mathvariant as they're already upright by default
            # This prevents wrapping operators in \mathrm{} which causes output pollution
            if style_info.mathvariant and local_tag == 'mi':
                elem.set('mathvariant', style_info.mathvariant)

            # Apply mathcolor (if color present)
            if style_info.color:
                elem.set('mathcolor', f'#{style_info.color}')

            # Apply font-family (custom attribute, will be handled by LaTeX converter)
            # Skip Cambria Math as it's the default math font
            if style_info.font_family and style_info.font_family != 'Cambria Math':
                elem.set('fontfamily', style_info.font_family)

            # Apply font-size (custom attribute)
            if style_info.font_size:
                elem.set('fontsize', str(style_info.font_size))

            position += 1

        # Convert back to string
        return ET.tostring(root, encoding='unicode')

    def has_styling(self, omml_xml: str) -> bool:
        """
        Quick check if OMML contains any styling information.

        Used for optimization: skip post-processing if no styling present.

        Args:
            omml_xml: OMML XML string

        Returns:
            True if styling (m:rPr, w:rPr, m:sty, m:scr, w:highlight) found
        """
        # Quick string search (much faster than parsing)
        return any(marker in omml_xml for marker in [
            'm:rPr', 'm:sty', 'm:scr',  # OMML math properties
            'w:rPr', 'w:color', 'w:b', 'w:i', 'w:highlight',  # Word properties
        ])

    def detect_uniform_background(self, omml_xml: str) -> Optional[str]:
        """
        Detect if entire equation has uniform background/highlight color.

        Checks all <w:highlight> elements in the OMML to determine if they
        all have the same color value, indicating the entire equation has
        a uniform background.

        Args:
            omml_xml: OMML XML string

        Returns:
            Background color name (e.g., 'yellow') if uniform, None otherwise

        Examples:
            >>> processor = StylingPostProcessor()
            >>> omml = '<m:oMath><m:r><w:rPr><w:highlight w:val="yellow"/></w:rPr><m:t>x</m:t></m:r></m:oMath>'
            >>> processor.detect_uniform_background(omml)
            'yellow'
        """
        try:
            root = ET.fromstring(omml_xml.encode('utf-8'))
        except ET.XMLSyntaxError:
            return None

        background_colors = set()
        has_elements_without_background = False

        # Find all <m:r> (run) elements and check for w:highlight
        for r_elem in root.iter(f'{{{OMML_NS}}}r'):
            # Get text content - skip empty runs
            t_elem = r_elem.find(f'{{{OMML_NS}}}t', OMML_NAMESPACES)
            if t_elem is None or not t_elem.text or not t_elem.text.strip():
                continue

            # Check for Word run properties
            w_rpr = r_elem.find(f'{{{WORD_NS}}}rPr', OMML_NAMESPACES)
            if w_rpr is None:
                has_elements_without_background = True
                continue

            # Check for highlight element
            highlight_elem = w_rpr.find(f'{{{WORD_NS}}}highlight', OMML_NAMESPACES)
            if highlight_elem is None:
                has_elements_without_background = True
                continue

            # Get highlight value
            highlight_val = highlight_elem.get(f'{{{WORD_NS}}}val')
            if highlight_val is None:
                highlight_val = highlight_elem.get('val')

            if highlight_val:
                background_colors.add(highlight_val.lower())
            else:
                has_elements_without_background = True

        # Also check structural elements (fractions, radicals) for highlight in ctrlPr
        for elem_type in ['f', 'rad', 'func', 'd', 'sSup', 'sSub', 'sSubSup']:
            for struct_elem in root.iter(f'{{{OMML_NS}}}{elem_type}'):
                # Find properties element (fPr, radPr, etc.)
                pr_elem = struct_elem.find(f'{{{OMML_NS}}}{elem_type}Pr', OMML_NAMESPACES)
                if pr_elem is None:
                    continue

                # Check ctrlPr for w:rPr with highlight
                ctrl_pr = pr_elem.find(f'{{{OMML_NS}}}ctrlPr', OMML_NAMESPACES)
                if ctrl_pr is not None:
                    w_rpr = ctrl_pr.find(f'{{{WORD_NS}}}rPr', OMML_NAMESPACES)
                    if w_rpr is not None:
                        highlight_elem = w_rpr.find(f'{{{WORD_NS}}}highlight', OMML_NAMESPACES)
                        if highlight_elem is not None:
                            highlight_val = highlight_elem.get(f'{{{WORD_NS}}}val')
                            if highlight_val is None:
                                highlight_val = highlight_elem.get('val')
                            if highlight_val:
                                background_colors.add(highlight_val.lower())

        # Uniform background only if:
        # 1. We found exactly one unique background color
        # 2. ALL content elements have that background (none are without)
        if len(background_colors) == 1 and not has_elements_without_background:
            return list(background_colors)[0]

        # Not uniform
        return None


def fix_mathml_table_structure(mathml_xml: str) -> str:
    """
    Fix malformed table structure from OMML2MML.xsl transformation.

    The Microsoft XSLT outputs <mtr><mrow>content</mrow></mtr> but MathML
    requires <mtr><mtd>content</mtd></mtr> for proper table structure.

    This function:
    1. Finds all <mtr> elements with <mrow> children (should be <mtd>)
    2. Wraps each <mrow> in <mtd> to create proper table cells
    3. Preserves all content and attributes

    Args:
        mathml_xml: MathML XML string (possibly malformed)

    Returns:
        Fixed MathML with proper table structure

    Examples:
        >>> mathml = '<math><mtable><mtr><mrow><mn>1</mn></mrow></mtr></mtable></math>'
        >>> fixed = fix_mathml_table_structure(mathml)
        >>> '<mtd>' in fixed
        True
    """
    try:
        root = ET.fromstring(mathml_xml.encode('utf-8'))
    except ET.XMLSyntaxError:
        # If parsing fails, return original
        return mathml_xml

    def _local_name(tag):
        """Extract local name without namespace."""
        return tag.split('}')[-1] if '}' in tag else tag

    def _make_tag(local_name, namespace=None):
        """Create tag with optional namespace."""
        if namespace:
            return f'{{{namespace}}}{local_name}'
        return local_name

    # Find all <mtr> elements (with or without namespace)
    # Collect all mtr elements first to avoid modification during iteration
    mtr_elements = [elem for elem in root.iter() if _local_name(elem.tag) == 'mtr']

    for mtr_elem in mtr_elements:
        # Check if mtr has direct <mrow> children (indicates malformed structure)
        mrow_children = [child for child in mtr_elem if _local_name(child.tag) == 'mrow']

        if mrow_children:
            # Get namespace from mtr element (or its parent mtable)
            mtr_ns = mtr_elem.tag.split('}')[0][1:] if '}' in mtr_elem.tag else None

            # Process each mrow child: wrap in mtd
            # IMPORTANT: We need to preserve order, so collect mtd elements first
            new_mtd_elements = []
            for mrow_child in mrow_children:
                # Create new mtd element with same namespace as mtr
                mtd_elem = ET.Element(_make_tag('mtd', mtr_ns))

                # Move all children from mrow to mtd
                for grandchild in list(mrow_child):
                    mtd_elem.append(grandchild)

                # Preserve mrow text content if any
                if mrow_child.text and mrow_child.text.strip():
                    if len(mtd_elem) == 0:
                        mtd_elem.text = mrow_child.text
                    else:
                        # Prepend text to first child
                        mtd_elem[0].text = (mrow_child.text or '') + (mtd_elem[0].text or '')

                new_mtd_elements.append(mtd_elem)

            # Now remove all mrow children and add mtd elements
            for mrow_child in mrow_children:
                mtr_elem.remove(mrow_child)

            for mtd_elem in new_mtd_elements:
                mtr_elem.append(mtd_elem)

    # Convert back to string
    return ET.tostring(root, encoding='unicode')


def apply_styling_to_mathml_from_omml(omml_xml: str, mathml_xml: str) -> str:
    """
    Convenience function to extract styling from OMML and apply to MathML.

    Also fixes table structure issues from OMML2MML.xsl transformation.
    Detects and adds uniform background color as a custom attribute.

    Args:
        omml_xml: Original OMML XML string
        mathml_xml: MathML XML string (from XSLT conversion)

    Returns:
        Enhanced MathML with styling attributes and fixed table structure

    Examples:
        >>> omml = '<m:oMath><m:r><m:rPr><m:sty m:val="b"/></m:rPr><m:t>Z</m:t></m:r></m:oMath>'
        >>> mathml = '<math><mi>Z</mi></math>'
        >>> enhanced = apply_styling_to_mathml_from_omml(omml, mathml)
        >>> 'mathvariant="bold"' in enhanced
        True
    """
    processor = StylingPostProcessor()

    # CRITICAL FIX: Fix table structure FIRST (before styling)
    # The OMML2MML.xsl outputs <mtr><mrow> instead of <mtr><mtd>
    mathml_xml = fix_mathml_table_structure(mathml_xml)

    # Detect uniform background color from OMML
    uniform_background = processor.detect_uniform_background(omml_xml)

    # Quick check: skip if no styling and no background
    if not processor.has_styling(omml_xml) and not uniform_background:
        return mathml_xml

    # Extract BOTH leaf and structural styling
    styling_map = processor.extract_styling_from_omml(omml_xml)
    structural_styling = processor.extract_structural_styling_from_omml(omml_xml)

    # Apply both types of styling
    result = processor.apply_styling_to_mathml(
        mathml_xml,
        styling_map,
        structural_styling
    )

    # Add uniform background as custom attribute on <math> root element
    # This will be detected by the LaTeX converter and applied as \bbox
    if uniform_background:
        result = _add_background_attribute(result, uniform_background)

    return result


def _add_background_attribute(mathml_xml: str, background_color: str) -> str:
    """
    Add data-background attribute to MathML root element.

    Args:
        mathml_xml: MathML XML string
        background_color: Background color name (e.g., 'yellow')

    Returns:
        MathML with data-background attribute added to root element
    """
    try:
        root = ET.fromstring(mathml_xml.encode('utf-8'))
    except ET.XMLSyntaxError:
        return mathml_xml

    # Add custom attribute to root element
    root.set('data-background', background_color)

    return ET.tostring(root, encoding='unicode')