"""Establishing a document from a nested tree, minting IDs, and structural checks.

Non-normative; see spec/03-nodes.md and spec/04-identity.md.
"""

import copy

from . import hashing
from .model import (
    AT_MOST_ONCE, CHILDREN, KINDS, LABELLED, LEAVES, MATH_KINDS, NONEMPTY, RUN_KINDS, RUNS, STYLES,
    VERBATIM_KINDS, Doc,
)
from .text import NODE_ID_RE, mint_node_id, norm


class Refusal(Exception):
    """A rule was broken. `rule` names it, as spec/22-conformance.md lists."""

    def __init__(self, rule, message, node=None):
        super().__init__("%s: %s" % (rule, message))
        self.rule = rule
        self.message = message
        self.node = node


def canonical_runs(runs):
    """Drop empty text runs and join adjacent text runs with the same styles."""
    out = []
    for r in runs:
        r = dict(r)
        if r["kind"] == "text":
            if "style" in r:
                r["style"] = sorted(set(r["style"]))
                if not r["style"]:
                    del r["style"]
            if r["text"] == "":
                continue
            if out and out[-1]["kind"] == "text" and out[-1].get("style") == r.get("style"):
                out[-1] = dict(out[-1], text=out[-1]["text"] + r["text"])
                continue
        out.append(r)
    return out


def _flatten(spec, parent, ord_, out):
    node = {k: v for k, v in spec.items() if k != "children"}
    node["_parent"] = parent
    node["ord"] = ord_
    out.append(node)
    idx = len(out) - 1
    for i, c in enumerate(spec.get("children", [])):
        _flatten(c, idx, i, out)
    return idx


def mint_ids(doc_key, flat, used, preassigned=None):
    """Mint IDs for flat nodes in document order. `flat` nodes carry `_parent` indices."""
    preassigned = preassigned or {}
    tmp = []
    for i, n in enumerate(flat):
        m = {k: v for k, v in n.items() if not k.startswith("_")}
        m["id"] = "tmp%d" % i
        m["parent"] = None if n["_parent"] is None else "tmp%d" % n["_parent"]
        tmp.append(m)
    _resolve_child_refs(tmp)
    doc = Doc(doc_key, tmp)
    seen = {}
    ids = {}
    for i, n in enumerate(tmp):
        if i in preassigned:
            ids[n["id"]] = preassigned[i]
            used.add(preassigned[i])
            continue
        th = hashing.PLAIN(norm(hashing.plain(doc, n["id"])))
        k = (n["kind"], th)
        occ = seen.get(k, 0)
        seen[k] = occ + 1
        nid = mint_node_id(doc_key, n["kind"], th, occ)
        while nid in used:
            occ += 1
            nid = mint_node_id(doc_key, n["kind"], th, occ)
        used.add(nid)
        ids[n["id"]] = nid
    result = []
    for n in tmp:
        m = dict(n)
        m["id"] = ids[n["id"]]
        m["parent"] = None if n["parent"] is None else ids[n["parent"]]
        if m["kind"] in RUN_KINDS:
            m["content"] = [
                dict(r, node=ids[r["node"]]) if r["kind"] == "equation" else r for r in m["content"]
            ]
        result.append(m)
    return result


def _resolve_child_refs(tmp):
    """Equation runs in a spec name their child by position: {"kind": "equation", "child": i}."""
    by_parent = {}
    for n in tmp:
        by_parent.setdefault(n["parent"], []).append(n)
    for n in tmp:
        if n["kind"] in RUN_KINDS:
            kids = sorted(by_parent.get(n["id"], []), key=lambda c: c["ord"])
            runs = []
            for r in n["content"]:
                if r["kind"] == "equation" and "child" in r:
                    r = {"kind": "equation", "node": kids[r["child"]]["id"]}
                runs.append(r)
            n["content"] = canonical_runs(runs)


def establish(doc_key, spec, used=None):
    """Turn a nested spec into flat node records with minted IDs, in document order."""
    flat = []
    _flatten(copy.deepcopy(spec), None, 0, flat)
    return mint_ids(doc_key, flat, used if used is not None else set())


def with_values(doc_key, nodes, hasher=hashing.PLAIN):
    doc = Doc(doc_key, nodes)
    vals = hashing.compute(doc, hasher)
    out = []
    for nid in doc.preorder():
        n = dict(doc.nodes[nid])
        for k in ("text_hash", "match_hash", "align_key", "own_hash", "node_hash"):
            n.pop(k, None)
        n.update(vals[nid])
        out.append(n)
    return out


def validate_structure(doc):
    """Check the tree against the child rules of spec/03-nodes.md."""
    if doc.root is None:
        raise Refusal("doc.tree.root", "no root node")
    roots = [x for x in doc.order if doc.nodes[x].get("parent") is None]
    if len(roots) != 1 or doc.nodes[doc.root]["kind"] != "document":
        raise Refusal("doc.tree.root", "exactly one root, of kind document, is required")
    if doc.preorder() != doc.order:
        raise Refusal("doc.tree.order", "nodes are not in document order")
    seen = set()
    for nid in doc.order:
        n = doc.nodes[nid]
        if nid in seen:
            raise Refusal("doc.id.unique", "duplicate ID " + nid, nid)
        seen.add(nid)
        if not NODE_ID_RE.match(nid):
            raise Refusal("doc.id.form", "ID is not seven letter-first base32 characters", nid)
        k = n["kind"]
        if k not in KINDS:
            raise Refusal("doc.node.kind", "unknown kind " + k, nid)
        kids = doc.kids(nid)
        if [doc.nodes[c]["ord"] for c in kids] != list(range(len(kids))):
            raise Refusal("doc.tree.ord", "child order is not 0..n-1", nid)
        counts = {}
        for i, c in enumerate(kids):
            ck = doc.nodes[c]["kind"]
            if ck not in CHILDREN[k]:
                raise Refusal("doc.tree.child", "%s may not contain %s" % (k, ck), c)
            counts[ck] = counts.get(ck, 0) + 1
            if ck == "heading" and i != 0:
                raise Refusal("doc.tree.child", "a heading must be a section's first child", c)
        for ck, cnt in counts.items():
            if cnt > 1 and (ck in AT_MOST_ONCE or ck == "heading"):
                raise Refusal("doc.tree.child", "%s may occur once in %s" % (ck, k), nid)
        if k == "table" and counts.get("table-body", 0) != 1:
            raise Refusal("doc.tree.child", "a table has exactly one table-body", nid)
        if k in NONEMPTY and not kids:
            raise Refusal("doc.tree.empty", "%s must have children" % k, nid)
        if n.get("label") is not None and k not in LABELLED:
            raise Refusal("doc.node.label", "%s does not take a label" % k, nid)
        _check_content(doc, nid)


def _check_content(doc, nid):
    n = doc.nodes[nid]
    k = n["kind"]
    if n.get("withheld"):
        if "content" in n:
            raise Refusal("doc.policy.withheld", "a withheld node keeps no content", nid)
        return
    if k in RUN_KINDS:
        runs = n.get("content")
        if not isinstance(runs, list) or not runs:
            raise Refusal("doc.node.content", "runs are required", nid)
        if canonical_runs(runs) != runs:
            raise Refusal("doc.run.canonical", "runs are not in canonical form", nid)
        eq_refs = []
        for r in runs:
            if r["kind"] not in RUNS:
                raise Refusal("doc.run.kind", "unknown run kind " + r["kind"], nid)
            st = r.get("style", [])
            if any(s not in STYLES for s in st) or ("sup" in st and "sub" in st):
                raise Refusal("doc.run.style", "bad style set", nid)
            if r["kind"] == "equation":
                eq_refs.append(r["node"])
        kids = doc.kids(nid)
        if eq_refs != kids:
            raise Refusal("doc.run.equation", "equation runs must name the sentence's children in order", nid)
    elif k in MATH_KINDS:
        c = n.get("content") or {}
        if c.get("notation") not in ("tex", "typst") or not isinstance(c.get("source"), str):
            raise Refusal("doc.node.content", "equations need a notation and a source", nid)
    elif k in VERBATIM_KINDS:
        c = n.get("content") or {}
        if not isinstance(c.get("text"), str):
            raise Refusal("doc.node.content", "verbatim nodes need text", nid)
    elif "content" in n:
        raise Refusal("doc.node.content", "%s holds no content" % k, nid)
    if k == "section" and not isinstance(n.get("attrs", {}).get("level"), int):
        raise Refusal("doc.node.attrs", "a section needs a level", nid)


def leaf_kinds():
    return LEAVES