"""Corrections on an established document.

Non-normative; see spec/11-corrections.md.
"""

import copy

from . import hashing
from .model import LABELLED, LEAVES, MATH_KINDS, RUN_KINDS, VERBATIM_KINDS, Doc, CHILDREN
from .tree import Refusal, canonical_runs, mint_ids, validate_structure, _flatten
from . import jcs
from .text import digits, record_id, sha256_hex

MERGEABLE = RUN_KINDS | {"paragraph", "list", "item", "quote", "caption", "figure-note", "table-note"}
SETTABLE = {
    "section": {"level"},
    "table": {"columns"},
    "cell": {"column", "rowspan", "colspan", "role"},
    "row": {"header_path"},
    "list": {"ordered"},
    "item": {"marker"},
    "theorem": {"variant"},
}
EDITABLE_RUN_FIELD = {"text": "text", "link": "text", "cite": "text", "xref": "text", "fnmark": "text", "math": "source"}


class State:
    """The current state of a document: nodes, the ID table and the corrections applied."""

    def __init__(self, key, nodes, ids=None, corrections=None):
        self.key = key
        self.nodes = [dict(n) for n in nodes]
        self.ids = [dict(e) for e in (ids or [])]
        self.corrections = [dict(c) for c in (corrections or [])]

    @property
    def rev(self):
        return len(self.corrections)

    def doc(self):
        return Doc(self.key, self.nodes)

    def used_ids(self):
        return {n["id"] for n in self.nodes} | {e["id"] for e in self.ids}


def _strings(node):
    """Editable strings of a leaf: (path, value) pairs."""
    k = node["kind"]
    if k in RUN_KINDS:
        out = []
        for i, r in enumerate(node["content"]):
            f = EDITABLE_RUN_FIELD.get(r["kind"])
            if f:
                out.append(((i, f), r[f]))
        return out
    if k in MATH_KINDS:
        return [(("source",), node["content"]["source"])]
    if k in VERBATIM_KINDS:
        return [(("text",), node["content"]["text"])]
    return []


def _occurrences(s, find):
    return [i for i in range(len(s) - len(find) + 1) if s.startswith(find, i)]


def _set_string(node, path, value):
    if len(path) == 2:
        i, f = path
        node["content"][i] = dict(node["content"][i], **{f: value})
    else:
        node["content"] = dict(node["content"], **{path[0]: value})


def _has_crop(c):
    return any(e.get("kind") == "crop" for e in c.get("evidence", []))


def _rebuild(nodes):
    """Recompute document order and `ord` from parent links, keeping sibling order as listed."""
    by_id = {n["id"]: n for n in nodes}
    kids = {}
    root = None
    for n in nodes:
        if n.get("parent") is None:
            root = n["id"]
        else:
            kids.setdefault(n["parent"], []).append(n["id"])
    for p in kids:
        kids[p].sort(key=lambda c: by_id[c]["ord"])
    out = []
    stack = [root]
    while stack:
        x = stack.pop()
        out.append(by_id[x])
        for i, c in enumerate(kids.get(x, [])):
            by_id[c]["ord"] = i
        stack.extend(reversed(kids.get(x, [])))
    return out


def _subtree(doc, nid):
    return doc.preorder(nid)


def _place(nodes, nid, parent, after):
    """Give `nid` a fractional position under `parent` after sibling `after` (None: first)."""
    by_id = {n["id"]: n for n in nodes}
    sibs = sorted([n for n in nodes if n.get("parent") == parent and n["id"] != nid], key=lambda n: n["ord"])
    for i, s in enumerate(sibs):
        s["ord"] = 2 * i + 2
    if after is None:
        pos = 0
    else:
        if after not in by_id or by_id[after].get("parent") != parent:
            raise Refusal("correction.target", "`place_after` is not a child of the parent")
        pos = by_id[after]["ord"] + 1
    by_id[nid]["parent"] = parent
    by_id[nid]["ord"] = pos


def apply(state, c, hasher=hashing.PLAIN):
    """Apply correction `c` to `state` and return the recorded correction. Raises Refusal."""
    c = copy.deepcopy(c)
    doc = state.doc()
    before_vals = hashing.compute(doc, hasher)
    nodes = [copy.deepcopy(n) for n in state.nodes]
    by_id = {n["id"]: n for n in nodes}
    op = c["op"]
    target = c.get("target")
    if op != "insert" or "target" in c:
        if target not in by_id:
            raise Refusal("correction.target", "no live node %s" % target)
    if target == doc.root and op in ("delete", "move", "merge", "reclassify"):
        raise Refusal("correction.target", "the root cannot be " + op + "d")
    record = {"format": "scienting-correction/1", "index": state.rev + 1}
    retired, created = [], []
    path = doc.ancestors(target) if target in by_id else None
    touched = []

    if op == "replace":
        node = by_id[target]
        find, repl = c.get("find", ""), c["replace"]
        if not find:
            raise Refusal("correction.find", "find must not be empty")
        hits = []
        for pth, s in _strings(node):
            if "run" in c and (len(pth) != 2 or pth[0] != c["run"]):
                continue
            hits += [(pth, s, i) for i in _occurrences(s, find)]
        if len(hits) != 1:
            raise Refusal("correction.find", "find occurs %d times, not exactly once" % len(hits), target)
        pth, s, i = hits[0]
        if isinstance(repl, str):
            _set_string(node, pth, s[:i] + repl + s[i + len(find):])
        else:
            if len(pth) != 2 or node["content"][pth[0]]["kind"] != "text":
                raise Refusal("correction.replace", "runs may replace text only inside a text run", target)
            if any(r["kind"] == "equation" for r in repl):
                raise Refusal("correction.replace", "equation runs are added by insert", target)
            run = node["content"][pth[0]]
            style = {"style": run["style"]} if "style" in run else {}
            mid = [dict({"kind": "text", "text": s[:i]}, **style)] + list(repl) + [
                dict({"kind": "text", "text": s[i + len(find):]}, **style)
            ]
            node["content"] = node["content"][: pth[0]] + mid + node["content"][pth[0] + 1:]
        if node["kind"] in RUN_KINDS:
            node["content"] = canonical_runs(node["content"])
        touched = [target]

    elif op == "insert" and "node" in c:
        if not _has_crop(c):
            raise Refusal("correction.evidence", "an insert needs crop evidence")
        parent = c["parent"]
        if parent not in by_id:
            raise Refusal("correction.target", "no live parent %s" % parent)
        flat = []
        _flatten(copy.deepcopy(c["node"]), None, 0, flat)
        new = mint_ids(state.key, flat, state.used_ids())
        new[0]["parent"] = parent
        nodes.extend(new)
        _place(nodes, new[0]["id"], parent, c.get("place_after"))
        nodes = _rebuild(nodes)
        created = [n["id"] for n in new]
        tmp = Doc(state.key, nodes)
        touched = [x for x in tmp.leaves(new[0]["id"])]
        path = tmp.ancestors(new[0]["id"])
        target = new[0]["id"]

    elif op == "insert":
        if not _has_crop(c):
            raise Refusal("correction.evidence", "an insert needs crop evidence")
        node = by_id[target]
        if node["kind"] not in RUN_KINDS:
            raise Refusal("correction.insert", "text is inserted into runs", target)
        r = node["content"][c["run"]]
        f = EDITABLE_RUN_FIELD.get(r["kind"])
        if not f or not (0 <= c["offset"] <= len(r[f])):
            raise Refusal("correction.insert", "bad run or offset", target)
        node["content"][c["run"]] = dict(r, **{f: r[f][: c["offset"]] + c["text"] + r[f][c["offset"]:]})
        node["content"] = canonical_runs(node["content"])
        touched = [target]

    elif op == "delete":
        gone = _subtree(doc, target)
        touched = [x for x in gone if doc.nodes[x]["kind"] in LEAVES]
        nodes = _rebuild([n for n in nodes if n["id"] not in set(gone)])
        retired = [(x, None) for x in gone]

    elif op == "split":
        node = by_id[target]
        if node["kind"] != "sentence":
            raise Refusal("correction.split", "only sentences are split", target)
        k = c["offset"]
        first, second, pos = [], [], 0
        total = len(hashing.plain_runs(node["content"], doc))
        if not 0 < k < total:
            raise Refusal("correction.split", "offset must fall inside the sentence", target)
        for r in node["content"]:
            ln = len(hashing.plain_runs([r], doc))
            if pos + ln <= k:
                first.append(r)
            elif pos >= k:
                second.append(r)
            elif r["kind"] == "text":
                cut = k - pos
                first.append(dict(r, text=r["text"][:cut]))
                second.append(dict(r, text=r["text"][cut:]))
            else:
                raise Refusal("correction.split", "the offset falls inside a %s run" % r["kind"], target)
            pos += ln
        _strip_edge(first, end=True)
        _strip_edge(second, end=False)
        first, second = canonical_runs(first), canonical_runs(second)
        if not first or not second:
            raise Refusal("correction.split", "both parts must keep text", target)
        new_spec = {"kind": "sentence", "content": second}
        eq_second = [r["node"] for r in second if r["kind"] == "equation"]
        flat = [dict(new_spec, _parent=None, ord=0)]
        tmp_nodes = mint_ids(state.key, [dict(flat[0], content=[
            {"kind": "text", "text": hashing.plain_runs(second, doc)}])], state.used_ids())
        new_id = tmp_nodes[0]["id"]
        node["content"] = first
        new_node = {"id": new_id, "kind": "sentence", "parent": node["parent"], "ord": 0, "content": second}
        if "fidelity" in node:
            new_node["fidelity"] = node["fidelity"]
        nodes.append(new_node)
        _place(nodes, new_id, node["parent"], target)
        for i, e in enumerate(eq_second):
            by_id[e]["parent"] = new_id
            by_id[e]["ord"] = i
        nodes = _rebuild(nodes)
        created = [new_id]
        touched = [target, new_id]

    elif op == "merge":
        other = c["with"]
        node = by_id[target]
        if other not in by_id:
            raise Refusal("correction.target", "no live node %s" % other)
        o = by_id[other]
        if node["kind"] != o["kind"] or node["kind"] not in MERGEABLE:
            raise Refusal("correction.merge", "merge joins two nodes of one mergeable kind", target)
        if o.get("parent") != node.get("parent") or o["ord"] != node["ord"] + 1:
            raise Refusal("correction.merge", "merge joins a node with its next sibling", target)
        if node["kind"] in RUN_KINDS:
            joiner = c.get("joiner", " ")
            mid = [{"kind": "text", "text": joiner}] if joiner else []
            node["content"] = canonical_runs(node["content"] + mid + o["content"])
            touched = [target, other]
        else:
            touched = [x for x in doc.leaves(other) if x != other]
        base = len(doc.kids(target))
        for i, kid in enumerate(doc.kids(other)):
            by_id[kid]["parent"] = target
            by_id[kid]["ord"] = base + i
        nodes = _rebuild([n for n in nodes if n["id"] != other])
        retired = [(other, target)]

    elif op == "move":
        parent = c["parent"]
        if parent not in by_id or parent in set(_subtree(doc, target)):
            raise Refusal("correction.move", "bad destination", target)
        _place(nodes, target, parent, c.get("place_after"))
        nodes = _rebuild(nodes)
        touched = [x for x in doc.leaves(target)]
        record["to_path"] = Doc(state.key, nodes).ancestors(target)

    elif op == "reclassify":
        node = by_id[target]
        new_kind = c["kind"]
        old = node["kind"]
        same_shape = (old in RUN_KINDS) == (new_kind in RUN_KINDS) and (old in MATH_KINDS) == (
            new_kind in MATH_KINDS) and (old in VERBATIM_KINDS) == (new_kind in VERBATIM_KINDS)
        if new_kind not in CHILDREN or not same_shape:
            raise Refusal("correction.reclassify", "%s cannot become %s" % (old, new_kind), target)
        node["kind"] = new_kind
        if new_kind not in LABELLED:
            node.pop("label", None)
        record["from"] = old
        touched = [x for x in doc.leaves(target)]

    elif op == "relabel":
        node = by_id[target]
        if node["kind"] not in LABELLED:
            raise Refusal("correction.relabel", "%s takes no label" % node["kind"], target)
        record["from"] = node.get("label")
        node["label"] = c["label"]
        touched = [target] if node["kind"] in LEAVES else []

    elif op == "set":
        node = by_id[target]
        attr = c["attr"]
        if attr not in SETTABLE.get(node["kind"], ()):
            raise Refusal("correction.set", "%s cannot be set on %s" % (attr, node["kind"]), target)
        attrs = dict(node.get("attrs", {}))
        record["from"] = attrs.get(attr)
        attrs[attr] = c["value"]
        node["attrs"] = attrs
        touched = [target] if node["kind"] in LEAVES else []

    else:
        raise Refusal("correction.op", "unknown operation " + op)

    after_doc = Doc(state.key, nodes)
    validate_structure(after_doc)
    after_vals = hashing.compute(after_doc, hasher)

    # Digits: a change to a digit needs crop evidence, and no leaf outside `touched` may change its digits.
    needs_crop = False
    if op == "replace":
        rep = c["replace"] if isinstance(c["replace"], str) else hashing.plain_runs(c["replace"], doc)
        needs_crop = digits(c["find"]) != digits(rep)
    elif op == "delete":
        needs_crop = any(digits(hashing.plain(doc, x)) for x in touched)
    elif op == "relabel":
        needs_crop = digits(record.get("from") or "") != digits(c.get("label") or "")
    if needs_crop and not _has_crop(c):
        raise Refusal("correction.digits", "a digit changed without crop evidence", target)
    for nid in doc.leaves(doc.root):
        if nid in after_doc.nodes and nid not in touched:
            if digits(hashing.plain(doc, nid)) != digits(hashing.plain(after_doc, nid)):
                raise Refusal("correction.digits", "an undeclared digit change", nid)

    for k in ("op", "target", "find", "replace", "run", "offset", "text", "parent", "place_after", "node", "with",
              "joiner", "kind", "label", "attr", "value", "evidence", "author", "reason", "at"):
        if k in c:
            record[k] = c[k]
    if op == "insert" and "node" in c:
        record["target"] = target
    record["before"] = before_vals[target]["node_hash"] if target in before_vals else None
    record["after"] = after_vals[target]["node_hash"] if target in after_vals else None
    record["path"] = path
    record["touched"] = sorted(set(touched))
    if created:
        record["created"] = created
    if retired:
        record["retired"] = [r[0] for r in retired]
    record["id"] = record_id(sha256_hex(jcs.dumpb({k: v for k, v in record.items() if k != "id"})))

    old_kinds = {n["id"]: n["kind"] for n in state.nodes}
    new_ids = list(state.ids)
    for rid, succ in retired:
        new_ids.append({"id": rid, "kind": "retired", "successor": succ, "node_kind": old_kinds[rid],
                        "rev": record["index"]})
    new_state = State(state.key, [after_doc.nodes[x] for x in after_doc.preorder()], new_ids,
                      state.corrections + [record])
    return new_state, record


def _strip_edge(runs, end):
    from .text import is_ws
    if not runs:
        return
    i = len(runs) - 1 if end else 0
    r = runs[i]
    if r["kind"] != "text":
        return
    t = r["text"]
    if end:
        while t and is_ws(t[-1]):
            t = t[:-1]
    else:
        while t and is_ws(t[0]):
            t = t[1:]
    runs[i] = dict(r, text=t)