"""Packages: writing, reading and validating a scienting-doc/1 package folder.

Non-normative; see spec/12-package.md and spec/22-conformance.md.
"""

import base64
import json
import os
import posixpath
import re

from . import crypto, hashing, jcs, receipts as R
from .anchors import ADDRESS_RE
from .model import LEAVES, MATH_KINDS, RUN_KINDS, VERBATIM_KINDS, Doc
from .schema import validate as schema_validate
from .text import letter_first, record_id, sha256_hex
from .tree import Refusal, validate_structure

JSONL = {
    "nodes.jsonl": "node.schema.json",
    "pages.jsonl": "page.schema.json",
    "locations.jsonl": "location.schema.json",
    "references.jsonl": "reference.schema.json",
    "corrections.jsonl": "correction.schema.json",
    "receipts.jsonl": None,
    "ids.jsonl": "id-entry.schema.json",
    "suspects.jsonl": "suspect.schema.json",
}
JSON = {
    "document.json": "document.schema.json",
    "witnesses.json": "witnesses.schema.json",
    "provenance.json": "provenance.schema.json",
    "roster.json": "roster.schema.json",
}
REQUIRED = ["document.json", "nodes.jsonl", "witnesses.json", "provenance.json"]


def line(rec):
    return jcs.dumps(rec) + "\n"


def jsonl(recs):
    return "".join(line(r) for r in recs).encode("utf-8")


def witness_id(sha):
    return letter_first(bytes.fromhex(sha), 7)


SHA256_RE = re.compile(r"[0-9a-f]{64}")


def safe_path(p):
    if not isinstance(p, str) or not p or p.startswith("/") or "\\" in p or "\x00" in p:
        return False
    parts = p.split("/")
    return all(x not in ("", ".", "..") for x in parts) and posixpath.normpath(p) == p


def write(root, files, absent=()):
    """Write `files` ({path: bytes}) and a MANIFEST.json listing them and the `absent` entries.

    Returns the package hash.
    """
    for path in list(files) + [a["path"] for a in absent]:
        if not safe_path(path) or path == "MANIFEST.json":
            raise Refusal("package.path.unsafe", "unsafe path %r" % path)
    os.makedirs(root, exist_ok=True)
    entries = [dict(a) for a in absent]
    for path in sorted(files, key=lambda p: p.encode("utf-8")):
        data = files[path]
        if data is None:
            continue
        full = os.path.join(root, path)
        os.makedirs(os.path.dirname(full), exist_ok=True)
        with open(full, "wb") as f:
            f.write(data)
        entries.append({"path": path, "sha256": sha256_hex(data), "size": len(data)})
    entries.sort(key=lambda e: e["path"].encode("utf-8"))
    doc = json.loads(files["document.json"])
    manifest = {"format": "scienting-package/1", "doc": "scienting-doc/1", "key": doc["key"], "files": entries}
    mb = (jcs.dumps(manifest) + "\n").encode("utf-8")
    with open(os.path.join(root, "MANIFEST.json"), "wb") as f:
        f.write(mb)
    return sha256_hex(mb)


def read_jsonl(path):
    out = []
    with open(path, encoding="utf-8") as f:
        for i, ln in enumerate(f, 1):
            if not ln.endswith("\n"):
                raise Refusal("package.file.newline", "%s:%d: every line ends with a newline" % (path, i))
            rec = json.loads(ln)
            if line(rec) != ln:
                raise Refusal("package.file.canonical", "%s:%d: line is not canonical JSON" % (path, i))
            out.append(rec)
    return out


def read_json(path):
    with open(path, encoding="utf-8") as f:
        text = f.read()
    rec = json.loads(text)
    if line(rec) != text:
        raise Refusal("package.file.canonical", "%s: file is not canonical JSON" % path)
    return rec


class Report:
    def __init__(self):
        self.errors = []

    def add(self, rule, where, message):
        self.errors.append({"rule": rule, "where": where, "message": message})

    def __bool__(self):
        return not self.errors


def _all_files(root):
    out = []
    for d, _, fs in os.walk(root):
        for f in fs:
            out.append(os.path.relpath(os.path.join(d, f), root).replace(os.sep, "/"))
    return sorted(out)


def validate(root, store, key=None):
    """Validate the package at `root`. `store` is a schema store; `key` a hashing key for keyed packages."""
    rep = Report()
    try:
        _validate(root, store, key, rep)
    except Refusal as e:
        rep.add(e.rule, e.node or "", e.message)
    return rep


def _validate(root, store, key, rep):
    mpath = os.path.join(root, "MANIFEST.json")
    if not os.path.exists(mpath):
        raise Refusal("package.manifest.missing", "MANIFEST.json is missing")
    manifest = read_json(mpath)
    for e in schema_validate(store, "manifest.schema.json", manifest):
        rep.add("package.manifest.schema", "MANIFEST.json", e)
    listed = {}
    paths = [f["path"] for f in manifest.get("files", [])]
    if paths != sorted(paths, key=lambda p: p.encode("utf-8")) or len(set(paths)) != len(paths):
        rep.add("package.manifest.order", "MANIFEST.json", "files must be sorted by path and unique")
    for f in manifest.get("files", []):
        p = f["path"]
        listed[p] = f
        if not safe_path(p):
            rep.add("package.path.unsafe", p, "unsafe path")
            continue
        full = os.path.join(root, p)
        if f.get("present", True) is False:
            if not p.startswith("assets/"):
                rep.add("package.manifest.absent", p, "only assets may be listed without their bytes")
            if os.path.exists(full):
                rep.add("package.manifest.absent", p, "a file listed as absent is present")
            continue
        if not os.path.exists(full):
            rep.add("package.manifest.missing-file", p, "listed file is missing")
            continue
        with open(full, "rb") as fh:
            data = fh.read()
        if sha256_hex(data) != f["sha256"] or len(data) != f["size"]:
            rep.add("package.manifest.hash", p, "SHA-256 or size does not match the manifest")
    for p in _all_files(root):
        if p != "MANIFEST.json" and p not in listed:
            rep.add("package.manifest.unlisted", p, "file is not listed in the manifest")
    for p in REQUIRED:
        if p not in listed:
            rep.add("package.file.required", p, "required file is missing")
    if not rep:
        return

    data = {}
    for p, sch in JSON.items():
        if p in listed:
            data[p] = read_json(os.path.join(root, p))
            for e in schema_validate(store, sch, data[p]):
                rep.add("package.schema", p, e)
    for p, sch in JSONL.items():
        if p in listed:
            data[p] = read_jsonl(os.path.join(root, p))
            for i, rec in enumerate(data[p], 1):
                s = sch
                if p == "receipts.jsonl":
                    s = "receipt.schema.json" if rec.get("format") == "scienting-receipt/1" else "revocation.schema.json"
                for e in schema_validate(store, s, rec):
                    rep.add("package.schema", "%s:%d" % (p, i), e)
            if p != "nodes.jsonl" and not data[p]:
                rep.add("package.file.empty", p, "an optional JSONL file with no records is omitted")
        else:
            data[p] = []
    if not rep:
        return

    document = data["document.json"]
    if document["key"] != manifest["key"]:
        rep.add("package.manifest.key", "MANIFEST.json", "key differs from document.json")
    nodes = data["nodes.jsonl"]
    doc = Doc(document["key"], nodes)
    try:
        validate_structure(doc)
    except Refusal as e:
        rep.add(e.rule, "nodes.jsonl " + (e.node or ""), e.message)
        return

    # Policy and hashing.
    withheld = [n["id"] for n in nodes if n.get("withheld")]
    scheme = document["hashing"]["scheme"]
    if document["policy"] == "full" and withheld:
        rep.add("doc.policy.withheld", "nodes.jsonl", "a full package withholds no text")
    if document["policy"] == "checks-only":
        bearing = [n["id"] for n in nodes if n["kind"] in LEAVES]
        if set(bearing) != set(withheld):
            rep.add("doc.policy.withheld", "nodes.jsonl", "a checks-only package withholds every leaf's content")
        if scheme == "plain":
            rep.add("doc.policy.keyed", "document.json", "a checks-only package carries keyed hashes")
    if scheme == "hmac-team" and key is not None and document["hashing"]["key_id"] != crypto.key_id(key):
        rep.add("doc.hash.key", "document.json", "key_id does not match the key supplied")

    # References inside runs.
    kinds = {n["id"]: n["kind"] for n in nodes}
    for n in nodes:
        for r in n.get("content", []) if n["kind"] in RUN_KINDS else []:
            if r["kind"] == "cite":
                for ref in r["refs"]:
                    if kinds.get(ref) != "reference":
                        rep.add("doc.run.cite", n["id"], "cite names %s, which is not a reference" % ref)
            if r["kind"] == "fnmark" and r.get("footnote") is not None and kinds.get(r["footnote"]) != "footnote":
                rep.add("doc.run.fnmark", n["id"], "fnmark names a node that is not a footnote")
            if r["kind"] == "xref" and r.get("target") and "#" not in r["target"]:
                if r["target"] not in kinds:
                    rep.add("doc.run.xref", n["id"], "xref names no node in this document")
        if n["kind"] == "table":
            for col in n.get("attrs", {}).get("columns", []):
                for h in col["header_path"]:
                    if kinds.get(h) != "cell":
                        rep.add("doc.table.header", n["id"], "a header path names %s, not a cell" % h)
        if n["kind"] == "row":
            for h in n.get("attrs", {}).get("header_path", []):
                if kinds.get(h) != "cell":
                    rep.add("doc.table.header", n["id"], "a row header path names %s, not a cell" % h)

    # Hashes.
    hasher = hashing.PLAIN if scheme == "plain" else (hashing.Hasher(key) if key else None)
    if hasher is not None and not withheld:
        vals = hashing.compute(doc, hasher)
        for n in nodes:
            for k, v in vals[n["id"]].items():
                if n[k] != v:
                    rep.add("doc.hash." + k.replace("_", "-"), n["id"], "%s does not recompute" % k)
        if vals[doc.root]["node_hash"] != document["root"]:
            rep.add("doc.hash.root", "document.json", "root does not equal the root node's node_hash")
        if hashing.meta_hash(document, hasher) != document["meta_hash"]:
            rep.add("doc.hash.meta", "document.json", "meta_hash does not recompute")
    else:
        vals = {n["id"]: {k: n[k] for k in ("text_hash", "match_hash", "align_key", "own_hash", "node_hash")}
                for n in nodes}
        if n_root_hash(nodes, doc) != document["root"]:
            rep.add("doc.hash.root", "document.json", "root does not equal the root node's node_hash")

    # IDs.
    live = set(kinds)
    ids = data["ids.jsonl"]
    known = set(live)
    for i, e in enumerate(ids, 1):
        if e["id"] in live:
            rep.add("doc.id.retired-live", "ids.jsonl:%d" % i, "%s is both live and in the ID table" % e["id"])
        known.add(e["id"])
    for i, e in enumerate(ids, 1):
        s = e.get("successor")
        if s is not None and s not in known:
            rep.add("doc.id.successor", "ids.jsonl:%d" % i, "successor %s is unknown" % s)
        if e["kind"] == "alias" and s not in live:
            rep.add("doc.id.alias", "ids.jsonl:%d" % i, "an alias names a live node")
    order = sorted(ids, key=lambda e: (0 if e["kind"] == "alias" else 1, e.get("rev", 0), e["id"]))
    if order != ids:
        rep.add("doc.order.ids", "ids.jsonl", "aliases first by ID, then retirements by revision and ID")

    # Corrections.
    cors = data["corrections.jsonl"]
    for i, c in enumerate(cors, 1):
        if c["index"] != i:
            rep.add("doc.correction.index", "corrections.jsonl:%d" % i, "indexes count from 1 without gaps")
        if c["id"] != record_id(sha256_hex(jcs.dumpb({k: v for k, v in c.items() if k != "id"}))):
            rep.add("doc.correction.id", "corrections.jsonl:%d" % i, "id does not derive from the record")
    if document["policy"] == "checks-only" and (cors or data["suspects.jsonl"]):
        rep.add("doc.policy.withheld", "corrections.jsonl", "a checks-only package carries no corrections or suspects")
    if document["policy"] == "full" and document["rev"] != len(cors):
        rep.add("doc.correction.rev", "document.json", "rev must equal the number of corrections")
    retired_by = {e["id"]: e for e in ids if e["kind"] == "retired"}
    for c in cors:
        for r in c.get("retired", []):
            if r not in retired_by:
                rep.add("doc.id.retired-missing", c["id"], "%s was retired but is not in ids.jsonl" % r)
    if not withheld and hasher is not None:
        last = {}
        for c in cors:
            if c.get("target") in live and kinds[c["target"]] in LEAVES:
                last[c["target"]] = c
            for x in c.get("touched", []):
                if x != c.get("target"):
                    last.pop(x, None)
        for nid, c in last.items():
            if c.get("after") and c["after"] != vals[nid]["node_hash"] and not doc.kids(nid):
                rep.add("doc.correction.after", c["id"], "the last correction of %s does not end at its hash" % nid)

    # Geometry.
    wits = {w["id"]: w for w in data["witnesses.json"]["witnesses"]}
    for w in wits.values():
        if w["id"] != witness_id(w["sha256"]):
            rep.add("doc.witness.id", w["id"], "a witness ID derives from its SHA-256")
    fps = set()
    for i, p in enumerate(data["pages.jsonl"], 1):
        if p["witness"] not in wits:
            rep.add("doc.page.witness", "pages.jsonl:%d" % i, "unknown witness")
        fps.add(p["page_fp"])
    porder = sorted(data["pages.jsonl"], key=lambda p: (p["witness"], p["page"]))
    if porder != data["pages.jsonl"]:
        rep.add("doc.order.pages", "pages.jsonl", "pages sort by witness and page")
    pos = {nid: i for i, nid in enumerate(doc.order)}
    for i, loc in enumerate(data["locations.jsonl"], 1):
        if loc["node"] not in live:
            rep.add("doc.location.node", "locations.jsonl:%d" % i, "unknown node")
        if loc["page_fp"] not in fps:
            rep.add("doc.location.page", "locations.jsonl:%d" % i, "page_fp is not a page of any witness")
    lorder = sorted(data["locations.jsonl"], key=lambda l: (pos.get(l["node"], -1), l["page"], l["page_fp"]))
    if lorder != data["locations.jsonl"]:
        rep.add("doc.order.locations", "locations.jsonl", "locations sort by node order, then page")
    for i, ref in enumerate(data["references.jsonl"], 1):
        if kinds.get(ref["node"]) != "reference":
            rep.add("doc.reference.node", "references.jsonl:%d" % i, "names no reference node")
    rorder = sorted(data["references.jsonl"], key=lambda r: pos.get(r["node"], -1))
    if rorder != data["references.jsonl"]:
        rep.add("doc.order.references", "references.jsonl", "references follow document order")

    sorder = sorted(data["suspects.jsonl"], key=lambda x: x["id"])
    if sorder != data["suspects.jsonl"]:
        rep.add("doc.order.suspects", "suspects.jsonl", "suspects sort by ID")

    # Receipts and fidelity.
    roster = {k["id"]: k for k in data.get("roster.json", {}).get("keys", [])}
    recs = data["receipts.jsonl"]
    rorder = sorted(recs, key=lambda r: (r["at"], r["id"]))
    if rorder != recs:
        rep.add("doc.order.receipts", "receipts.jsonl", "receipts and revocations sort by time, then ID")
    by_id = {r["id"]: r for r in recs}
    for r in recs:
        if r["id"] != record_id(R.record_hash(r)):
            rep.add("doc.receipt.id", r["id"], "id does not derive from the record")
        if r["format"] == "scienting-receipt/1":
            if r["doc"] != document["key"]:
                rep.add("doc.receipt.doc", r["id"], "receipt names another document")
            if r["root"] != crypto.merkle_root(crypto.receipt_items(r)):
                rep.add("doc.receipt.root", r["id"], "Merkle root does not recompute")
            ids_listed = [x[0] for x in r["nodes"]]
            if ids_listed != sorted(ids_listed) or len(set(ids_listed)) != len(ids_listed):
                rep.add("doc.receipt.nodes", r["id"], "nodes sort by ID without repeats")
        if roster and not R.verify_signature(r, roster):
            rep.add("doc.receipt.signature", r["id"], "signature does not verify")
        qa = r.get("qa")
        if qa:
            orig = by_id.get(qa["of"])
            if not orig:
                rep.add("doc.qa.of", r["id"], "the spot-checked receipt is not in this package")
                continue
            ck = roster.get(qa["key"])
            seed = base64.b64decode(qa["seed"])
            if roster:
                if not ck or ck["scheme"] != "ed25519" or "spot-check" not in ck["purposes"]:
                    rep.add("doc.qa.key", r["id"], "the seed key is not a registered Ed25519 spot-check key")
                elif not ck["registered_at"] <= orig["at"]:
                    rep.add("doc.qa.key", r["id"], "the spot-check key was registered after the receipt")
                elif not crypto.ed25519_verify(base64.b64decode(ck["public"]), R.seed_payload(orig), seed):
                    rep.add("doc.qa.seed", r["id"], "the seed is not the checker's signature over the receipt")
            if qa["sample"] != R.sample(orig, seed, qa["rate"], qa["minimum"]):
                rep.add("doc.qa.sample", r["id"], "the sample does not follow from the seed")
    if roster and not withheld:
        if R.distribution(R.effective_levels(nodes, vals, recs, roster)) != document["fidelity"]:
            rep.add("doc.fidelity.report", "document.json", "the fidelity distribution does not recompute")


def read_state(root):
    """The current state of the package at `root`, for the gate (spec/11 §11.4): nodes without their values,
    the ID table and the corrections applied so far. Validate the package first."""
    from .corrections import State
    present = {f["path"] for f in read_json(os.path.join(root, "MANIFEST.json"))["files"]}

    def jl(name):
        return read_jsonl(os.path.join(root, name)) if name in present else []

    values = ("text_hash", "match_hash", "align_key", "own_hash", "node_hash")
    nodes = [{k: v for k, v in n.items() if k not in values} for n in jl("nodes.jsonl")]
    return State(read_json(os.path.join(root, "document.json"))["key"], nodes, jl("ids.jsonl"), jl("corrections.jsonl"))


def n_root_hash(nodes, doc):
    return {n["id"]: n for n in nodes}[doc.root]["node_hash"]