"""ML-DSA-65, the module-lattice digital signature algorithm of FIPS 204, in plain Python.

Non-normative; see spec/08-fidelity.md §8.4.
This is the reference form of the standard, slow and clear, for the signature scheme `ml-dsa-65`: key generation
from a 32-byte seed, the deterministic signing variant (the hedging bytes are zero), and verification. Keys and
signatures are the standard's byte encodings, so a production verifier can use any FIPS 204 library. The pure
(non pre-hashed) variant with an empty context is used, as the standard's ML-DSA.Sign defines it.
"""

import hashlib

Q = 8380417
N = 256
D = 13
ZETA = 1753
K, L = 6, 5
ETA = 4
TAU = 49
LAMBDA = 192
GAMMA1 = 1 << 19
GAMMA2 = (Q - 1) // 32
BETA = TAU * ETA
OMEGA = 55
PK_SIZE = 32 + 32 * K * (23 - D)
SK_SIZE = 32 + 32 + 64 + 32 * ((L + K) * 4 + D * K)
SIG_SIZE = LAMBDA // 4 + L * 32 * 20 + OMEGA + K
F_INV256 = 8347681  # 256^-1 mod q


def _bitrev8(m):
    return int("{:08b}".format(m)[::-1], 2)


ZETAS = [pow(ZETA, _bitrev8(m), Q) for m in range(256)]


def _shake256(data, n):
    return hashlib.shake_256(data).digest(n)


def _shake128(data, n):
    return hashlib.shake_128(data).digest(n)


# --- bit packing (FIPS 204 §7.1) ----------------------------------------------------------------------------


def _bitlen(x):
    return x.bit_length()


def simple_bit_pack(w, b):
    """Coefficients in [0, b] as bitlen(b)-bit little-endian fields."""
    bits = _bitlen(b)
    acc = 0
    for i, c in enumerate(w):
        acc |= c << (bits * i)
    return acc.to_bytes(32 * bits, "little")


def simple_bit_unpack(v, b):
    bits = _bitlen(b)
    acc = int.from_bytes(v, "little")
    mask = (1 << bits) - 1
    return [(acc >> (bits * i)) & mask for i in range(N)]


def bit_pack(w, a, b):
    """Coefficients in [-a, b] stored as b - w_i in bitlen(a + b) bits."""
    bits = _bitlen(a + b)
    acc = 0
    for i, c in enumerate(w):
        acc |= (b - c) << (bits * i)
    return acc.to_bytes(32 * bits, "little")


def bit_unpack(v, a, b):
    bits = _bitlen(a + b)
    acc = int.from_bytes(v, "little")
    mask = (1 << bits) - 1
    return [b - ((acc >> (bits * i)) & mask) for i in range(N)]


def hint_bit_pack(h):
    y = bytearray(OMEGA + K)
    index = 0
    for i in range(K):
        for j in range(N):
            if h[i][j]:
                y[index] = j
                index += 1
        y[OMEGA + i] = index
    return bytes(y)


def hint_bit_unpack(y):
    h = [[0] * N for _ in range(K)]
    index = 0
    for i in range(K):
        if y[OMEGA + i] < index or y[OMEGA + i] > OMEGA:
            return None
        first = index
        while index < y[OMEGA + i]:
            if index > first and y[index - 1] >= y[index]:
                return None
            h[i][y[index]] = 1
            index += 1
    if any(y[i] != 0 for i in range(index, OMEGA)):
        return None
    return h


# --- arithmetic ---------------------------------------------------------------------------------------------


def _mod_pm(r, m):
    """r mod± m: the representative in (-m/2, m/2]."""
    r %= m
    return r - m if r > m // 2 else r


def power2round(r):
    rp = r % Q
    r0 = _mod_pm(rp, 1 << D)
    return (rp - r0) >> D, r0


def decompose(r):
    rp = r % Q
    r0 = _mod_pm(rp, 2 * GAMMA2)
    if rp - r0 == Q - 1:
        return 0, r0 - 1
    return (rp - r0) // (2 * GAMMA2), r0


def high_bits(r):
    return decompose(r)[0]


def low_bits(r):
    return decompose(r)[1]


def make_hint(z, r):
    return 1 if high_bits(r) != high_bits(r + z) else 0


def use_hint(h, r):
    m = (Q - 1) // (2 * GAMMA2)
    r1, r0 = decompose(r)
    if h and r0 > 0:
        return (r1 + 1) % m
    if h and r0 <= 0:
        return (r1 - 1) % m
    return r1


def _inf_norm(w):
    return max(abs(_mod_pm(c, Q)) for c in w)


def ntt(w):
    w = list(w)
    m, length = 0, 128
    while length >= 1:
        start = 0
        while start < N:
            m += 1
            z = ZETAS[m]
            for j in range(start, start + length):
                t = z * w[j + length] % Q
                w[j + length] = (w[j] - t) % Q
                w[j] = (w[j] + t) % Q
            start += 2 * length
        length //= 2
    return w


def ntt_inv(w):
    w = list(w)
    m, length = 256, 1
    while length < N:
        start = 0
        while start < N:
            m -= 1
            z = (-ZETAS[m]) % Q
            for j in range(start, start + length):
                t = w[j]
                w[j] = (t + w[j + length]) % Q
                w[j + length] = z * (t - w[j + length]) % Q
            start += 2 * length
        length *= 2
    return [c * F_INV256 % Q for c in w]


def _mul(a, b):
    return [x * y % Q for x, y in zip(a, b)]


def _add(a, b):
    return [(x + y) % Q for x, y in zip(a, b)]


def _sub(a, b):
    return [(x - y) % Q for x, y in zip(a, b)]


# --- sampling (FIPS 204 §7.3) -------------------------------------------------------------------------------


def _coeff_from_three_bytes(b0, b1, b2):
    z = ((b2 & 0x7F) << 16) | (b1 << 8) | b0
    return z if z < Q else None


def rej_ntt_poly(rho):
    out, j, n = [], 0, 3 * 300
    stream = _shake128(rho, n)
    pos = 0
    while j < N:
        if pos + 3 > len(stream):
            n *= 2
            stream = _shake128(rho, n)
        c = _coeff_from_three_bytes(stream[pos], stream[pos + 1], stream[pos + 2])
        pos += 3
        if c is not None:
            out.append(c)
            j += 1
    return out


def _coeff_from_half_byte(b):
    return 4 - b if b < 9 else None  # eta = 4


def rej_bounded_poly(rho):
    out, j, n = [], 0, 200
    stream = _shake256(rho, n)
    pos = 0
    while j < N:
        if pos >= len(stream):
            n *= 2
            stream = _shake256(rho, n)
        z = stream[pos]
        pos += 1
        z0 = _coeff_from_half_byte(z & 15)
        z1 = _coeff_from_half_byte(z >> 4)
        if z0 is not None:
            out.append(z0)
            j += 1
        if z1 is not None and j < N:
            out.append(z1)
            j += 1
    return out


def expand_a(rho):
    return [[rej_ntt_poly(rho + bytes([s, r])) for s in range(L)] for r in range(K)]


def expand_s(rho_prime):
    s1 = [rej_bounded_poly(rho_prime + r.to_bytes(2, "little")) for r in range(L)]
    s2 = [rej_bounded_poly(rho_prime + (r + L).to_bytes(2, "little")) for r in range(K)]
    return s1, s2


def expand_mask(rho, mu):
    c = 1 + _bitlen(GAMMA1 - 1)
    return [bit_unpack(_shake256(rho + (mu + r).to_bytes(2, "little"), 32 * c), GAMMA1 - 1, GAMMA1) for r in range(L)]


def sample_in_ball(rho):
    c = [0] * N
    n = 8 + TAU * 4
    stream = _shake256(rho, n)
    signs = int.from_bytes(stream[:8], "little")
    pos = 8
    for i in range(N - TAU, N):
        while True:
            if pos >= len(stream):
                n *= 2
                stream = _shake256(rho, n)
            j = stream[pos]
            pos += 1
            if j <= i:
                break
        c[i] = c[j]
        c[j] = 1 if (signs >> (i + TAU - N)) & 1 == 0 else Q - 1
    return c


# --- encodings (FIPS 204 §7.2) ------------------------------------------------------------------------------


def pk_encode(rho, t1):
    return rho + b"".join(simple_bit_pack(t, (1 << (23 - D)) - 1) for t in t1)


def pk_decode(pk):
    rho, rest = pk[:32], pk[32:]
    step = 32 * (23 - D)
    return rho, [simple_bit_unpack(rest[i * step:(i + 1) * step], (1 << (23 - D)) - 1) for i in range(K)]


def sk_encode(rho, k, tr, s1, s2, t0):
    out = rho + k + tr
    out += b"".join(bit_pack(s, ETA, ETA) for s in s1)
    out += b"".join(bit_pack(s, ETA, ETA) for s in s2)
    out += b"".join(bit_pack(t, (1 << (D - 1)) - 1, 1 << (D - 1)) for t in t0)
    return out


def sk_decode(sk):
    rho, k, tr = sk[:32], sk[32:64], sk[64:128]
    pos = 128
    s1, s2, t0 = [], [], []
    for _ in range(L):
        s1.append(bit_unpack(sk[pos:pos + 128], ETA, ETA))
        pos += 128
    for _ in range(K):
        s2.append(bit_unpack(sk[pos:pos + 128], ETA, ETA))
        pos += 128
    for _ in range(K):
        t0.append(bit_unpack(sk[pos:pos + 32 * D], (1 << (D - 1)) - 1, 1 << (D - 1)))
        pos += 32 * D
    return rho, k, tr, s1, s2, t0


def sig_encode(c_tilde, z, h):
    return c_tilde + b"".join(bit_pack(zi, GAMMA1 - 1, GAMMA1) for zi in z) + hint_bit_pack(h)


def sig_decode(sigma):
    c_tilde = sigma[:LAMBDA // 4]
    pos = LAMBDA // 4
    z = []
    for _ in range(L):
        z.append(bit_unpack(sigma[pos:pos + 32 * 20], GAMMA1 - 1, GAMMA1))
        pos += 32 * 20
    h = hint_bit_unpack(sigma[pos:])
    return c_tilde, z, h


def w1_encode(w1):
    return b"".join(simple_bit_pack(w, (Q - 1) // (2 * GAMMA2) - 1) for w in w1)


# --- the scheme (FIPS 204 §6) -------------------------------------------------------------------------------


def _matvec(a_hat, v_hat):
    out = []
    for i in range(K):
        acc = [0] * N
        for j in range(L):
            acc = _add(acc, _mul(a_hat[i][j], v_hat[j]))
        out.append(acc)
    return out


def keygen(seed):
    """ML-DSA.KeyGen_internal: (public key, secret key) bytes for a 32-byte seed."""
    if len(seed) != 32:
        raise ValueError("the seed is 32 bytes")
    h = _shake256(seed + bytes([K, L]), 128)
    rho, rho_prime, kk = h[:32], h[32:96], h[96:]
    a_hat = expand_a(rho)
    s1, s2 = expand_s(rho_prime)
    t = [_add(ntt_inv(row), s2i) for row, s2i in zip(_matvec(a_hat, [ntt(s) for s in s1]), s2)]
    t1, t0 = [], []
    for poly in t:
        hi, lo = zip(*(power2round(c) for c in poly))
        t1.append(list(hi))
        t0.append(list(lo))
    pk = pk_encode(rho, t1)
    tr = _shake256(pk, 64)
    return pk, sk_encode(rho, kk, tr, s1, s2, t0)


def public_key(seed):
    return keygen(seed)[0]


def _message_representative(message, ctx=b""):
    if len(ctx) > 255:
        raise ValueError("the context is at most 255 bytes")
    return bytes([0, len(ctx)]) + ctx + message


def sign(sk, message, ctx=b"", rnd=bytes(32)):
    """ML-DSA.Sign with the deterministic variant (rnd = 0): the 3309-byte signature."""
    rho, kk, tr, s1, s2, t0 = sk_decode(sk)
    s1_hat = [ntt(s) for s in s1]
    s2_hat = [ntt(s) for s in s2]
    t0_hat = [ntt(t) for t in t0]
    a_hat = expand_a(rho)
    mu = _shake256(tr + _message_representative(message, ctx), 64)
    rho_pp = _shake256(kk + rnd + mu, 64)
    kappa = 0
    while True:
        y = expand_mask(rho_pp, kappa)
        w = [ntt_inv(row) for row in _matvec(a_hat, [ntt(yi) for yi in y])]
        w1 = [[high_bits(c) for c in wi] for wi in w]
        c_tilde = _shake256(mu + w1_encode(w1), LAMBDA // 4)
        c = sample_in_ball(c_tilde)
        c_hat = ntt(c)
        cs1 = [ntt_inv(_mul(c_hat, s)) for s in s1_hat]
        cs2 = [ntt_inv(_mul(c_hat, s)) for s in s2_hat]
        z = [_add(yi, ci) for yi, ci in zip(y, cs1)]
        r0 = [[low_bits(c) for c in _sub(wi, ci)] for wi, ci in zip(w, cs2)]
        kappa += L
        if max(_inf_norm(zi) for zi in z) >= GAMMA1 - BETA or max(max(abs(c) for c in r) for r in r0) >= GAMMA2 - BETA:
            continue
        ct0 = [ntt_inv(_mul(c_hat, t)) for t in t0_hat]
        h = []
        for i in range(K):
            base = _add(_sub(w[i], cs2[i]), ct0[i])
            h.append([make_hint((-ct0[i][j]) % Q, base[j]) for j in range(N)])
        if max(_inf_norm(t) for t in ct0) >= GAMMA2 or sum(sum(hi) for hi in h) > OMEGA:
            continue
        z_signed = [[_mod_pm(c, Q) for c in zi] for zi in z]
        return sig_encode(c_tilde, z_signed, h)


def verify(pk, message, signature, ctx=b""):
    """ML-DSA.Verify: True for a valid signature of `message` under `pk`. Never raises."""
    try:
        if len(pk) != PK_SIZE or len(signature) != SIG_SIZE or len(ctx) > 255:
            return False
        rho, t1 = pk_decode(pk)
        c_tilde, z, h = sig_decode(signature)
        if h is None:
            return False
        a_hat = expand_a(rho)
        tr = _shake256(pk, 64)
        mu = _shake256(tr + _message_representative(message, ctx), 64)
        c_hat = ntt(sample_in_ball(c_tilde))
        az = _matvec(a_hat, [ntt(zi) for zi in z])
        ct1 = [_mul(c_hat, ntt([c << D for c in t])) for t in t1]
        w_approx = [ntt_inv(_sub(a, b)) for a, b in zip(az, ct1)]
        w1 = [[use_hint(h[i][j], w_approx[i][j]) for j in range(N)] for i in range(K)]
        if max(_inf_norm(zi) for zi in z) >= GAMMA1 - BETA:
            return False
        return _shake256(mu + w1_encode(w1), LAMBDA // 4) == c_tilde
    except (ValueError, TypeError, IndexError):
        return False