#!/usr/bin/env python3
"""
rehearsal_build.py -- builds the published rehearsal set for manual revision 2.0.

This is the sealing ceremony in miniature, with the one difference that matters:
here the factors are written down and published. A rehearsal protects nothing.

Produces
  rehearsal-lock.json      a 2,048-bit AB-L1 lock, four waypoints, 88 milestones, four tiny
                           capsules, the factors, the waypoint values and a relay proof per leg
  rehearsal-carrier.png    an 888 x 888 synthetic carrier holding one AB-R1 record

Needs: afterberry_ref.py beside it, 'cryptography', 'Pillow', and two safe primes
(openssl prime -generate -safe -bits 1024) in p.txt and q.txt.
"""

import json
import os
import sys

from PIL import Image, ImageDraw, ImageFont

import afterberry_ref as ab

LOCK_ID = "rehearsal-1"
LEG = 2_222_222                                   # squarings between waypoints
MILESTONES = 88                                   # the real lock will carry 8,888
MESSAGES = {
    1: "rehearsal, ring i. nothing of berry is in here. only proof that a door can be built before there is a house.",
    2: "rehearsal, ring ii. you waited, or you were clever. the lock cannot tell the difference and does not mind.",
    3: "rehearsal, ring iii. the real rings will be slower. bring a chair, or a civilisation.",
    4: "rehearsal, ring iv. this is where 2400 will be. hello from 2026. the eights say hi.",
}


def main(out_dir):
    p, q = (int(open(f).read()) for f in ("p.txt", "q.txt"))
    for prime in (p, q):
        assert ab.is_probable_prime(prime) and ab.is_probable_prime((prime - 1) // 2), "need safe primes"
    n = p * q
    bits = n.bit_length()
    assert bits == 2048

    # -- the lock ----------------------------------------------------------
    doc = {
        "profile": "AB-L1",
        "lock_id": LOCK_ID,
        "status": "REHEARSAL. The factors are published below. This lock protects nothing.",
        "bits": bits,
        "n_hex": "%x" % n,
        "waypoints": [{"ring": k, "t": str(k * LEG), "check_sha256": "00" * 32} for k in (1, 2, 3, 4)],
    }
    lock = ab.Lock(doc)
    doc["x_hex"] = "%x" % lock.x
    doc["lock_digest_sha256"] = lock.digest.hex()

    # the maker's shortcut: with the factors in hand every waypoint is immediate
    ys = {ring: y for ring, _, y in ab.door_capability(lock, p, q)}
    for w in doc["waypoints"]:
        w["check_sha256"] = lock.check(w["ring"], ys[w["ring"]]).hex()

    # milestones: evenly spaced checks along the chain, so that a runner who goes wrong --
    # or is handed a false checkpoint -- finds out within one interval
    phi = (p - 1) * (q - 1)
    interval = 4 * LEG // MILESTONES
    doc["milestones"] = {
        "interval": str(interval),
        "checks_sha256": [lock.milestone(j, pow(lock.x, pow(2, j * interval, phi), n)).hex()
                          for j in range(1, MILESTONES + 1)],
    }

    # a relay proof per leg, so the waypoints can be checked without either door
    proofs, a, done = [], lock.x, 0
    for w in doc["waypoints"]:
        b, T = ys[w["ring"]], int(w["t"]) - done
        ell = ab.hash_to_prime(lock, a, b, T)
        pi = lock.canon(pow(a, ((1 << T) // ell) % phi, n))
        assert ab.relay_verify(lock, a, b, T, pi)
        proofs.append({"from_t": str(done), "to_t": w["t"], "ell_hex": "%x" % ell, "pi_hex": "%x" % pi})
        a, done = b, int(w["t"])

    # -- four single-chunk capsules ---------------------------------------
    doc["capsules"] = []
    for ring, text in MESSAGES.items():
        capsule_id, plain = ring, text.encode("utf-8")
        core = ab.header_core(ring, capsule_id, 1, 1, len(plain), lock.digest)
        content_key = os.urandom(32)                               # used once, never stored
        kek = lock.capsule_kek(lock.ring_key(ring, ys[ring]), capsule_id)
        capsule = ab.seal_capsule(plain, content_key, kek, core)
        assert ab.open_capsule(capsule, kek) == plain
        doc["capsules"].append({
            "ring": ring, "capsule_id": capsule_id,
            "sha256": ab.sha256(capsule).hex(), "capsule_hex": capsule.hex(),
        })
        del content_key

    doc["rehearsal_only"] = {
        "p_hex": "%x" % p, "q_hex": "%x" % q,
        "y_hex": ["%x" % ys[k] for k in (1, 2, 3, 4)],
        "relay_proofs": proofs,
    }
    json.dump(doc, open(os.path.join(out_dir, "rehearsal-lock.json"), "w"), indent=1)

    # -- one record, one carrier -------------------------------------------
    # rehearsal capsules are unsharded, so the capsule file stands in for the shard file
    first = bytes.fromhex(doc["capsules"][0]["capsule_hex"])
    record = ab.record_build(photo=0, capsule=1, shard=1, wrapped_key=first[64:104],
                             shard_sha256=ab.sha256(first), flags=0b11)
    codeword = ab.rs_encode(record)

    size = 888
    img = Image.new("RGB", (size, size))
    px = img.load()
    noise = ab._Stream(ab.sha256(b"afterberry rehearsal carrier"))
    for y in range(size):
        for x in range(size):
            g = noise.next()
            base = 14 + (y * 22) // size
            px[x, y] = (base + (g & 7), base - 4 + ((g >> 3) & 7), base + 10 + ((g >> 6) & 7))
    draw = ImageDraw.Draw(img)
    font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", 640)
    draw.text((size // 2, size // 2 - 20), "8", font=font, fill=(60, 235, 170), anchor="mm")
    pixels = bytearray(img.tobytes())
    changed = ab.carrier_embed(pixels, size, size, 0, codeword)
    path = os.path.join(out_dir, "rehearsal-carrier.png")
    ab.png_write(path, size, size, pixels, codeword.hex())

    # read it back three ways
    w, h, back, text = ab.png_read(path)
    assert bytes(back) == bytes(pixels) == Image.open(path).convert("RGB").tobytes()
    assert ab.rs_decode(ab.carrier_extract(back, w, h, 0)) == record
    assert bytes.fromhex(text["afterberry"]) == codeword
    print("lock digest   ", lock.digest.hex())
    print("record        ", record.hex())
    print("samples moved ", changed)
    print("carrier sha256", ab.sha256(open(path, "rb").read()).hex())
    print("pixel digest  ", ab.pixel_digest(w, h, back))


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else ".")