#!/usr/bin/env python3
"""Independent verifier of the carousel corpus (extensions/carousel.md).

Written from the extension's text. It shares no code with wai-rs, which generates the corpus,
and never runs it. It takes from the second implementations of the texts this extension builds
on what they define: group and object receipts from the JWP corpus's
(`../jwp-conformance/jwp_verify.py`, jwp-receipts §2.2, §3), and a claim's steps 1 to 5 from the
claim corpus's (`../claim-conformance/claim_verify.py`, jwp-receipts §3.1), whose step 4 holds
the schedule class's rules with §5.1 and §5.2's division of an object into source blocks.
Everything else this extension adds is written here:

- §4.3: a schedule against the group, its object receipts and their envelopes;
- §5.3: what a receiver needs of a block, K′ read from Appendix A of the extension's text;
- §6: the declarations and schedules that stand, each entry's availability within a loss bound,
  and statements offered;
- §7: the passes a receiver listens to.

For every case it requires the outcome the corpus states, field for field; that every reason
of §9's table appears in the corpus; then it mutates copies of the corpus, one defect each, and
requires every copy to fail; and it drops each of its own checks in turn and requires the corpus
to catch the weakened verifier.

    python3 carousel_verify.py <live-conformance/carousel>

Exit status 0 only if every case verifies, every mutated copy fails and every weakened verifier
is caught. Python 3.8 or later; the `cryptography` package, which the claim verifier imports.
"""

import copy
import json
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
for d in ("claim-conformance", "jwp-conformance"):
    sys.path.insert(0, os.path.join(HERE, "..", d))
import claim_verify as CV  # noqa: E402
import jwp_verify as JV  # noqa: E402

EXTENSION = os.path.join(HERE, "..", "extensions", "carousel.md")

# A weakened run drops one of these checks (see `WEAKENINGS`); None checks everything.
DROPPED = None


def drops(name):
    return DROPPED == name


def extension():
    return open(EXTENSION, encoding="utf-8").read()


def reasons_table():
    """§9's codes, in order."""
    text = extension()
    at = text.index("\n## 9. Reasons")
    end = text.find("\n## ", at + 1)
    rows = text[at:end if end >= 0 else len(text)].splitlines()
    return [r.split("`")[1] for r in rows if r.startswith("| `")]


def k_primes():
    """Appendix A: the column K′ of RFC 6330 Table 2."""
    text = extension()
    a, b = text.index("<!-- k-prime:begin -->"), text.index("<!-- k-prime:end -->")
    return [int(x) for x in re.findall(r"[0-9]+", text[a + len("<!-- k-prime:begin -->"):b])]


K_PRIME = k_primes()


def k_prime(k):
    """§5.3: the least K′ at least k."""
    return next(v for v in K_PRIME if v >= k)


def need(ok, reason):
    if not ok:
        raise CV.Reject(6, reason)


def included(o, g):
    try:
        JV.inclusion(o, g)
        return True
    except JV.Refuse:
        return False


def accepted(o, g, envelope):
    try:
        JV.accept(o, g, envelope)
        return True
    except JV.Refuse:
        return False


# ------------------------------------------------------------------ §4.3: step 6

def check_transmission(c, g):
    """bearer-receipts §4.3."""
    need(JV.link(g).hex() == c["subject"], "transmission_subject_mismatch")
    need(g["signer_pubkey"] == c["signer_pubkey"], "transmission_signer_mismatch")


def check_schedule(c, g, objects):
    """§4.3 against a group that passed jwp-receipts §3 steps 1 and 2, and the object receipts
    held under it with their envelopes where held: [(object, envelope or None)]."""
    start = JV.seconds(g["group_start"])
    need(start is not None, "context_invalid")
    entries = c["body"]["objects"]
    need(drops("subject") or JV.link(g).hex() == c["subject"], "schedule_subject_mismatch")
    need(drops("signer") or g["signer_pubkey"] == c["signer_pubkey"], "schedule_signer_mismatch")
    need(drops("range") or all(e["leaf_index"] < g["leaf_count"] for e in entries), "schedule_leaf_out_of_range")
    need(drops("before") or all(e["first_ms"] // 1000 >= start for e in entries), "schedule_before_group")
    for e in entries:
        for o, envelope in objects:
            if o["object_id"] != e["leaf_index"] or not included(o, g):
                continue
            need(drops("content") or o["content_hash"].hex() == e["content_hash"], "schedule_content_mismatch")
            if envelope is not None and accepted(o, g, envelope):
                need(drops("length") or len(envelope) == e["oti"]["transfer_length"], "schedule_length_mismatch")


# ------------------------------------------------------------------ §5.3, §6.3

def needed(e, k, h):
    """§5.3: every source symbol for no-code, K′ + h for RaptorQ."""
    return k if e["fec"] == "no-code" else k_prime(k) + h


def missed(loss, m):
    """§2: ⌊L · m / 1,000,000⌋."""
    return loss * m // 1_000_000


def pass_end(e, j):
    return e["first_ms"] + (j - 1) * e["period_ms"] + e["duration_ms"]


def block_pass(e, k, loss, h):
    """§6.3: the first pass after which a receiver within the bound holds a block of k."""
    m = k + e["repair"]
    x = missed(loss, m)
    need_ = needed(e, k, h)
    if e["fec"] == "raptorq" and e["esi"] == "advance":
        j = CV.ceil_div(need_, m - x)
    elif e["fec"] == "raptorq":
        j = 1 if m - x >= need_ else None
    else:
        j = 1 if x == 0 or drops("no_code_loss") else None
    return j if j is not None and j <= e["count"] else None


def availability(e, loss, h):
    blocks = CV.schedule_blocks(e)
    passes = [block_pass(e, k, loss, h) for k, _ in blocks]
    j = None if None in passes else max(passes)
    return {
        "pass": j,
        "available_ms": None if j is None else pass_end(e, j),
        "bytes_sent": CV.schedule_bytes(e, blocks),
        "failure_one_in": [100, 10_000, 1_000_000][h] if e["fec"] == "raptorq" else None,
    }


# ------------------------------------------------------------------ §6: the verdict

def bounds(loss, h):
    if loss > 999_999:
        return "loss_bound_invalid"
    if h not in (0, 1, 2):
        return "overhead_invalid"
    return None


def verdict(case):
    bad = bounds(case["loss_ppm"], case["overhead"])
    if bad:
        return {"outcome": "refuse", "reason": bad}
    try:
        g = JV.read_group(case["group"])
        JV.verify_group(g)
    except JV.Refuse as r:
        return {"outcome": "refuse", "reason": r.reason}

    # Step 1: the objects, the same on every bearer.
    objs, overdicts = [], []
    for o in case["objects"]:
        try:
            ob = JV.read_object(o["receipt"])
        except JV.Refuse as r:
            overdicts.append("reject reason=" + r.reason)
            continue
        env = None if o["envelope"] is None else bytes.fromhex(o["envelope"])
        objs.append((ob, env))
        try:
            if env is None:
                JV.inclusion(ob, g)
                overdicts.append("inclusion-verified")
            else:
                JV.accept(ob, g, env)
                overdicts.append("delivered-verified")
        except JV.Refuse as r:
            overdicts.append("reject reason=" + r.reason)

    # Step 2: each claim, §3.1 steps 1-5, then step 6.
    held, claims, standing_claims = {}, [], []
    for wire in case["claims"]:
        try:
            c = CV.verify_claim(wire)
        except CV.Reject as r:
            claims.append(f"reject step={r.step} reason={r.reason}")
            standing_claims.append(None)
            continue
        previous = held.get(c["signer_pubkey"])
        held[c["signer_pubkey"]] = c
        try:
            CV.check_chain(c, previous)
            if c["class"] == "transmission":
                check_transmission(c, g)
            elif c["class"] == "schedule":
                check_schedule(c, g, objs)
            else:
                raise CV.Reject(6, "context_mismatch")
        except CV.Reject as r:
            claims.append(f"reject step={r.step} reason={r.reason}")
            standing_claims.append(None)
            continue
        claims.append("stands" if c["class"] == "schedule" else "declared")
        standing_claims.append(c)

    def by_instance(cls):
        out = {}
        for i, c in enumerate(standing_claims):
            if c is not None and c["class"] == cls:
                out.setdefault(c["body"]["bearer_id"], []).append(i)
        return out

    # Step 3: the declarations.
    declared, contradicted = {}, set()
    for bid, idx in by_instance("transmission").items():
        bodies = [standing_claims[i]["body"] for i in idx]
        if all(b == bodies[0] for b in bodies):
            declared[bid] = bodies[0]["bearer"]
        else:
            contradicted.add(bid)
            for i in idx:
                claims[i] = "set-aside reason=bearer_contradicted"

    # Step 4: the schedules.
    schedules = {}
    for bid, idx in by_instance("schedule").items():
        bodies = [standing_claims[i]["body"] for i in idx]
        if not drops("schedule_contradiction") and any(b != bodies[0] for b in bodies):
            aside = "schedule_contradicted"
        elif bid in contradicted:
            aside = "bearer_contradicted"
        elif bid not in declared and not drops("undeclared"):
            aside = "bearer_not_declared"
        elif declared.get(bid) == "bidirectional" and not drops("bidirectional"):
            aside = "schedule_on_bidirectional"
        else:
            aside = None
        if aside:
            for i in idx:
                claims[i] = "set-aside reason=" + aside
        else:
            schedules[bid] = bodies[0]

    # Step 5.
    loss, h = case["loss_ppm"], case["overhead"]
    lines = []
    for bid in sorted(schedules):
        for e in schedules[bid]["objects"]:
            lines.append({"bearer_id": bid, "leaf_index": e["leaf_index"], **availability(e, loss, h)})

    # §6.4.
    def held_by(bid, i, by):
        line = next((x for x in lines if x["bearer_id"] == bid and x["leaf_index"] == i), None)
        if line is None:
            return "object_not_scheduled"
        if line["available_ms"] is None:
            return "availability_not_within_schedule"
        if line["available_ms"] > by and not drops("deadline"):
            return "availability_after_deadline"
        return None

    statement = None
    s = case["statement"]
    if s is not None:
        if s["kind"] == "object":
            r = held_by(s["bearer_id"], s["index"], s["by_ms"])
        else:
            r = next((r for r in (held_by(s["bearer_id"], i, s["by_ms"]) for i in range(g["leaf_count"])) if r), None)
        statement = r or "holds"
    return {"outcome": "verdict", "objects": overdicts, "claims": claims, "lines": lines, "statement": statement}


# ------------------------------------------------------------------ §7: a plan

def entry_of(v):
    """A plan's entry, read as a schedule's entry is at jwp-receipts §3.1 steps 2 and 4."""
    body = {"bearer_id": "urn:example:plan", "objects": [v], "attests": "scheduled"}
    CV.check_body("schedule", body)
    CV.schedule_entry_rules(v)
    return v


def plan(case):
    e = entry_of(case["entry"])
    loss, h, t, held = case["loss_ppm"], case["overhead"], case["at_ms"], case["held"]
    bad = bounds(loss, h)
    if bad:
        return {"outcome": "refuse", "reason": bad}
    ks = [k for k, c in CV.schedule_blocks(e) for _ in range(c)]
    if len(held) != len(ks) or (e["fec"] == "no-code" and any(n > k for n, k in zip(held, ks))
                                and not drops("holding")):
        return {"outcome": "refuse", "reason": "holding_invalid"}
    n = list(held)
    open_ = lambda: [b for b, k in enumerate(ks) if n[b] < needed(e, k, h)]  # noqa: E731
    first = 1 if t <= e["first_ms"] else (2 if e["count"] == 1 else CV.ceil_div(t - e["first_ms"], e["period_ms"]) + 1)
    if drops("late_pass"):
        first = 1 if t <= e["first_ms"] else (t - e["first_ms"]) // max(1, e["period_ms"]) + 1
    if not open_():
        return {"outcome": "plan", "passes": [], "complete_ms": t, "listen_ms": 0}
    passes = []
    for j in range(first, e["count"] + 1):
        blocks = open_()
        if not blocks:
            break
        gains = {}
        for b in blocks:
            m = ks[b] + e["repair"]
            x = missed(loss, m)
            if e["fec"] == "raptorq" and e["esi"] == "advance":
                gains[b] = m - x
            else:
                gains[b] = max(0, m - n[b] - (0 if drops("again") else x))
        if not any(gains.values()):
            break
        for b, gain in gains.items():
            n[b] += gain
        passes.append(j)
    complete = None if open_() else pass_end(e, passes[-1])
    return {
        "outcome": "plan",
        "passes": [{"pass": j, "start_ms": e["first_ms"] + (j - 1) * e["period_ms"], "end_ms": pass_end(e, j)} for j in passes],
        "complete_ms": complete,
        "listen_ms": len(passes) * e["duration_ms"],
    }


# ------------------------------------------------------------------ the corpus

def strings(v, out):
    if isinstance(v, str):
        out.append(v)
    elif isinstance(v, list):
        for x in v:
            strings(x, out)
    elif isinstance(v, dict):
        for x in v.values():
            strings(x, out)


def claim_outcome(wire):
    try:
        CV.verify_claim(wire)
        return "accept"
    except CV.Reject as r:
        return f"reject step={r.step} reason={r.reason}"


def check(doc, quiet=False):
    """The failures of a corpus document: each case's outcome, and every reason of §9."""
    fails = []
    for kind, run in (("verdicts", verdict), ("claims", lambda c: claim_outcome(c["claim"])), ("plans", plan)):
        for c in doc[kind]:
            try:
                got = run(c)
            except CV.Reject as r:
                got = {"error": f"reject step={r.step} reason={r.reason}"}
            except (KeyError, TypeError, ValueError, IndexError, AttributeError, StopIteration, ZeroDivisionError) as e:
                got = {"error": repr(e)}
            if got != c["outcome"]:
                fails.append(c["name"])
                if not quiet:
                    print(f"FAIL {c['name']}\n  stated {json.dumps(c['outcome'])}\n  got    {json.dumps(got)}")
    stated = []
    for kind in ("verdicts", "claims", "plans"):
        for c in doc[kind]:
            strings(c["outcome"], stated)
    for code in reasons_table():
        if not any(s == code or s.endswith("reason=" + code) for s in stated):
            fails.append("no case gives " + code)
    return fails


def mutations(doc):
    """Copies of the corpus, one defect each, every one of which must fail."""
    def case(d, name):
        return next(c for k in ("verdicts", "claims", "plans") for c in d[k] if c["name"] == name)

    def entry(c, i):
        w = json.loads(c["claims"][i])
        return w, w["body"]["objects"]

    def m_pass(d):
        case(d, "loss_30_percent")["outcome"]["lines"][0]["pass"] = 1

    def m_bytes(d):
        case(d, "lossless")["outcome"]["lines"][0]["bytes_sent"] += 64

    def m_failure(d):
        case(d, "overhead_1")["outcome"]["lines"][0]["failure_one_in"] = 100

    def m_loss(d):
        # 88%: object 0 gains 3 a pass, and needs 7.
        case(d, "loss_85_percent_last_pass")["loss_ppm"] = 880_000

    def m_repair(d):
        c = case(d, "lossless")
        w, es = entry(c, 1)
        es[0]["repair"] = 7
        c["claims"][1] = json.dumps(w, separators=(",", ":"), sort_keys=True)

    def m_statement(d):
        case(d, "object_after_deadline")["statement"]["by_ms"] += 1

    def m_plan(d):
        case(d, "too_late")["outcome"]["complete_ms"] = case(d, "too_late")["outcome"]["passes"][0]["end_ms"]

    def m_held(d):
        # 19 short of a block whose pass adds 18: two passes.
        case(d, "holding_part")["held"][0] = 1

    def m_accept(d):
        case(d, "bytes_greatest")["outcome"] = "reject step=4 reason=schedule_passes_invalid"

    def m_reason_gone(d):
        d["verdicts"].remove(case(d, "schedule_on_bidirectional"))

    return [m_pass, m_bytes, m_failure, m_loss, m_repair, m_statement, m_plan, m_held, m_accept, m_reason_gone]


WEAKENINGS = ["subject", "signer", "range", "before", "content", "length", "no_code_loss", "schedule_contradiction",
              "undeclared", "bidirectional", "deadline", "holding", "late_pass", "again"]


def main(argv):
    global DROPPED
    if len(argv) != 2:
        print(__doc__)
        return 2
    doc = json.load(open(os.path.join(argv[1], "vectors.json"), encoding="utf-8"))
    extra = sorted(set(os.listdir(argv[1])) - {"vectors.json", "README.md"})
    fails = check(doc)
    if extra:
        fails.append(f"files the corpus does not hold: {extra}")
    if len(K_PRIME) != 477 or any(a >= b for a, b in zip(K_PRIME, K_PRIME[1:])) or K_PRIME[-1] != CV.K_PRIME_MAX:
        fails.append("Appendix A is not 477 increasing values ending at K'max")
    print(f"{len(doc['verdicts'])} verdict cases, {len(doc['claims'])} claim cases, {len(doc['plans'])} plans: "
          f"{len(fails)} failure(s)")
    caught = 0
    for m in mutations(doc):
        d = copy.deepcopy(doc)
        m(d)
        if check(d, quiet=True):
            caught += 1
        else:
            fails.append(f"mutation {m.__name__} was not caught")
    print(f"{caught} of {len(mutations(doc))} mutated copies fail")
    weak = 0
    for w in WEAKENINGS:
        DROPPED = w
        if check(doc, quiet=True):
            weak += 1
        else:
            fails.append(f"a verifier without its {w} check passes the corpus")
    DROPPED = None
    print(f"{weak} of {len(WEAKENINGS)} weakened verifiers are caught")
    for f in fails:
        print("FAIL", f)
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))