"""The framing maps ser and inst: round trip, injectivity, and rejection.

Lemma 2.14 of the paper turns on two properties of Definition 2.12: that the
length fields make the decomposition of ser(Omega) unique, so that ser is
injective and inst inverts it, and that inst is partial, rejecting a string that
is not a framed instance, in place of returning some other triple. Both are
checked here on shapes chosen to stress the framing: empty serializations and
tapes, payloads whose leading bytes look like length fields, memory images at
both extremes, and every truncation of a real instance.
"""
import hashlib
import json
import os
import random
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from selfrep import ser, inst, lam, HOST_PATH, M_STAR, tau_star

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def shapes():
    rng = random.Random(5)
    host = open(HOST_PATH, "rb").read()
    m0 = bytes(256)
    m1 = bytes([255]) * 256
    mr = bytes(rng.randrange(256) for _ in range(256))
    # A payload whose first eight bytes are a plausible length field, and a tape
    # that begins the way a framed field begins: if inst read a length from the
    # wrong offset these would decode to the wrong split.
    trap = lam(1 << 40) + b"trap"
    out = [
        (b"", m0, b""),
        (b"", m0, b"\x00"),
        (b"\x00", m0, b""),
        (b"\x00" * 3, m1, b""),
        (trap, mr, b""),
        (b"", mr, trap),
        (trap, m0, trap),
        (bytes(range(256)), mr, bytes(range(256))),
        (host, bytes(M_STAR), tau_star(host, M_STAR)),
        (host, m0, b""),
    ]
    for _ in range(300):
        a = bytes(rng.randrange(256) for _ in range(rng.randrange(0, 300)))
        t = bytes(rng.randrange(256) for _ in range(rng.randrange(0, 300)))
        out.append((a, bytes(rng.randrange(256) for _ in range(256)), t))
    return out


def main() -> int:
    cases = shapes()
    rt_bad, digests = [], {}
    for sg, m, t in cases:
        b = ser(sg, m, t)
        if inst(b) != (sg, m, t):
            rt_bad.append((len(sg), len(t)))
        digests.setdefault(hashlib.sha256(b).hexdigest(), []).append((len(sg), len(m), len(t)))
    collisions = {k: v for k, v in digests.items() if len(v) > 1}
    print(f"  inst(ser(Omega)) = Omega on {len(cases)} instances: "
          f"{'exact' if not rt_bad else f'{len(rt_bad)} FAILED'}")
    print(f"  ser injective over them (distinct instances, distinct bytes): "
          f"{'yes' if not collisions else f'NO {list(collisions.values())[:2]}'}")

    # Truncations. A string shorter than the framed serialization and memory
    # image is not an instance and must be rejected; one that reaches at least
    # that far is the instance with a correspondingly shorter tape, and must
    # parse to exactly that.
    host = open(HOST_PATH, "rb").read()
    tau = tau_star(host, M_STAR)
    full = ser(host, bytes(M_STAR), tau)
    minimal = 16 + len(host) + 256
    probes = sorted({p for p in
                     [0, 1, 7, 8, 9, 15, 16, 17,
                      8 + len(host) - 1, 8 + len(host), 8 + len(host) + 7,
                      8 + len(host) + 8, minimal - 256, minimal - 1, minimal,
                      minimal + 1, len(full) - 1, len(full)]
                     if 0 <= p <= len(full)})
    wrongly_accepted, wrongly_rejected = [], []
    for k in probes:
        try:
            got = inst(full[:k])
        except (ValueError, IndexError, struct.error):
            if k >= minimal:
                wrongly_rejected.append(k)
            continue
        if k < minimal:
            wrongly_accepted.append((k, len(got[1])))
        elif got != (host, bytes(M_STAR), full[minimal:k]):
            wrongly_accepted.append((k, "wrong triple"))
    short = sum(1 for k in probes if k < minimal)
    print(f"  {short} truncations below a complete memory image rejected: "
          f"{'yes' if not wrongly_accepted else f'NO {wrongly_accepted}'}")
    print(f"  {len(probes) - short} longer prefixes parse to the instance with "
          f"that tape: {'yes' if not wrongly_rejected else f'NO {wrongly_rejected}'}")
    accepted = wrongly_accepted + wrongly_rejected

    # a memory field of the wrong length must be rejected
    wrong = []
    for size in (0, 1, 255, 257, 512):
        b = lam(len(host)) + host + lam(size) + bytes(size) + tau
        try:
            inst(b)
            wrong.append(size)
        except (ValueError, IndexError, struct.error):
            pass
    print(f"  memory field of length != 256 rejected: "
          f"{'yes' if not wrong else f'NO for {wrong}'}")

    # the full instance is accepted and recovers exactly
    exact = inst(full) == (host, bytes(M_STAR), tau)
    print(f"  the distributed instance parses back exactly: {exact}")

    ok = not rt_bad and not collisions and not accepted and not wrong and exact
    json.dump({"instances": len(cases), "roundtrip_failures": len(rt_bad),
               "digest_collisions": len(collisions),
               "truncations_accepted": len(accepted),
               "bad_memory_lengths_accepted": len(wrong),
               "distributed_instance_exact": exact},
              open(runs_path("paper_framing.json"), "w"), indent=1)
    return 0 if ok else 1


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