"""Simple notation and its canonical linear form.

Non-normative; the rules are spec/05-hashes.md, "The strict projection".
A math source or a styled script either folds to one canonical linear form, or is complex.
"""

import unicodedata

from .text import collapse_ws, is_ws, remove_ws, trim

MINUS = "\u2212"
DOT = "\u22c5"
PRIME = "\u2032"

# Characters that fold to one canonical character inside notation.
CHAR_FOLD = {
    "-": MINUS,
    "\u2013": MINUS,
    "\u2212": MINUS,
    "\u00b7": DOT,
    "\u2219": DOT,
    "\u22c5": DOT,
    "'": PRIME,
    "\u2032": PRIME,
}

OPERATORS = set("+=/<>.,()") | {MINUS, DOT, PRIME, "\u00d7", "\u2264", "\u2265", "\u00b1"}

TEX_GREEK = {
    "alpha": "\u03b1", "beta": "\u03b2", "gamma": "\u03b3", "delta": "\u03b4",
    "epsilon": "\u03f5", "varepsilon": "\u03b5", "zeta": "\u03b6", "eta": "\u03b7",
    "theta": "\u03b8", "vartheta": "\u03d1", "iota": "\u03b9", "kappa": "\u03ba",
    "varkappa": "\u03f0", "lambda": "\u03bb", "mu": "\u03bc", "nu": "\u03bd",
    "xi": "\u03be", "pi": "\u03c0", "varpi": "\u03d6", "rho": "\u03c1",
    "varrho": "\u03f1", "sigma": "\u03c3", "varsigma": "\u03c2", "tau": "\u03c4",
    "upsilon": "\u03c5", "phi": "\u03d5", "varphi": "\u03c6", "chi": "\u03c7",
    "psi": "\u03c8", "omega": "\u03c9",
    "Gamma": "\u0393", "Delta": "\u0394", "Theta": "\u0398", "Lambda": "\u039b",
    "Xi": "\u039e", "Pi": "\u03a0", "Sigma": "\u03a3", "Upsilon": "\u03a5",
    "Phi": "\u03a6", "Psi": "\u03a8", "Omega": "\u03a9",
}
TEX_SYMBOLS = {
    "times": "\u00d7", "cdot": DOT, "pm": "\u00b1", "le": "\u2264", "leq": "\u2264",
    "ge": "\u2265", "geq": "\u2265", "lt": "<", "gt": ">", "prime": PRIME,
    "AA": "\u00c5",
}
TEX_WRAPPERS = {"mathrm", "text", "textrm", "mathup"}
TEX_DECLARATIONS = {"rm"}
TEX_SPACES = {"quad", "qquad", "enspace", "thinspace", "medspace", "thickspace"}
TEX_SPACE_SYMBOLS = {",", ";", ":", "!", " "}

TYPST_GREEK = {
    "alpha": "\u03b1", "beta": "\u03b2", "gamma": "\u03b3", "delta": "\u03b4",
    "epsilon": "\u03b5", "zeta": "\u03b6", "eta": "\u03b7", "theta": "\u03b8",
    "iota": "\u03b9", "kappa": "\u03ba", "lambda": "\u03bb", "mu": "\u03bc",
    "nu": "\u03bd", "xi": "\u03be", "omicron": "\u03bf", "pi": "\u03c0",
    "rho": "\u03c1", "sigma": "\u03c3", "tau": "\u03c4", "upsilon": "\u03c5",
    "phi": "\u03c6", "chi": "\u03c7", "psi": "\u03c8", "omega": "\u03c9",
    "Alpha": "\u0391", "Beta": "\u0392", "Gamma": "\u0393", "Delta": "\u0394",
    "Epsilon": "\u0395", "Zeta": "\u0396", "Eta": "\u0397", "Theta": "\u0398",
    "Iota": "\u0399", "Kappa": "\u039a", "Lambda": "\u039b", "Mu": "\u039c",
    "Nu": "\u039d", "Xi": "\u039e", "Omicron": "\u039f", "Pi": "\u03a0",
    "Rho": "\u03a1", "Sigma": "\u03a3", "Tau": "\u03a4", "Upsilon": "\u03a5",
    "Phi": "\u03a6", "Chi": "\u03a7", "Psi": "\u03a8", "Omega": "\u03a9",
    "epsilon.alt": "\u03f5", "theta.alt": "\u03d1", "phi.alt": "\u03d5",
    "pi.alt": "\u03d6", "rho.alt": "\u03f1", "sigma.alt": "\u03c2",
    "kappa.alt": "\u03f0", "beta.alt": "\u03d0",
}
TYPST_SYMBOLS = {
    "plus": "+", "minus": MINUS, "times": "\u00d7", "dot": DOT, "dot.op": DOT,
    "dot.c": DOT, "plus.minus": "\u00b1", "lt": "<", "gt": ">", "lt.eq": "\u2264",
    "gt.eq": "\u2265", "eq": "=", "prime": PRIME, "slash": "/",
}
TYPST_WRAPPERS = {"upright"}
TYPST_SPACES = {"thin", "med", "thick", "quad", "wide", "space"}


class NotSimple(Exception):
    pass


# The Latin and Greek letter blocks of spec/05 \u00a75.4: Basic Latin through Latin Extended-B, Greek and Coptic without
# the Coptic letters, Latin Extended Additional, and Greek Extended.
LETTER_RANGES = ((0x0041, 0x024F), (0x0370, 0x03E1), (0x03F0, 0x03FF), (0x1E00, 0x1EFF), (0x1F00, 0x1FFF))


def is_letter(ch):
    """A letter of General Category L in the Latin or Greek blocks, or the micro sign."""
    if ch == "\u00b5":
        return True
    if not unicodedata.category(ch).startswith("L"):
        return False
    cp = ord(ch)
    return any(lo <= cp <= hi for lo, hi in LETTER_RANGES)


def fold_char(ch):
    return CHAR_FOLD.get(ch, ch)


def is_simple_char(ch):
    ch = fold_char(ch)
    return is_letter(ch) or unicodedata.category(ch) == "Nd" or ch in OPERATORS


def script_type(ch):
    d = unicodedata.decomposition(ch)
    if d.startswith("<super>"):
        return "^", "".join(chr(int(h, 16)) for h in d.split()[1:])
    if d.startswith("<sub>"):
        return "_", "".join(chr(int(h, 16)) for h in d.split()[1:])
    return None, None


class Atom:
    __slots__ = ("base", "primes", "sub", "sup")

    def __init__(self, base=""):
        self.base = base
        self.primes = 0
        self.sub = None
        self.sup = None


def _canon(atoms):
    out = []
    for a in atoms:
        if a.sup is not None and a.sup and all(
            x.base == PRIME and x.sub is None and x.sup is None and x.primes == 0 for x in a.sup
        ):
            a.primes += len(a.sup)
            a.sup = None
        s = a.base + PRIME * a.primes
        if a.sub is not None:
            s += "_{" + _canon(a.sub) + "}"
        if a.sup is not None:
            s += "^{" + _canon(a.sup) + "}"
        out.append(s)
    return "".join(out)


def _attach(atoms, op, arg):
    if not atoms:
        atoms.append(Atom(""))
    a = atoms[-1]
    if op == "_":
        if a.sub is not None:
            raise NotSimple("double subscript")
        a.sub = arg
    else:
        if a.sup is not None:
            raise NotSimple("double superscript")
        a.sup = arg


def _char_atom(ch):
    if not is_simple_char(ch):
        raise NotSimple("character %r" % ch)
    return Atom(fold_char(ch))


# --- TeX ---------------------------------------------------------------------


def _tex_tokens(src):
    toks = []
    i, n = 0, len(src)
    while i < n:
        ch = src[i]
        if ch == "\\":
            j = i + 1
            if j < n and src[j].isascii() and src[j].isalpha():
                while j < n and src[j].isascii() and src[j].isalpha():
                    j += 1
                toks.append(("cw", src[i + 1:j]))
                while j < n and is_ws(src[j]):
                    j += 1
                i = j
                continue
            if j < n:
                toks.append(("cs", src[j]))
                i = j + 1
                continue
            raise NotSimple("trailing backslash")
        if is_ws(ch) or ch == "~":
            i += 1
            continue
        toks.append(("ch", ch))
        i += 1
    return toks


class _TexParser:
    def __init__(self, toks):
        self.t = toks
        self.i = 0

    def peek(self):
        return self.t[self.i] if self.i < len(self.t) else None

    def take(self):
        tok = self.peek()
        self.i += 1
        return tok

    def seq(self, closing):
        atoms = []
        while True:
            tok = self.peek()
            if tok is None:
                if closing:
                    raise NotSimple("unclosed group")
                return atoms
            if tok == ("ch", "}"):
                if not closing:
                    raise NotSimple("unbalanced brace")
                self.take()
                return atoms
            self.take()
            kind, val = tok
            if kind == "ch" and val in "_^":
                _attach(atoms, val, self.arg())
            elif kind == "ch" and val == "'":
                if not atoms:
                    atoms.append(Atom(""))
                atoms[-1].primes += 1
            elif kind == "ch" and val == "{":
                inner = self.seq(True)
                nxt = self.peek()
                if nxt in (("ch", "_"), ("ch", "^")) and len(inner) > 1:
                    raise NotSimple("script on a multi-atom group")
                atoms.extend(inner)  # an empty group leaves a following script to the previous atom or an empty base
            else:
                atoms.extend(self.command_or_char(kind, val))

    def command_or_char(self, kind, val):
        if kind == "ch":
            return [_char_atom(val)]
        if kind == "cs":
            if val in TEX_SPACE_SYMBOLS:
                return []
            raise NotSimple("control symbol \\" + val)
        if val in TEX_GREEK:
            return [Atom(TEX_GREEK[val])]
        if val in TEX_SYMBOLS:
            return [Atom(TEX_SYMBOLS[val])]
        if val in TEX_SPACES or val in TEX_DECLARATIONS:
            return []
        if val in TEX_WRAPPERS:
            return self.arg()
        if val in ("left", "right"):
            tok = self.take()
            if tok == ("ch", "(" if val == "left" else ")"):
                return [Atom(tok[1])]
            raise NotSimple("\\left or \\right with another delimiter")
        raise NotSimple("command \\" + val)

    def arg(self):
        tok = self.take()
        if tok is None:
            raise NotSimple("missing argument")
        kind, val = tok
        if tok == ("ch", "{"):
            return self.seq(True)
        if kind == "ch" and val in "}_^'":
            raise NotSimple("bad argument")
        return self.command_or_char(kind, val)


def tex_atoms(src):
    return _TexParser(_tex_tokens(src)).seq(False)


# --- Typst -------------------------------------------------------------------


def _typst_tokens(src):
    toks = []
    i, n = 0, len(src)
    while i < n:
        ch = src[i]
        if is_ws(ch):
            i += 1
            continue
        if ch == '"':
            j = i + 1
            buf = []
            while j < n and src[j] != '"':
                if src[j] == "\\":
                    raise NotSimple("escape in string")
                buf.append(src[j])
                j += 1
            if j >= n:
                raise NotSimple("unclosed string")
            toks.append(("str", "".join(buf)))
            i = j + 1
            continue
        if unicodedata.category(ch).startswith("L"):
            j = i
            while j < n and (
                unicodedata.category(src[j]).startswith("L")
                or (src[j] == "." and j + 1 < n and src[j + 1].isascii() and src[j + 1].isalpha() and j > i)
            ):
                j += 1
            toks.append(("id", src[i:j]))
            i = j
            continue
        if unicodedata.category(ch) == "Nd":
            j = i
            while j < n and unicodedata.category(src[j]) == "Nd":
                j += 1
            if j + 1 < n and src[j] == "." and unicodedata.category(src[j + 1]) == "Nd":
                j += 1
                while j < n and unicodedata.category(src[j]) == "Nd":
                    j += 1
            toks.append(("num", src[i:j]))
            i = j
            continue
        two = src[i:i + 2]
        if two in ("<=", ">="):
            toks.append(("ch", "\u2264" if two == "<=" else "\u2265"))
            i += 2
            continue
        if two in ("->", "=>", "!=", "<-", "==", ":="):
            raise NotSimple("operator " + two)
        toks.append(("ch", ch))
        i += 1
    return toks


class _TypstParser:
    def __init__(self, toks):
        self.t = toks
        self.i = 0

    def peek(self):
        return self.t[self.i] if self.i < len(self.t) else None

    def take(self):
        tok = self.peek()
        self.i += 1
        return tok

    def seq(self, closing):
        atoms = []
        while True:
            tok = self.peek()
            if tok is None:
                if closing:
                    raise NotSimple("unclosed parenthesis")
                return atoms
            if closing and tok == ("ch", ")"):
                self.take()
                return atoms
            self.take()
            kind, val = tok
            if kind == "ch" and val in "_^":
                _attach(atoms, val, self.arg())
            elif kind == "ch" and val == "'":
                if not atoms:
                    atoms.append(Atom(""))
                atoms[-1].primes += 1
            elif kind == "ch" and val == "(":
                atoms.append(Atom("("))
                atoms.extend(self.seq(True))
                atoms.append(Atom(")"))
            else:
                atoms.extend(self.single(kind, val))

    def single(self, kind, val):
        if kind == "num":
            return [Atom(c) for c in val]
        if kind == "str":
            return [_char_atom(c) for c in remove_ws(val)]
        if kind == "id":
            if len(val) == 1:
                return [_char_atom(val)]
            if val in TYPST_GREEK:
                return [Atom(TYPST_GREEK[val])]
            if val in TYPST_SYMBOLS:
                return [Atom(TYPST_SYMBOLS[val])]
            if val in TYPST_SPACES:
                return []
            if val in TYPST_WRAPPERS and self.peek() == ("ch", "("):
                self.take()
                return self.seq(True)
            raise NotSimple("identifier " + val)
        if val in "/*#&[]{}\\$!:;|~@":
            raise NotSimple("character %r" % val)
        return [_char_atom(val)]

    def arg(self):
        tok = self.take()
        if tok is None:
            raise NotSimple("missing argument")
        kind, val = tok
        if tok == ("ch", "("):
            return self.seq(True)
        if kind == "ch" and val in ")_^'":
            raise NotSimple("bad argument")
        return self.single(kind, val)


def typst_atoms(src):
    return _TypstParser(_typst_tokens(src)).seq(False)


# --- public ------------------------------------------------------------------


def fold_math(notation, source):
    """Return the canonical linear form of a math source, or None when it is complex."""
    src = unicodedata.normalize("NFC", source)
    try:
        atoms = tex_atoms(src) if notation == "tex" else typst_atoms(src)
        return _canon(atoms)
    except NotSimple:
        return None


def complex_key(notation, source):
    return notation + ":" + remove_ws(unicodedata.normalize("NFC", source))


def _fold_script_chars(chars):
    """Fold a run of Unicode script characters of one type, or None if any base is not simple."""
    op = None
    bases = []
    for ch in chars:
        t, b = script_type(ch)
        if t is None or (op is not None and t != op):
            return None
        op = t
        bases.append(b)
    base = "".join(bases)
    if not all(is_simple_char(c) for c in base):
        return None
    return op + "{" + "".join(fold_char(c) for c in base) + "}"


def _fold_chars(s, in_notation):
    """Fold the characters of text: script characters always, operators only inside notation."""
    out = []
    i, n = 0, len(s)
    while i < n:
        ch = s[i]
        t, b = script_type(ch)
        if t is not None:
            j = i
            while j < n and script_type(s[j])[0] == t:
                j += 1
            folded = _fold_script_chars(s[i:j])
            out.append(folded if folded is not None else s[i:j])
            i = j
            continue
        if in_notation:
            out.append(fold_char(ch))
        elif ch in ("-", "\u2013") and i + 1 < n and unicodedata.category(s[i + 1]) == "Nd":
            out.append(MINUS)
        else:
            out.append(ch)
        i += 1
    return "".join(out)


def fold_plain_text(text):
    return _fold_chars(unicodedata.normalize("NFC", text), in_notation=False)


def fold_styled_script(op, text):
    """The canonical form of text styled as a subscript (`_`) or superscript (`^`)."""
    t = unicodedata.normalize("NFC", text)
    s = remove_ws(t)
    simple = all(is_simple_char(c) or script_type(c)[0] is not None for c in s)
    if simple:
        inner = _fold_chars(s, in_notation=True)
        if op == "^" and inner and all(c == PRIME for c in inner):
            return inner
    else:
        inner = trim(collapse_ws(t))
    return op + "{" + inner + "}"