# Copyright (c) Opendatalab. All rights reserved.
"""
Office Math Markup Language (OMML)

Adapted from https://github.com/xiilei/dwml/blob/master/dwml/omml.py
On 23/01/2025
"""
import re

import lxml.etree as ET
from loguru import logger
from pylatexenc.latexencode import UnicodeToLatexEncoder

from .latex_dict import (
    ALN,
    ARR,
    BACKSLASH,
    BLANK,
    BRK,
    CHARS,
    CHR,
    CHR_BO,
    CHR_DEFAULT,
    D_DEFAULT,
    F_DEFAULT,
    FUNC,
    FUNC_PLACE,
    LIM_FUNC,
    LIM_TO,
    LIM_UPP,
    POS,
    POS_DEFAULT,
    RAD,
    RAD_DEFAULT,
    SUB,
    SUP,
    D,
    F,
    M,
    T,
)

OMML_NS = "{http://schemas.openxmlformats.org/officeDocument/2006/math}"

# Mapping from OMML <m:scr> values to LaTeX math font commands.
# Used in do_r to convert math script/font style to appropriate LaTeX commands.
SCR_TO_LATEX = {
    "script":       "\\mathscr{{{0}}}",       # 手写体/花体 — \mathscr covers both upper and lowercase
    "fraktur":      "\\mathfrak{{{0}}}",       # 德国哥特体 — \mathfrak for upper and lowercase
    "double-struck": "\\mathbb{{{0}}}",        # 双线体/黑板粗体 — \mathbb
    "sans-serif":   "\\mathsf{{{0}}}",         # 无衬线体
    "monospace":    "\\mathtt{{{0}}}",         # 等宽字体
}

LOWER_GROUP_LIMITS = ("\\underbrace{", "\\underbracket{", "\\underparen{")
UPPER_GROUP_LIMITS = ("\\overbrace{", "\\overbracket{", "\\overparen{")


def load(stream):
    tree = ET.parse(stream)
    for omath in tree.findall(OMML_NS + "oMath"):
        yield oMath2Latex(omath)


def load_string(string):
    root = ET.fromstring(string)
    for omath in root.findall(OMML_NS + "oMath"):
        yield oMath2Latex(omath)


def escape_latex(strs):
    last = None
    new_chr = []
    strs = strs.replace(r"\\", "\\")
    for c in strs:
        if (c in CHARS) and (last != BACKSLASH):
            new_chr.append(BACKSLASH + c)
        else:
            new_chr.append(c)
        last = c
    return BLANK.join(new_chr)


def get_val(key, default=None, store=CHR):
    if key is not None:
        return key if not store else store.get(key, key)
    else:
        return default


def _normalize_latex_delimiter(delimiter):
    """将 Word OMML 定界符字符转换为 LaTeX/KaTeX 可渲染的定界符。"""
    if delimiter in ("\u2225", "\u2016"):
        return r"\|"
    return delimiter


class Tag2Method:
    def call_method(self, elm, stag=None):
        getmethod = self.tag2meth.get
        if stag is None:
            stag = elm.tag.replace(OMML_NS, "")
        method = getmethod(stag)
        if method:
            return method(self, elm)
        else:
            return None

    def process_children_list(self, elm, include=None):
        """
        process children of the elm,return iterable
        """
        for _e in list(elm):
            if OMML_NS not in _e.tag:
                continue
            stag = _e.tag.replace(OMML_NS, "")
            if include and (stag not in include):
                continue
            t = self.call_method(_e, stag=stag)
            if t is None:
                t = self.process_unknow(_e, stag)
                if t is None:
                    continue
            yield (stag, t, _e)

    def process_children_dict(self, elm, include=None):
        """
        process children of the elm,return dict
        """
        latex_chars = dict()
        for stag, t, e in self.process_children_list(elm, include):
            latex_chars[stag] = t
        return latex_chars

    def process_children(self, elm, include=None):
        """
        process children of the elm,return string
        """
        return BLANK.join(
            (
                t if not isinstance(t, Tag2Method) else str(t)
                for stag, t, e in self.process_children_list(elm, include)
            )
        )

    def process_unknow(self, elm, stag):
        return None


class Pr(Tag2Method):
    text = ""

    __val_tags = ("chr", "pos", "begChr", "endChr", "type")

    __innerdict = None  # can't use the __dict__

    """ common properties of element"""

    def __init__(self, elm):
        self.__innerdict = {}
        self.text = self.process_children(elm)

    def __str__(self):
        return self.text

    def __unicode__(self):
        return self.__str__(self)

    def __getattr__(self, name):
        return self.__innerdict.get(name, None)

    def do_brk(self, elm):
        self.__innerdict["brk"] = BRK
        return BRK

    def do_common(self, elm):
        stag = elm.tag.replace(OMML_NS, "")
        if stag in self.__val_tags:
            t = elm.get(f"{OMML_NS}val")
            self.__innerdict[stag] = t
        return None

    tag2meth = {
        "brk": do_brk,
        "chr": do_common,
        "pos": do_common,
        "begChr": do_common,
        "endChr": do_common,
        "type": do_common,
    }


class oMath2Latex(Tag2Method):
    """
    Convert oMath element of omml to latex
    """

    _t_dict = T

    __direct_tags = ("box", "sSub", "sSup", "sSubSup", "num", "den", "deg", "e")
    u = UnicodeToLatexEncoder(
        replacement_latex_protection="braces-all",
        unknown_char_policy="keep",
        unknown_char_warning=False,
    )

    def __init__(self, element):
        self._latex = self.process_children(element)

    def __str__(self):
        return self.latex.replace("  ", " ")

    def __unicode__(self):
        return self.__str__(self)

    def process_unknow(self, elm, stag):
        if stag in self.__direct_tags:
            return self.process_children(elm)
        elif stag[-2:] == "Pr":
            return Pr(elm)
        else:
            return None

    @property
    def latex(self):
        return self._latex

    def _apply_limit_marker(self, base_text, limit_text):
        if not isinstance(limit_text, str):
            return None

        latex_template = CHR.get(limit_text)
        if latex_template and "{0}" in latex_template:
            return latex_template.format(base_text)
        return None

    def _format_limit_like(self, base_text, limit_text, *, upper):
        marker_wrapped = self._apply_limit_marker(base_text, limit_text)
        if marker_wrapped is not None:
            return marker_wrapped

        if upper:
            if isinstance(base_text, str) and base_text.lstrip().startswith(UPPER_GROUP_LIMITS):
                return f"{base_text}{SUP.format(limit_text)}"
            return LIM_UPP.format(lim=limit_text, text=base_text)

        latex_s = LIM_FUNC.get(base_text)
        if latex_s:
            return latex_s.format(lim=limit_text)

        if isinstance(base_text, str) and base_text.lstrip().startswith(LOWER_GROUP_LIMITS):
            return f"{base_text}{SUB.format(limit_text)}"

        return f"\\underset{{{limit_text}}}{{{base_text}}}"

    def do_acc(self, elm):
        """
        the accent function
        """
        c_dict = self.process_children_dict(elm)
        latex_s = get_val(
            c_dict["accPr"].chr, default=CHR_DEFAULT.get("ACC_VAL"), store=CHR
        )
        return latex_s.format(c_dict["e"])

    def do_bar(self, elm):
        """
        the bar function
        """
        c_dict = self.process_children_dict(elm)
        pr = c_dict["barPr"]
        latex_s = get_val(pr.pos, default=POS_DEFAULT.get("BAR_VAL"), store=POS)
        return pr.text + latex_s.format(c_dict["e"])

    def do_d(self, elm):
        """
        the delimiter object
        """
        c_dict = self.process_children_dict(elm)
        pr = c_dict["dPr"]
        null = D_DEFAULT.get("null")

        s_val = _normalize_latex_delimiter(
            get_val(pr.begChr, default=D_DEFAULT.get("left"), store=T)
        )
        e_val = _normalize_latex_delimiter(
            get_val(pr.endChr, default=D_DEFAULT.get("right"), store=T)
        )
        delim = pr.text + D.format(
            left=null if not s_val else escape_latex(s_val),
            text=c_dict["e"],
            right=null if not e_val else escape_latex(e_val),
        )
        return delim

    def do_spre(self, elm):
        """
        the Pre-Sub-Superscript object -- Not support yet
        """

    def do_sub(self, elm):
        text = self.process_children(elm)
        return SUB.format(text)

    def do_sup(self, elm):
        text = self.process_children(elm)
        return SUP.format(text)

    def do_f(self, elm):
        """
        the fraction object
        """
        c_dict = self.process_children_dict(elm)
        pr = c_dict.get("fPr")
        if pr is None:
            # Handle missing fPr element gracefully
            logger.debug("Missing fPr element in fraction, using default formatting")
            latex_s = F_DEFAULT
            return latex_s.format(
                num=c_dict.get("num"),
                den=c_dict.get("den"),
            )
        latex_s = get_val(pr.type, default=F_DEFAULT, store=F)
        return pr.text + latex_s.format(num=c_dict.get("num"), den=c_dict.get("den"))

    def do_func(self, elm):
        """
        the Function-Apply object (Examples:sin cos)
        """
        c_dict = self.process_children_dict(elm)
        func_name = c_dict.get("fName")
        return func_name.replace(FUNC_PLACE, c_dict.get("e"))

    def do_fname(self, elm):
        """
        the func name
        """
        latex_chars = []
        for stag, t, e in self.process_children_list(elm):
            if stag == "r":
                if FUNC.get(t):
                    latex_chars.append(FUNC[t])
                else:
                    logger.warning("Function not supported, will default to text: %s", t)
                    if isinstance(t, str):
                        latex_chars.append(t)
            elif isinstance(t, str):
                latex_chars.append(t)
        t = BLANK.join(latex_chars)
        return t if FUNC_PLACE in t else t + FUNC_PLACE  # do_func will replace this

    def do_groupchr(self, elm):
        """
        the Group-Character object
        """
        c_dict = self.process_children_dict(elm)
        pr = c_dict["groupChrPr"]
        latex_s = get_val(pr.chr)
        return pr.text + latex_s.format(c_dict["e"])

    def do_rad(self, elm):
        """
        the radical object
        """
        c_dict = self.process_children_dict(elm)
        text = c_dict.get("e")
        deg_text = c_dict.get("deg")
        if deg_text:
            return RAD.format(deg=deg_text, text=text)
        else:
            return RAD_DEFAULT.format(text=text)

    def do_eqarr(self, elm):
        """
        the Array object.

        Handles two cases:
        1. Single-row eqArr with a right-aligned equation tag encoded as
           ``\\#(n)`` at the end of the row (OMML column-alignment syntax).
           The ``#`` column separator and the equation number ``(n)`` are
           converted to LaTeX ``\\tag{n}`` so KaTeX can render the tag,
           and the unnecessary ``\\begin{array}{c}...\\end{array}`` wrapper
           is omitted.
        2. Single-row eqArr without a tag: content is returned as-is (no
           array wrapper needed).
        3. Multi-row eqArr: kept as ``\\begin{array}{c}...\\end{array}``.
        """
        rows = [t for stag, t, e in self.process_children_list(elm, include=("e",))]

        if len(rows) == 1:
            row = rows[0]
            # Detect the OMML equation-tag pattern: the text element "#(n)" is
            # stored verbatim inside the row; do_r converts "#" via pylatexenc
            # to "\# " (escaped hash with surrounding spaces due to brace-
            # protection stripping).  Match that at the end of the row,
            # allowing optional whitespace between "\#" and the opening "(".
            tag_match = re.search(r'\\#\s*\(([^)]*)\)\s*$', row)
            if tag_match:
                formula = row[:tag_match.start()].rstrip()
                tag_content = tag_match.group(1)
                return f'{formula}\\tag{{{tag_content}}}'
            # Single row without tag — no array wrapper required.
            return row

        return ARR.format(text=BRK.join(rows))

    def do_limlow(self, elm):
        """
        the Lower-Limit object
        """
        t_dict = self.process_children_dict(elm, include=("e", "lim"))
        return self._format_limit_like(
            t_dict.get("e", ""),
            t_dict.get("lim", ""),
            upper=False,
        )

    def do_limupp(self, elm):
        """
        the Upper-Limit object
        """
        t_dict = self.process_children_dict(elm, include=("e", "lim"))
        return self._format_limit_like(
            t_dict.get("e", ""),
            t_dict.get("lim", ""),
            upper=True,
        )

    def do_lim(self, elm):
        """
        the lower limit of the limLow object and the upper limit of the limUpp function
        """
        return self.process_children(elm).replace(LIM_TO[0], LIM_TO[1])

    def do_m(self, elm):
        """
        the Matrix object
        """
        rows = []
        for stag, t, e in self.process_children_list(elm):
            if stag == "mPr":
                pass
            elif stag == "mr":
                rows.append(t)
        return M.format(text=BRK.join(rows))

    def do_mr(self, elm):
        """
        a single row of the matrix m
        """
        return ALN.join(
            [t for stag, t, e in self.process_children_list(elm, include=("e",))]
        )

    def do_nary(self, elm):
        """
        the n-ary object
        """
        res = []
        bo = ""
        for stag, t, e in self.process_children_list(elm):
            if stag == "naryPr":
                # if <m:naryPr> contains no <m:chr>, the n-ary represents an integral
                bo = get_val(t.chr, default="\\int", store=CHR_BO)
            else:
                res.append(t)
        return bo + BLANK.join(res)

    def process_unicode(self, s):
        if s in CHARS:
            return BACKSLASH + s

        # Check T dictionary first for known math-mode symbols.
        # The T dictionary holds explicit math-mode LaTeX mappings and takes precedence
        # over pylatexenc, which uses text-mode mappings by default and therefore produces
        # text-mode commands like \textperiodcentered (for U+00B7 ·) that are invalid
        # inside math environments.
        t_result = self._t_dict.get(s)
        if t_result is not None:
            return t_result

        out_latex_str = self.u.unicode_to_latex(s)

        # pylatexenc常把数学字符包成 {\ensuremath{...}}，这里只剥离外层包装，
        # 不能删除内部LaTeX命令的闭合花括号。
        if out_latex_str.startswith(r"{\ensuremath{") and out_latex_str.endswith("}}"):
            out_latex_str = out_latex_str[len(r"{\ensuremath{") : -2]
        elif out_latex_str.startswith(r"\ensuremath{") and out_latex_str.endswith("}"):
            out_latex_str = out_latex_str[len(r"\ensuremath{") : -1]
        elif (
            s.startswith("{") is False
            and out_latex_str.startswith("{")
            and s.endswith("}") is False
            and out_latex_str.endswith("}")
        ):
            out_latex_str = f" {out_latex_str[1:-1]} "

        # Do NOT wrap remaining content in \text{}.
        # Previously this code matched any string starting with "\text" and wrapped it
        # again, producing invalid constructs like \text{ \textperiodcentered } for
        # textcomp symbols.  Characters that truly need text mode should be mapped in
        # the T dictionary above; for all others we keep the pylatexenc output as-is.

        return out_latex_str


    def do_r(self, elm):
        """
        Get text from 'r' element,And try convert them to latex symbols
        @todo text style support , (sty)
        @todo \text (latex pure text support)
        """
        _str = []
        _base_str = []
        found_text = elm.findtext(f"./{OMML_NS}t")
        if found_text:
            for s in found_text:
                out_latex_str = self.process_unicode(s)
                _str.append(out_latex_str)
                _base_str.append(s)

        proc_str = escape_latex(BLANK.join(_str))
        base_proc_str = BLANK.join(_base_str)

        if "{" not in base_proc_str and "\\{" in proc_str:
            proc_str = proc_str.replace("\\{", "{")

        if "}" not in base_proc_str and "\\}" in proc_str:
            proc_str = proc_str.replace("\\}", "}")

        # Handle <m:scr> math font style (script, fraktur, double-struck, etc.)
        # OMML encodes math alphabets via <m:rPr><m:scr m:val="..."/> rather than
        # Unicode math-alphabet codepoints, so we must apply the LaTeX wrapper here.
        rPr = elm.find(f"{OMML_NS}rPr")
        if rPr is not None:
            scr_elem = rPr.find(f"{OMML_NS}scr")
            if scr_elem is not None:
                scr_val = scr_elem.get(f"{OMML_NS}val")
                latex_template = SCR_TO_LATEX.get(scr_val)
                if latex_template and proc_str.strip():
                    proc_str = latex_template.format(proc_str.strip())

        return proc_str

    tag2meth = {
        "acc": do_acc,
        "r": do_r,
        "bar": do_bar,
        "sub": do_sub,
        "sup": do_sup,
        "f": do_f,
        "func": do_func,
        "fName": do_fname,
        "groupChr": do_groupchr,
        "d": do_d,
        "rad": do_rad,
        "eqArr": do_eqarr,
        "limLow": do_limlow,
        "limUpp": do_limupp,
        "lim": do_lim,
        "m": do_m,
        "mr": do_mr,
        "nary": do_nary,
    }
