"""The plain, strict and lossy projections, and the five content values.

Non-normative; see spec/05-hashes.md.
"""

import hashlib
import hmac
import unicodedata

from . import jcs
from .model import LABELLED, MATH_KINDS, RUN_KINDS, VERBATIM_KINDS
from .notation import complex_key, fold_math, fold_plain_text, fold_styled_script
from .text import collapse_ws, is_ws, norm, remove_ws, trim

OBJ = "\ufffc"
MARK_OPEN = "\ufff9"
MARK_CLOSE = "\ufffb"


class Hasher:
    """SHA-256, or HMAC-SHA-256 under a key: every content value uses the same function."""

    def __init__(self, key=None):
        self.key = key

    def __call__(self, data):
        if isinstance(data, str):
            data = data.encode("utf-8")
        if self.key is None:
            return hashlib.sha256(data).hexdigest()
        return hmac.new(self.key, data, hashlib.sha256).hexdigest()


PLAIN = Hasher()


# --- runs --------------------------------------------------------------------


def _eq_source(doc, run):
    eq = doc.nodes[run["node"]]
    return eq["content"]["source"], eq["content"]["notation"]


def plain_runs(runs, doc=None):
    out = []
    for r in runs:
        k = r["kind"]
        if k in ("text", "link", "cite", "xref", "fnmark"):
            out.append(r["text"])
        elif k == "math":
            out.append("$" + r["source"] + "$")
        elif k == "equation":
            src, _ = _eq_source(doc, r)
            out.append("$$" + src + "$$")
        else:
            raise ValueError("run kind " + k)
    return "".join(out)


def _strict_math(notation, source, complexes):
    folded = fold_math(notation, source)
    if folded is None:
        complexes.append(complex_key(notation, source))
        return OBJ
    return folded


def _finish_strict(s):
    s = unicodedata.normalize("NFKC", s)
    return trim(collapse_ws(s))


def strict_runs_raw(runs, doc, complexes):
    out = []
    for r in runs:
        k = r["kind"]
        if k == "text":
            st = r.get("style", [])
            if "sup" in st:
                out.append(fold_styled_script("^", r["text"]))
            elif "sub" in st:
                out.append(fold_styled_script("_", r["text"]))
            else:
                out.append(fold_plain_text(r["text"]))
        elif k == "link":
            out.append(fold_plain_text(r["text"]))
        elif k == "math":
            out.append(_strict_math(r["notation"], r["source"], complexes))
        elif k == "equation":
            src, notation = _eq_source(doc, r)
            out.append(_strict_math(notation, src, complexes))
        elif k in ("cite", "xref", "fnmark"):
            out.append(MARK_OPEN + remove_ws(unicodedata.normalize("NFC", r["text"])) + MARK_CLOSE)
        else:
            raise ValueError("run kind " + k)
    return "".join(out)


def lossy_from_strict(s):
    out = []
    stack = []
    i, n = 0, len(s)
    while i < n:
        ch = s[i]
        if ch in "_^" and i + 1 < n and s[i + 1] == "{":
            stack.append("s")
            i += 2
            continue
        if ch == "{":
            stack.append("t")
            out.append(ch)
        elif ch == "}":
            if stack:
                if stack.pop() == "t":
                    out.append(ch)
            else:
                out.append(ch)
        elif ch == MARK_OPEN:
            j = s.find(MARK_CLOSE, i)
            i = n if j < 0 else j + 1
            continue
        else:
            out.append(ch)
        i += 1
    t = unicodedata.normalize("NFKC", "".join(out).casefold())
    return "".join(
        c for c in t if not is_ws(c) and unicodedata.category(c) != "Pd" and c != "\u00ad"
    )


def prose_length(lossy):
    return sum(1 for c in lossy if c != OBJ)


# --- nodes -------------------------------------------------------------------


def plain(doc, nid):
    n = doc.nodes[nid]
    k = n["kind"]
    if k in RUN_KINDS:
        return plain_runs(n["content"], doc)
    if k in MATH_KINDS:
        return "$$" + n["content"]["source"] + "$$"
    if k in VERBATIM_KINDS:
        return n["content"]["text"]
    return " ".join(plain(doc, x) for x in doc.text_leaves(nid))


def strict(doc, nid):
    """The strict canonical projection and the node's complex math keys, in order."""
    n = doc.nodes[nid]
    k = n["kind"]
    complexes = []
    if k in RUN_KINDS:
        s = strict_runs_raw(n["content"], doc, complexes)
    elif k in MATH_KINDS:
        s = _strict_math(n["content"]["notation"], n["content"]["source"], complexes)
    elif k in VERBATIM_KINDS:
        s = n["content"]["text"]
    else:
        parts = []
        for x in doc.text_leaves(nid):
            p, c = strict(doc, x)
            parts.append(p)
            complexes.extend(c)
        return _finish_strict(" ".join(parts)), complexes
    return _finish_strict(s), complexes


def lossy(doc, nid):
    return lossy_from_strict(strict(doc, nid)[0])


def _runs_for_own(runs):
    out = []
    for r in runs:
        if r["kind"] == "equation":
            out.append({"kind": "equation"})
        else:
            out.append(r)
    return out


def own_content(doc, nid, text_hash_of=None):
    n = doc.nodes[nid]
    k = n["kind"]
    o = {"kind": k}
    attrs = n.get("attrs", {})
    if k in LABELLED:
        o["label"] = n.get("label")
    if k in RUN_KINDS:
        o["runs"] = _runs_for_own(n["content"])
    if k in MATH_KINDS:
        o["notation"] = n["content"]["notation"]
        o["source"] = n["content"]["source"]
    if k in VERBATIM_KINDS:
        o["language"] = n["content"].get("language")
        o["text"] = n["content"]["text"]
    if k == "section":
        o["level"] = attrs["level"]
        heads = [c for c in doc.kids(nid) if doc.nodes[c]["kind"] == "heading"]
        o["heading"] = text_hash_of(heads[0]) if heads else None
    if k == "table":
        o["columns"] = attrs.get("columns", [])
    if k == "cell":
        o["column"] = attrs.get("column")
        o["rowspan"] = attrs.get("rowspan", 1)
        o["colspan"] = attrs.get("colspan", 1)
        o["role"] = attrs.get("role", "data")
    if k == "theorem":
        o["variant"] = attrs.get("variant")
    return o


class _Entry:
    """One node's cached work: its five values, and the projections its parent joins."""

    __slots__ = ("sig", "values", "plain", "joined_plain", "joined_strict", "complexes")

    def __init__(self, sig, values, plain, joined_plain, joined_strict, complexes):
        self.sig = sig
        self.values = values
        self.plain = plain
        # What this node gives a container's projections: its text leaves joined, or None when it has none.
        self.joined_plain = joined_plain
        self.joined_strict = joined_strict
        self.complexes = complexes


def _sig(doc, nid):
    """Everything a node's values depend on besides its children's values."""
    n = doc.nodes[nid]
    return (n["kind"], n.get("label"), n.get("attrs"), n.get("content"), tuple(doc.kids(nid)))


class Index:
    """The five values of every node of one document, kept so that a change rehashes only what it touched.

    A node's values depend only on its own subtree, so after a change only the nodes whose own record or
    children changed, and their ancestors, are rehashed. A container's projections are its children's
    cached projections joined, so no container walks its subtree again.
    """

    def __init__(self, hasher=PLAIN):
        self.hasher = hasher
        self.entries = {}
        self.values = {}

    @classmethod
    def build(cls, doc, hasher=PLAIN):
        ix = cls(hasher)
        ix._rehash(doc, set(doc.nodes))
        return ix

    def update(self, doc):
        """A new Index for `doc`, a later state of this one's document; this Index is left unchanged."""
        ix = Index(self.hasher)
        ix.entries = {nid: e for nid, e in self.entries.items() if nid in doc.nodes}
        ix.values = {nid: v for nid, v in self.values.items() if nid in doc.nodes}
        dirty = {nid for nid in doc.nodes if nid not in ix.entries or ix.entries[nid].sig != _sig(doc, nid)}
        for nid in list(dirty):
            p = doc.nodes[nid].get("parent")
            while p is not None and p not in dirty:
                dirty.add(p)
                p = doc.nodes[p].get("parent")
        ix._rehash(doc, dirty)
        return ix

    def plain(self, nid):
        """The node's plain projection."""
        return self.entries[nid].plain

    def _rehash(self, doc, dirty):
        h = self.hasher
        for nid in reversed(doc.preorder()):
            if nid not in dirty:
                continue
            n = doc.nodes[nid]
            k = n["kind"]
            kids = doc.kids(nid)
            if k in RUN_KINDS or k in MATH_KINDS or k in VERBATIM_KINDS:
                p = plain(doc, nid)
                st, cx = strict(doc, nid)
                # An equation line is never a text leaf; a sentence's equations are part of its own runs.
                joined = (None, None, []) if k == "equation-line" else (p, st, cx)
            else:
                ps, ss, cx = [], [], []
                for c in kids:
                    e = self.entries[c]
                    if e.joined_plain is not None:
                        ps.append(e.joined_plain)
                        ss.append(e.joined_strict)
                        cx.extend(e.complexes)
                p = " ".join(ps)
                st = _finish_strict(" ".join(ss))
                joined = (p, " ".join(ss), cx) if ps else (None, None, [])
            own = h(jcs.dumpb(own_content(doc, nid, lambda x: self.values[x]["text_hash"])))
            self.values[nid] = {
                "text_hash": h(norm(p)),
                "match_hash": h(st),
                "align_key": h(lossy_from_strict(st)),
                "own_hash": own,
                "node_hash": h(jcs.dumpb({"own": own, "children": [self.values[c]["node_hash"] for c in kids]})),
            }
            self.entries[nid] = _Entry(_sig(doc, nid), self.values[nid], p, *joined)


def compute(doc, hasher=PLAIN):
    """The five content values of every node, bottom up. Returns {id: {...}}."""
    return Index.build(doc, hasher).values


def meta_object(document):
    return {k: document.get(k) for k in ("work", "version", "relations", "meta", "language", "license")}


def meta_hash(document, hasher=PLAIN):
    return hasher(jcs.dumpb(meta_object(document)))