"""ML-DSA verification (FIPS 204), written from the standard's text with Python's integers and
nothing beyond the standard library: SHAKE128 and SHAKE256 are hashlib's (FIPS 202).

It verifies; it does not sign. It shares no code with wai-rs, which signs and verifies with a
separate implementation. Every algorithm is named by its number in FIPS 204:

- `verify(pk, m, sig, ctx)`         Algorithm 3, ML-DSA.Verify, with the context string `ctx`;
- `verify_internal(pk, mp, sig)`    Algorithm 8, ML-DSA.Verify_internal, over M' as bytes;
- `verify_mu(pk, mu, sig)`          Algorithm 8 from line 8, over a message representative μ
                                    computed elsewhere (the "external μ" interface).

Parameters are ML-DSA-65's (§4, Table 1); `PARAMS` holds the three sets so a test can name
another, and a key or signature of another length is refused (§3.6.2).

A verification takes a few milliseconds: Python's integers, no vectorisation. The matrix Â
(Algorithm 32) is cached per public key, as §3.6.3 allows.
"""
import hashlib

Q = 8380417
D = 13
N = 256
ZETA = 1753

# FIPS 204 §4, Table 1.
PARAMS = {
    "ML-DSA-44": dict(k=4, l=4, eta=2, tau=39, lam=128, gamma1=1 << 17, gamma2=(Q - 1) // 88, omega=80),
    "ML-DSA-65": dict(k=6, l=5, eta=4, tau=49, lam=192, gamma1=1 << 19, gamma2=(Q - 1) // 32, omega=55),
    "ML-DSA-87": dict(k=8, l=7, eta=2, tau=60, lam=256, gamma1=1 << 19, gamma2=(Q - 1) // 32, omega=75),
}
for _p in PARAMS.values():
    _p["beta"] = _p["tau"] * _p["eta"]


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


def pk_len(p):
    return 32 + 32 * p["k"] * (bitlen(Q - 1) - D)


def sig_len(p):
    return p["lam"] // 4 + p["l"] * 32 * (1 + bitlen(p["gamma1"] - 1)) + p["omega"] + p["k"]


def H(data, n):
    """SHAKE256 to n bytes (§3.7)."""
    return hashlib.shake_256(data).digest(n)


class _XOF:
    """Incremental squeezing (§3.7, SP 800-185's API): each call returns the next bytes of the
    one output stream. hashlib's `digest(n)` gives a prefix of that stream, so the stream is
    read again, longer, whenever more is needed."""

    def __init__(self, ctor, data):
        self.ctor, self.data, self.buf, self.at = ctor, data, b"", 0

    def squeeze(self, n):
        while self.at + n > len(self.buf):
            self.buf = self.ctor(self.data).digest(max(2 * len(self.buf), 1024))
        out = self.buf[self.at:self.at + n]
        self.at += n
        return out


# ---------------------------------------------------------------- §7.1, §7.2: encodings

def _bits(data):
    """BytesToBits (Algorithm 13), as one little-endian integer."""
    return int.from_bytes(data, "little")


def simple_bit_unpack(v, b):
    """Algorithm 18: 256 coefficients of bitlen(b) bits each."""
    c = bitlen(b)
    z = _bits(v)
    mask = (1 << c) - 1
    return [(z >> (i * c)) & mask for i in range(N)]


def bit_unpack(v, a, b):
    """Algorithm 19: coefficients b - x for x of bitlen(a + b) bits."""
    c = bitlen(a + b)
    z = _bits(v)
    mask = (1 << c) - 1
    return [b - ((z >> (i * c)) & mask) for i in range(N)]


def hint_bit_unpack(y, p):
    """Algorithm 21; None for ⊥."""
    k, omega = p["k"], p["omega"]
    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
    for i in range(index, omega):
        if y[i] != 0:
            return None
    return h


def pk_decode(pk, p):
    """Algorithm 23."""
    rho = pk[:32]
    w = 32 * (bitlen(Q - 1) - D)
    t1 = [simple_bit_unpack(pk[32 + i * w:32 + (i + 1) * w], (1 << (bitlen(Q - 1) - D)) - 1) for i in range(p["k"])]
    return rho, t1


def sig_decode(sig, p):
    """Algorithm 27; h is None for ⊥."""
    lam4 = p["lam"] // 4
    c_tilde = sig[:lam4]
    w = 32 * (1 + bitlen(p["gamma1"] - 1))
    z = [bit_unpack(sig[lam4 + i * w:lam4 + (i + 1) * w], p["gamma1"] - 1, p["gamma1"]) for i in range(p["l"])]
    h = hint_bit_unpack(sig[lam4 + p["l"] * w:], p)
    return c_tilde, z, h


def simple_bit_pack(w, b):
    """Algorithm 16."""
    c = bitlen(b)
    z = 0
    for i, x in enumerate(w):
        z |= x << (i * c)
    return z.to_bytes(32 * c, "little")


def w1_encode(w1, p):
    """Algorithm 28."""
    b = (Q - 1) // (2 * p["gamma2"]) - 1
    return b"".join(simple_bit_pack(w, b) for w in w1)


# ---------------------------------------------------------------- §7.3: sampling

def sample_in_ball(rho, p):
    """Algorithm 29."""
    tau = p["tau"]
    c = [0] * N
    x = _XOF(hashlib.shake_256, rho)
    s = x.squeeze(8)
    h = int.from_bytes(s, "little")
    for i in range(N - tau, N):
        j = x.squeeze(1)[0]
        while j > i:
            j = x.squeeze(1)[0]
        c[i] = c[j]
        c[j] = -1 if (h >> (i + tau - N)) & 1 else 1
    return c


def rej_ntt_poly(rho):
    """Algorithm 30, with CoeffFromThreeBytes (Algorithm 14)."""
    x = _XOF(hashlib.shake_128, rho)
    a = []
    while len(a) < N:
        s = x.squeeze(3)
        z = (s[2] & 0x7F) << 16 | s[1] << 8 | s[0]
        if z < Q:
            a.append(z)
    return a


def expand_a(rho, p):
    """Algorithm 32."""
    return [[rej_ntt_poly(rho + bytes([s, r])) for s in range(p["l"])] for r in range(p["k"])]


# ---------------------------------------------------------------- §7.4: hints

def decompose(r, p):
    """Algorithm 36."""
    g2 = 2 * p["gamma2"]
    rp = r % Q
    r0 = rp % g2
    if r0 > g2 // 2:
        r0 -= g2
    if rp - r0 == Q - 1:
        return 0, r0 - 1
    return (rp - r0) // g2, r0


def use_hint(h, r, p):
    """Algorithm 40."""
    m = (Q - 1) // (2 * p["gamma2"])
    r1, r0 = decompose(r, p)
    if h == 1 and r0 > 0:
        return (r1 + 1) % m
    if h == 1 and r0 <= 0:
        return (r1 - 1) % m
    return r1


# ---------------------------------------------------------------- §7.5, §7.6: the NTT

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


ZETAS = [pow(ZETA, _bitrev8(k), Q) for k in range(N)]  # zetas[k] = ζ^BitRev8(k) mod q (§7.5)


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


def ntt_inv(a):
    """Algorithm 42."""
    w = list(a)
    m, ln = 256, 1
    while ln < N:
        start = 0
        while start < N:
            m -= 1
            z = -ZETAS[m] % Q
            for j in range(start, start + ln):
                t = w[j]
                w[j] = (t + w[j + ln]) % Q
                w[j + ln] = (t - w[j + ln]) % Q
                w[j + ln] = z * w[j + ln] % Q
            start += 2 * ln
        ln *= 2
    f = 8347681  # 256^-1 mod q
    return [f * x % Q for x in w]


# ---------------------------------------------------------------- §5.3, §6.3: verification

_A_CACHE = {}


def _matrix(pk, rho, p):
    key = (bytes(pk), p["k"], p["l"])
    a = _A_CACHE.get(key)
    if a is None:
        a = _A_CACHE[key] = expand_a(rho, p)
    return a


def _inf_norm(v):
    """‖·‖∞ of coefficients in mod± representation (§2.3)."""
    m = 0
    for poly in v:
        for x in poly:
            x %= Q
            if x > Q // 2:
                x -= Q
            m = max(m, abs(x))
    return m


def verify_mu(pk, mu, sig, name="ML-DSA-65"):
    """Algorithm 8 with μ given (lines 6 and 7 done by the caller)."""
    p = PARAMS[name]
    if len(pk) != pk_len(p) or len(sig) != sig_len(p) or len(mu) != 64:
        return False
    rho, t1 = pk_decode(pk, p)
    c_tilde, z, h = sig_decode(sig, p)
    if h is None:
        return False
    a_hat = _matrix(pk, rho, p)
    c_hat = ntt(sample_in_ball(c_tilde, p))
    z_hat = [ntt(zi) for zi in z]
    t1_hat = [ntt([x << D for x in ti]) for ti in t1]
    w1 = []
    for r in range(p["k"]):
        acc = [0] * N
        for s in range(p["l"]):
            ar, zs = a_hat[r][s], z_hat[s]
            for j in range(N):
                acc[j] += ar[j] * zs[j]
        ct = t1_hat[r]
        w_approx = ntt_inv([(acc[j] - c_hat[j] * ct[j]) % Q for j in range(N)])
        w1.append([use_hint(h[r][j], w_approx[j], p) for j in range(N)])
    c_tilde2 = H(mu + w1_encode(w1, p), p["lam"] // 4)
    return _inf_norm(z) < p["gamma1"] - p["beta"] and c_tilde == c_tilde2


def verify_internal(pk, mp, sig, name="ML-DSA-65"):
    """Algorithm 8: tr = H(pk, 64), μ = H(tr ‖ M', 64), then the rest."""
    tr = H(bytes(pk), 64)
    return verify_mu(pk, H(tr + bytes(mp), 64), sig, name)


def verify(pk, m, sig, ctx=b"", name="ML-DSA-65"):
    """Algorithm 3: M' = 0 ‖ |ctx| ‖ ctx ‖ M."""
    if len(ctx) > 255:
        return False
    return verify_internal(pk, bytes([0, len(ctx)]) + bytes(ctx) + bytes(m), sig, name)