#!/usr/bin/env python3
"""Independent verifier of the per-session ad track corpus (extensions/live-compose.md §5.10).

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 §5.10 builds on what they
define: a composition descriptor read, its digest and signature, the layout of a room and its ad
frames (§3, §5) from the live-compose corpus's (`live_compose_verify.py`); a pod-input record read,
its frame, ad region and share from the stitched-presentation corpus's
(`stitched_presentation_verify.py`); and Ed25519 verification as SPEC §7 *Signatures* states it
from the routine every Python verifier of WAI's corpora shares (`../claim-conformance/_ed25519.py`).
Everything §5.10 adds is written here:

- §5.10.1: the session's ad track and its `track`;
- §5.10.2: the binding read, its bounds, its digest and its signature;
- §5.10.3: a binding verified against a descriptor and a record, and what an accepted one states;
- §5.10.4: the layout of a session around its binding.

For every case it requires the outcome the corpus states, member for member; that every reason of
§5.10.5's table appears in the corpus; that each accepted binding's ad region and share are the
descriptor's own ad frame's; and that the program's tile in a session's layout is the room's. 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 -I compose_ad_track_verify.py <live-conformance/compose-ad-track>

Exit status 0 only if every case verifies, every mutated copy fails and every weakened verifier is
caught. Python 3.8 or later; standard library only.
"""

import copy
import hashlib
import json
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import _ed25519 as ED  # noqa: E402
import live_compose_verify as LCV  # noqa: E402
import stitched_presentation_verify as SPV  # noqa: E402

EXTENSION = os.path.join(HERE, "..", "extensions", "live-compose.md")
MAGIC = b"WAT1"                                     # §5.10.2
DOMAIN = b"wai:live-compose-ad-track\x01"           # §5.10.2
LEN = 4 + 1 + 32 + 32 + 32 + 16 + 1 + 8 + 2 + 2 + 32 + 64
MAX_EDGE = 1920
SQUEEZE, DOUBLE_BOX = 3, 4
INVALID = "ad_track_invalid"
NOT_PERMITTED = "ad_track_not_permitted"
MISMATCH = "ad_track_mismatch"

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


def drops(name):
    return DROPPED == name


def reasons_table():
    """§5.10.5's codes, in order."""
    text = open(EXTENSION, encoding="utf-8").read()
    at = text.index("\n#### 5.10.5 Reasons")
    end = text.find("\n#", at + 1)
    return [r.split("`")[1] for r in text[at:end].splitlines() if r.startswith("| `")]


class Refused(Exception):
    def __init__(self, reason, at, member=None):
        super().__init__(reason)
        self.reason, self.at, self.member = reason, at, member


def sha(*parts):
    h = hashlib.sha256()
    for p in parts:
        h.update(p)
    return h.digest()


# ------------------------------------------------------------------ §5.10.1, §5.10.2

def track_of(room, session):
    """§5.10.1: the ad track's `track`: the owner `ad_` and the session in hexadecimal."""
    return LCV.track_digest(room, "ad_" + session.hex(), LCV.VIDEO, "0")


def read_binding(b):
    """§5.10.2: the binding, or Refused at `binding`."""
    def no(reason):
        raise Refused(reason, "binding")
    if len(b) < 4:
        no("truncated")
    if b[:4] != MAGIC:
        no("malformed")
    if len(b) < 5:
        no("truncated")
    if b[4] != 1:
        no("unknown_version")
    if len(b) < LEN:
        no("truncated")
    if len(b) > LEN:
        no("malformed")
    x = {"room": b[5:37], "descriptor": b[37:69], "record": b[69:101], "session": b[101:117], "slot": b[117],
         "start": struct.unpack("<Q", b[118:126])[0], "width": struct.unpack("<H", b[126:128])[0],
         "height": struct.unpack("<H", b[128:130])[0], "signer": b[130:162], "sig": b[162:226], "body": b[:162]}
    if not (1 <= x["width"] <= MAX_EDGE and 1 <= x["height"] <= MAX_EDGE) and not drops("bounds"):
        no(INVALID)
    if not drops("signature"):
        domain = b"" if drops("signature_domain") else SPV.POD_DOMAIN if drops("signature_pod_domain") else DOMAIN
        bad = ED.verify(x["signer"], domain + x["body"], x["sig"])
        if bad:
            no(bad)
    return x


def frame_of(p):
    """§5.10.3 step 7 `frame`: the record's frame a descriptor of an ad frame fixes."""
    sq = p["layout"] == SQUEEZE
    return {"kind": p["layout"], "width": p["width"], "height": p["height"], "corner": p["corner"],
            "window_w": p["pip_width"] if sq else 0, "window_h": p["pip_height"] if sq else 0,
            "gap": 0 if sq else p["gap"], "aspect": 0 if sq else p["aspect"]}


def same_frame(a, b):
    keys = [k for k in a if not (drops("frame_window") and k in ("window_w", "window_h"))
            and not (drops("frame_corner") and k == "corner")]
    return all(a[k] == b[k] for k in keys)


# ------------------------------------------------------------------ §5.10.3

def verify(dbytes, rbytes, bbytes):
    """§5.10.3: the accepted binding's facts, or Refused."""
    try:                                                                          # step 1
        c = LCV.read_descriptor(dbytes)
    except LCV.Refuse as e:
        raise Refused(e.reason, "descriptor")
    try:                                                                          # step 2
        rec = SPV.read_record(rbytes)
    except SPV.Bad as e:
        raise Refused(e.args[0], "record")
    if not drops("record_signature"):
        bad = ED.verify(rec["signer"], SPV.POD_DOMAIN + rec["body"], rec["sig"])
        if bad:
            raise Refused(bad, "record")
    b = read_binding(bbytes)                                                      # step 3

    def no(reason, member=None):
        raise Refused(reason, "binding", member)
    if b["signer"] != rec["signer"] and not drops("signer"):                      # step 4
        no(NOT_PERMITTED)
    if b["room"] != c["room"] and not drops("room"):                              # step 5
        no("wrong_room")
    if b["descriptor"] != c["digest"] and not drops("descriptor"):               # step 6
        no("descriptor_mismatch")
    p = c["params"]                                                               # step 7
    if p["layout"] not in (SQUEEZE, DOUBLE_BOX) and not drops("layout"):
        no(MISMATCH, "layout")
    if b["record"] != sha(rec["body"]) and not drops("record"):
        no(MISMATCH, "record")
    if rec["session"] == bytes(16) and not drops("shared"):
        no(MISMATCH, "session")
    if rec["session"] != b["session"] and not drops("session"):
        no(MISMATCH, "session")
    if p["layout"] in (SQUEEZE, DOUBLE_BOX) and not same_frame(rec["frame"], frame_of(p)) and not drops("frame"):
        no(MISMATCH, "frame")
    if b["slot"] >= len(rec["ads"]) and not drops("slot"):
        no(MISMATCH, "slot")
    if b["start"] < c["effective"] and not drops("effective"):                    # step 8
        no("not_yet_effective")
    f = rec["frame"]
    duration = rec["ads"][min(b["slot"], len(rec["ads"]) - 1)]["duration_ms"]
    return {"outcome": "accepted", "digest": sha(b["body"]).hex(), "session": b["session"].hex(), "slot": b["slot"],
            "track": track_of(b["room"], b["session"]).hex(), "width": b["width"], "height": b["height"],
            "start": b["start"], "end": min(b["start"] + 1000 * duration, 2 ** 64 - 1),
            "share_bp": SPV.share_of(f), "region": [list(r) for r in SPV.frame_region(f)]}


def binding_outcome(case):
    try:
        return verify(bytes.fromhex(case["descriptor"]), bytes.fromhex(case["record"]), bytes.fromhex(case["binding"]))
    except Refused as r:
        out = {"outcome": "refuse", "reason": r.reason, "at": r.at}
        if r.member is not None:
            out["member"] = r.member
        return out


# ------------------------------------------------------------------ §5.10.4

def entrants_of(v):
    return [{"principal": bytes.fromhex(e["principal"]), "join": e["join"],
             "video": None if e["video"] is None else tuple(e["video"]), "audio": e["audio"],
             "attrs": [tuple(kv) for kv in e["attrs"]]} for e in v]


def with_track(tile):
    return dict(tile, track=None)


def session_layout(p, entrants, bound, t):
    """§5.10.4: the layout at `t` of a session holding `bound`, or none."""
    room = LCV.layout(p, entrants)
    plain = {"outcome": "layout", "order": room["order"], "tiles": [with_track(x) for x in room["tiles"]],
             "voices": room["voices"]}
    if bound is None or p["layout"] not in (SQUEEZE, DOUBLE_BOX):
        return plain
    if not (bound["start"] <= t or drops("span_start")) or not (t < bound["end"] or (drops("span_end") and t == bound["end"])):
        return plain
    ad = next((x["principal"] for x in room["tiles"] if x["place"] == "ad"), None)
    prog = next((x["principal"] for x in room["tiles"] if x["place"] == "program"), None)
    w, h = p["width"], p["height"]

    def track_tile(rect):
        mode = p["fill"] if p["fill"] != LCV.DEFAULT else LCV.COVER
        dst, src = LCV.show(rect, (bound["width"], bound["height"]), mode)
        return {"principal": None, "track": bound["track"], "place": "ad", "z": 0, "tile": list(rect), "dst": list(dst),
                "src": list(src)}
    if prog is not None and not drops("program_tile"):
        pt, at, _ = LCV.ad_frame(w, h, p["layout"], p["corner"], p["pip_width"], p["pip_height"], p["gap"], p["aspect"])
        m = next(e for e in entrants if e["principal"].hex() == prog)
        program = with_track(LCV.tile_of(p, m, pt, "program", 1 if p["layout"] == SQUEEZE else 0))
        tiles = [track_tile(at), program] if p["layout"] == SQUEEZE else [program, track_tile(at)]
    else:
        tiles = [track_tile((0, 0, w, h))]
    voices = [v for v in room["voices"] if v["principal"] != ad or drops("ad_voice")]
    return {"outcome": "layout", "order": room["order"], "tiles": tiles, "voices": voices}


def layout_outcome(case):
    try:
        return session_layout(case["params"], entrants_of(case["entrants"]), case["binding"], case["t"])
    except LCV.Refuse as e:
        return {"outcome": "refuse", "reason": e.reason}


RUNS = (("bindings", binding_outcome), ("layouts", layout_outcome))
KINDS = tuple(k for k, _ in RUNS)


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

def invariances(doc, fail):
    for c in doc["bindings"]:
        o = c["outcome"]
        if o["outcome"] != "accepted":
            continue
        # The ad region and its share are the descriptor's own ad frame's.
        p = LCV.read_descriptor(bytes.fromhex(c["descriptor"]))["params"]
        _, _, region = LCV.ad_frame(p["width"], p["height"], p["layout"], p["corner"], p["pip_width"], p["pip_height"],
                                    p["gap"], p["aspect"])
        if [list(r) for r in region] != o["region"] or LCV.share_bp(region, p["width"], p["height"]) != o["share_bp"]:
            fail(f"bindings/{c['name']}: its region is not the descriptor's ad frame's")
    for c in doc["layouts"]:
        o = c["outcome"]
        if o["outcome"] != "layout":
            continue
        # The program's tile is where the room's own frame puts it, whenever the room has an ad.
        room = LCV.layout(c["params"], entrants_of(c["entrants"]))
        mine = [x for x in o["tiles"] if x["place"] == "program"]
        theirs = [x for x in room["tiles"] if x["place"] == "program"]
        if mine and theirs and any(x["place"] == "ad" for x in room["tiles"]):
            if [with_track(x) for x in theirs] != mine:
                fail(f"layouts/{c['name']}: the program is not where the room's frame puts it")


def check(doc, quiet=False, full=True):
    fails = []

    def fail(msg):
        fails.append(msg)
        if not quiet:
            print("FAIL", msg)

    for kind, run in RUNS:
        cases = doc.get(kind) or []
        if not cases:
            fail(f"no {kind} cases")
        for c in cases:
            try:
                got = run(c)
            except (KeyError, TypeError, ValueError, AttributeError, IndexError, struct.error) as e:
                fail(f"{kind}/{c.get('name')}: {e!r}")
                continue
            if got != c["outcome"]:
                fail(f"{kind}/{c['name']}: stated {json.dumps(c['outcome'])} got {json.dumps(got)}")
    reasons = [c["outcome"].get("reason") for k in KINDS for c in doc.get(k) or []]
    for r in reasons_table():
        if r not in reasons:
            fail(f"no case gives {r}")
    if full:
        try:
            invariances(doc, fail)
        except (KeyError, TypeError, ValueError, AttributeError, LCV.Refuse) as e:
            fail(f"invariances: {e!r}")
    return fails


WEAKENINGS = (
    "bounds", "signature", "signature_domain", "signature_pod_domain", "record_signature", "signer", "room", "descriptor", "layout", "record",
    "shared", "session", "frame", "frame_window", "frame_corner", "slot", "effective", "span_start", "span_end",
    "program_tile", "ad_voice",
)


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

    def flip(h, at):
        b = bytearray.fromhex(h)
        b[at] ^= 1
        return b.hex()

    def m_member(d):
        case(d, "bindings", "frame_of_another_window")["outcome"]["member"] = "session"

    def m_reason(d):
        case(d, "bindings", "signer_not_the_records")["outcome"]["reason"] = "bad_signature"

    def m_share(d):
        case(d, "bindings", "squeeze")["outcome"]["share_bp"] -= 1

    def m_end(d):
        case(d, "bindings", "double_box")["outcome"]["end"] += 1

    def m_track(d):
        o = case(d, "bindings", "squeeze")["outcome"]
        o["track"] = flip(o["track"], 0)

    def m_binding_bytes(d):
        c = case(d, "bindings", "squeeze")
        c["binding"] = flip(c["binding"], 120)

    def m_record_bytes(d):
        c = case(d, "bindings", "double_box")
        c["record"] = flip(c["record"], 30)

    def m_at(d):
        case(d, "bindings", "record_signature")["outcome"]["at"] = "binding"

    def m_t(d):
        case(d, "layouts", "squeeze_last")["t"] += 1

    def m_tile(d):
        case(d, "layouts", "double_box_during")["outcome"]["tiles"][1]["tile"][0] += 2

    def m_voice(d):
        o = case(d, "layouts", "squeeze_during")["outcome"]
        o["voices"] = o["voices"] + [{"principal": d["keys"]["other"], "gain": 4096, "pan": 0}]

    def m_reason_gone(d):
        d["bindings"] = [c for c in d["bindings"] if c["outcome"].get("reason") != INVALID]

    def m_no_kind(d):
        d["layouts"] = []

    return [m_member, m_reason, m_share, m_end, m_track, m_binding_bytes, m_record_bytes, m_at, m_t, m_tile, m_voice,
            m_reason_gone, m_no_kind]


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}")
    print(f"{sum(len(doc[k]) for k in KINDS)} cases ({', '.join(f'{len(doc[k])} {k}' for k in KINDS)}): {len(fails)} failure(s)")
    caught = 0
    ms = mutations(doc)
    for m in ms:
        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(ms)} mutated copies fail")
    weak = 0
    for w in WEAKENINGS:
        DROPPED = w
        try:
            if check(doc, quiet=True, full=False):
                weak += 1
            else:
                fails.append(f"a verifier without its {w} check passes the corpus")
        finally:
            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))