#!/usr/bin/env python3
"""tempo_verify.py: a second implementation of playout timing (extensions/co-watch.md §15).

    tempo_verify.py [VECTORS.json]

Written from co-watch §5, §6 and §15, moq-streaming-format §4.4 and
draft-nandakumar-moq-tempo-00; it shares no code with joule-moq-room or wai-rs. It computes the
controlled position itself, and for moqt-conformance/tempo/vectors.json it reads every object's
properties from their bytes, expands every short timestamp, reads every wai_tempo member, holds
every timing claim to the room's controls, makes every change of delay as controls, and
computes every sink's target; then it changes cases and requires each change to be caught.

Exit 0 when everything holds, 1 with one line a finding otherwise. Standard library only.
"""
import copy
import json
import os
import sys

PLAY, PAUSE, SEEK, RATE = 1, 2, 3, 4
SHORT = 16384


class Refuse(Exception):
    pass


def varint(b):
    if not b:
        raise Refuse("tempo_malformed")
    ones = 0
    while ones < 8 and b[0] & (0x80 >> ones):
        ones += 1
    n = 9 if ones == 8 else ones + 1
    if len(b) < n:
        raise Refuse("tempo_malformed")
    if n == 9:
        return int.from_bytes(b[1:9], "big"), 9
    return int.from_bytes(bytes([b[0] & (0xFF >> n)]) + b[1:n], "big"), n


def pairs(b):
    out, at, ty = [], 0, 0
    while at < len(b):
        d, n = varint(b[at:])
        at, ty = at + n, ty + d
        if ty % 2 == 0:
            v, n = varint(b[at:])
            at += n
        else:
            ln, n = varint(b[at:])
            at += n
            if at + ln > len(b):
                raise Refuse("tempo_malformed")
            v, at = b[at:at + ln], at + ln
        out.append((ty, v))
    return out


def read_timing(immutable, mutable, types):
    """§15.2: (capture, delay, hop), each None where the object does not state it."""
    inner, outer = pairs(immutable), pairs(mutable)

    def only(ps, ty):
        got = [v for t, v in ps if t == ty]
        if len(got) > 1:
            raise Refuse("tempo_malformed")
        return got[0] if got else None

    d, c, h = types["playout_delay"], types["capture_timestamp"], types["hop_time"]
    if only(outer, d) is not None or only(outer, c) is not None or only(inner, h) is not None:
        raise Refuse("tempo_malformed")

    def one(b):
        if b is None:
            return None
        v, n = varint(b)
        if n != len(b):
            raise Refuse("tempo_malformed")
        return v

    delay = only(inner, d)
    if delay is not None:
        if len(delay) != 4:
            raise Refuse("tempo_malformed")
        delay = int.from_bytes(delay, "big")
    return one(only(inner, c)), delay, one(only(outer, h))


def expand(ts, now):
    """The draft's §4.1."""
    if ts >= SHORT:
        return ts
    full = now - now % SHORT + ts
    if full > now:
        full -= SHORT
    return full if full >= 0 else None


def read_catalog(v):
    """moq-streaming-format §4.4: (max delay, default delay, (objects per 720 s, start) or None)."""
    bad = Refuse("catalog_tempo_malformed")
    if not isinstance(v, dict) or set(v) - {"max_delay_ms", "playout_delay_ms", "objects_per_720s", "start_ms"}:
        raise bad

    def integer(k, top):
        if k not in v:
            return None
        x = v[k]
        if isinstance(x, bool) or not isinstance(x, int) or not 0 <= x <= top:
            raise bad
        return x

    mx = integer("max_delay_ms", (1 << 32) - 1)
    if mx is None:
        raise bad
    default = integer("playout_delay_ms", mx)
    n, start = integer("objects_per_720s", (1 << 53) - 1), integer("start_ms", (1 << 53) - 1)
    if (n is None) != (start is None) or n == 0:
        raise bad
    return mx, default, None if n is None else (n, start)


# co-watch §5.4, §6: the states of the controlled position, in order.

def states(controls, duration):
    """[(instant, playing, rate, anchor instant, anchor position)], the source starting at 0."""
    order = sorted(range(len(controls)), key=lambda i: (controls[i][1], controls[i][1] - 300000, i))
    st = [0, False, 1000000, 0, 0]
    out = [tuple(st)]
    for i in order:
        action, at, arg = controls[i]
        pos = position(st, at, duration)
        st = [at, st[1], st[2], at, pos]
        if action == PLAY:
            st[1] = True
        elif action == PAUSE:
            st[1] = False
        elif action == SEEK:
            st[4] = arg
        else:
            st[2] = arg
        out.append(tuple(st))
    return out


def position(st, x, duration):
    _, playing, rate, a, p = st
    return min(duration, p + (x - a) * rate // 1000000) if playing else p


def state_at(sts, x):
    return [s for s in sts if s[0] <= x][-1]


def position_at(sts, x, duration):
    return position(state_at(sts, x), x, duration)


def due(sts, duration, frm, q, ln):
    """§15.3."""
    end = q + max(ln, 1)
    for i, s in enumerate(sts):
        hi = sts[i + 1][0] if i + 1 < len(sts) else None
        if hi is not None and hi <= frm:
            continue
        lo = max(s[0], frm)
        at_lo = position(s, lo, duration)
        if at_lo >= end:
            continue
        if at_lo >= q:
            x = lo
        elif not s[1] or q > duration:
            continue
        else:
            x = s[3] + -(-(q - s[4]) * 1000000 // s[2])
        if hi is not None and x >= hi:
            continue
        if q <= position(s, x, duration) < end:
            return x
    return None


def check_claim(sts, duration, t0, c):
    """§15.4, in its order."""
    if c["delay_ms"] > c["max_delay_ms"]:
        raise Refuse("tempo_delay_out_of_bound")
    stated = c["capture_ms"] + c["delay_ms"]
    if stated >= 1 << 64:
        raise Refuse("tempo_malformed")
    if stated * 1000 < t0 or stated * 1000 >= 1 << 64:
        raise Refuse("tempo_before_room")
    x = due(sts, duration, c["from_us"], c["position_us"], c["length_us"])
    if x is None:
        raise Refuse("tempo_not_due")
    if -(-(t0 + x) // 1000) != stated:
        raise Refuse("tempo_off_position")
    return x


def change(sts, duration, at, delta):
    """§15.5: a change of delay as controls."""
    s = state_at(sts, at)
    if not s[1] or delta == 0:
        return []
    if delta > 0:
        return [[PAUSE, at, 0], [PLAY, at + delta * 1000, 0]]
    return [[SEEK, at, min(duration, position(s, at, duration) + -delta * 1000 * s[2] // 1000000)]]


def check(doc):
    fails = []
    types, t0, duration = doc["property_types"], doc["t0_us"], doc["duration_us"]
    for i, c in enumerate(doc["properties"]):
        try:
            got = read_timing(bytes.fromhex(c["immutable"]), bytes.fromhex(c["mutable"]), types)
        except Refuse as r:
            got = str(r)
        if got != (c["capture_ms"], c["delay_ms"], c["hop_ms"]):
            fails.append("properties %d: read %s" % (i, got))
    for c in doc["malformed"]:
        try:
            read_timing(bytes.fromhex(c["immutable"]), bytes.fromhex(c["mutable"]), types)
            fails.append("malformed %r was read" % c["what"])
        except Refuse as r:
            if str(r) != c["reason"]:
                fails.append("malformed %r: %s" % (c["what"], r))
    for c in doc["short"]:
        if expand(c["sent"], c["now_ms"]) != c["full"]:
            fails.append("short %s" % c)
    for c in doc["catalog"]:
        try:
            mx, default, aligned = read_catalog(c["wai_tempo"])
            got = "accept"
            times = [None if aligned is None else aligned[1] + (g + o) * 720000 // aligned[0] for g, o in [(0, 0), (0, 1), (7, 0), (7, 3), (21600, 0)]]
            if times != c.get("aligned"):
                fails.append("catalog %s: aligned capture times %s" % (c["wai_tempo"], times))
        except Refuse as r:
            got = str(r)
        if got != c["outcome"]:
            fails.append("catalog %s: %s, stated %s" % (c["wai_tempo"], got, c["outcome"]))
    scripts = [states(s["controls"], duration) for s in doc["scripts"]]
    held, refused = 0, set()
    for c in doc["claims"]:
        sts = scripts[c["script"]]
        try:
            x = check_claim(sts, duration, t0, c)
            got = {"due_us": x}
            held += 1
            if not c["position_us"] <= position_at(sts, x, duration) < c["position_us"] + c["length_us"]:
                fails.append("claim %r: the position at its instant is outside the object" % c["what"])
            if x > 0 and c["position_us"] <= position_at(sts, x - 1, duration) < c["position_us"] + c["length_us"] and x - 1 >= c["from_us"]:
                fails.append("claim %r: the position was inside the object an instant before" % c["what"])
        except Refuse as r:
            got = {"reason": str(r)}
            refused.add(str(r))
        if got != c["outcome"]:
            fails.append("claim %r: %s, stated %s" % (c["what"], got, c["outcome"]))
    if held < 8 or not {"tempo_delay_out_of_bound", "tempo_before_room", "tempo_not_due", "tempo_off_position"} <= refused:
        fails.append("the claims do not reach every outcome")
    for c in doc["changes"]:
        base = doc["scripts"][c["script"]]["controls"]
        sts = scripts[c["script"]]
        added = change(sts, duration, c["at_us"], c["delta_ms"])
        after = states(base + added, duration)
        got = (added, due(sts, duration, 0, c["position_us"], 40000), due(after, duration, 0, c["position_us"], 40000))
        if got != (c["controls"], c["due_before_us"], c["due_after_us"]):
            fails.append("change at %d by %d: %s" % (c["at_us"], c["delta_ms"], got))
        if c["delta_ms"] > 0 and added:
            d = c["delta_ms"] * 1000
            if got[1] is not None and got[1] >= c["at_us"] and got[2] != got[1] + d:
                fails.append("change at %d by %d: the shift is not the delay" % (c["at_us"], c["delta_ms"]))
            for x in (c["at_us"] + d, c["at_us"] + d + 1234567):
                if position_at(after, x, duration) != position_at(sts, x - d, duration):
                    fails.append("change at %d by %d: the position at %d is not the earlier one" % (c["at_us"], c["delta_ms"], x))
    for c in doc["targets"]:
        p, d, b = c["delay_ms"], c["d_oneway"], c["d_baseline"]
        got = (c["capture_ms"] + p, c["receive_ms"] + p - (d - b), c["hop_ms"] + p + b)
        if got != (c["all_synced"], c["sink_unsynced"], c["relays_only"]):
            fails.append("targets %s" % (got,))
    for k, least in (("properties", 5), ("malformed", 10), ("short", 8), ("catalog", 16), ("claims", 20), ("changes", 7), ("targets", 3)):
        if len(doc[k]) < least:
            fails.append("%s: fewer than %d" % (k, least))
    return fails


def main(argv):
    here = os.path.dirname(os.path.abspath(__file__))
    path = argv[1] if len(argv) > 1 else os.path.join(here, "tempo", "vectors.json")
    with open(path, encoding="utf-8") as fh:
        doc = json.load(fh)
    fails = check(doc)
    off = next(i for i, c in enumerate(doc["claims"]) if c["outcome"].get("reason") == "tempo_off_position")
    ok = next(i for i, c in enumerate(doc["claims"]) if "due_us" in c["outcome"])
    bad_cat = next(i for i, c in enumerate(doc["catalog"]) if c["outcome"] != "accept")
    edits = [
        ("a held claim's capture time", lambda d: d["claims"][ok].__setitem__("capture_ms", d["claims"][ok]["capture_ms"] + 1)),
        ("a held claim's instant", lambda d: d["claims"][ok]["outcome"].__setitem__("due_us", d["claims"][ok]["outcome"]["due_us"] - 1)),
        ("a refused claim stated to hold", lambda d: d["claims"][off].__setitem__("outcome", {"due_us": 3000000})),
        ("a held claim's bound", lambda d: d["claims"][ok].__setitem__("max_delay_ms", d["claims"][ok]["delay_ms"] - 1)),
        ("a control's instant", lambda d: d["scripts"][1]["controls"][1].__setitem__(1, 3000001)),
        ("a control's action", lambda d: d["scripts"][0]["controls"][0].__setitem__(0, PAUSE)),
        ("the room's t0", lambda d: d.__setitem__("t0_us", d["t0_us"] + 1000)),
        ("a change's controls", lambda d: d["changes"][0]["controls"][1].__setitem__(1, d["changes"][0]["controls"][1][1] + 1)),
        ("a change's instant after", lambda d: d["changes"][0].__setitem__("due_after_us", d["changes"][0]["due_before_us"])),
        ("a property's bytes", lambda d: d["properties"][0].__setitem__("delay_ms", 201)),
        ("a malformed property made good", lambda d: d["malformed"][0].__setitem__("immutable", d["properties"][0]["immutable"])),
        ("a short timestamp", lambda d: d["short"][1].__setitem__("full", d["short"][1]["full"] + SHORT)),
        ("a refused member stated to read", lambda d: d["catalog"][bad_cat].__setitem__("outcome", "accept")),
        ("an aligned capture time", lambda d: d["catalog"][2]["aligned"].__setitem__(2, 0)),
        ("a sink's target", lambda d: d["targets"][0].__setitem__("sink_unsynced", d["targets"][0]["sink_unsynced"] + 1)),
    ]
    for what, e in edits:
        changed = copy.deepcopy(doc)
        e(changed)
        if not check(changed):
            fails.append("a change went unnoticed: %s" % what)
    for f in fails:
        print("FAIL " + f)
    print("tempo: %d claims under %d scripts, %d changes of delay, %d catalog members, %d edits; %d failure(s)" % (len(doc["claims"]), len(doc["scripts"]), len(doc["changes"]), len(doc["catalog"]), len(edits), len(fails)))
    return 1 if fails else 0


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