#!/usr/bin/env python3
"""Pratt primality certificates in crux's binary TLV format.

leaf (prime 2):  00
node (prime p):  01 | u32 n | n x (u32 len | child) | u32 len | p bytes | u32 len | a bytes
where children certify the prime factors q of p-1 (with multiplicity) and a is a
witness: a^(p-1) == 1 (mod p) and a^((p-1)/q) != 1 for every distinct q.
"""
import struct

from sympy import factorint, isprime

u32 = lambda v: struct.pack("<I", v)


def ib(n, order):
    b = n.to_bytes(max(1, (n.bit_length() + 7) // 8), "big")
    return b if order == "big" else b[::-1]


def witness(p):
    fs = factorint(p - 1)
    for a in range(2, p):
        if pow(a, p - 1, p) == 1 and all(pow(a, (p - 1) // q, p) != 1 for q in fs):
            return a
    raise ValueError(p)


def cert(p, order="big", a=None):
    assert isprime(p)
    if p == 2:
        return b"\x00"
    kids = []
    for q, e in sorted(factorint(p - 1).items()):
        kids += [cert(q, order)] * e
    a = a or witness(p)
    A, B = ib(p, order), ib(a, order)
    return (b"\x01" + u32(len(kids)) + b"".join(u32(len(k)) + k for k in kids)
            + u32(len(A)) + A + u32(len(B)) + B)