"""
Hidden Number Problem solver for SIREN's biased-nonce ECDSA.

Per signature i on message m_i:
    s_i = k_i^-1 (z_i + r_i * d) mod N
    k_i = A_i * 2^SUFFIX_BITS + x_i,   0 <= x_i < 2^SUFFIX_BITS
    A_i = public_pitch(m_i)  -- known/computable from SONG_ID + m_i

  =>  x_i == a_i + b_i * d   (mod N)
      b_i = s_i^-1 * r_i mod N
      a_i = s_i^-1 * z_i - A_i * 2^SUFFIX_BITS mod N

L = { (y_1,...,y_m) in Z^m : y_i == d*b_i (mod N) for some integer d }
  is a rank-m lattice. Using sample 1 as pivot (c_i = b_i * b_1^-1 mod N):
      row0 = (1, c_2, ..., c_m)
      row_i = N * e_i          for i = 2..m
  is a valid basis of L (standard HNF-style normalization).

We center x_i -> x_i' = x_i - 2^(t-1) so the target is symmetric, then
solve CVP(L, u) with u = (-a_1', ..., -a_m') via Kannan embedding
(append (u, S) as one more basis row, S ~ 2^(t-1), search the reduced
basis for a vector (e_1,...,e_m, +-S)). Recovering x_1 from any hit
lets us solve for d directly: d = (x_1 - a_1) * b_1^-1 mod N.
No unbounded coordinate (d itself) ever enters the lattice, which is
what makes this construction well-conditioned.
"""
import hashlib
from fpylll import IntegerMatrix, LLL, BKZ


def public_pitch(song_id, msg, pitch_bits, nbits):
    material = (song_id + ":" + msg).encode()
    h = int.from_bytes(hashlib.sha256(material).digest(), "big")
    return h >> (nbits - pitch_bits)


def msg_hash(msg, N):
    h = int.from_bytes(hashlib.sha256(msg.encode()).digest(), "big")
    return h % N


def build_samples(sigs, song_id, N, pitch_bits, suffix_bits):
    """sigs: list of (msg, r, s). Returns list of (a_i, b_i)."""
    samples = []
    for msg, r, s in sigs:
        z = msg_hash(msg, N)
        A = public_pitch(song_id, msg, pitch_bits, N.bit_length())
        s_inv = pow(s, -1, N)
        b_i = (s_inv * r) % N
        a_i = (s_inv * z - (A << suffix_bits)) % N
        samples.append((a_i, b_i))
    return samples


def _center(v, N):
    return v - N if v > N // 2 else v


def solve_hnp(samples, N, suffix_bits, block_size=0, pivot=0):
    """
    samples: list of (a_i, b_i) with x_i = a_i + b_i*d mod N, 0 <= x_i < 2^suffix_bits.
    Returns list of candidate d values (int) to verify against the pubkey.
    block_size: 0 disables BKZ (LLL only); otherwise runs BKZ with that block size.
    pivot: index of the sample used to normalize the basis (must have b_pivot != 0).
    """
    m = len(samples)
    if m < 3:
        return []

    half = 1 << (suffix_bits - 1)  # centering offset, also the embedding factor S

    order = [pivot] + [i for i in range(m) if i != pivot]
    a = [samples[i][0] for i in order]
    b = [samples[i][1] for i in order]

    a_c = [(ai - half) % N for ai in a]  # centered targets a_i'

    b1_inv = pow(b[0], -1, N)
    c = [(bi * b1_inv) % N for bi in b[1:]]  # c_2..c_m

    dim = m + 1  # m (lattice L) + 1 (Kannan embedding)
    Mat = IntegerMatrix(dim, dim)

    # row 0: (1, c_2, ..., c_m, 0)
    Mat[0, 0] = 1
    for j, cj in enumerate(c):
        Mat[0, 1 + j] = cj
    # rows 1..m-1: N * e_i for coords 2..m
    for i in range(1, m):
        Mat[i, i] = N
    # row m (target/embedding row): (-a_1', ..., -a_m', S)
    for j in range(m):
        Mat[m, j] = (-a_c[j]) % N
        Mat[m, j] -= N if Mat[m, j] > N // 2 else 0  # keep entries small/signed
    Mat[m, m] = half

    LLL.reduction(Mat)
    if block_size and block_size > 2:
        try:
            BKZ.reduction(Mat, BKZ.Param(block_size=block_size))
        except Exception:
            pass

    candidates = []
    b1 = b[0]
    a1 = a[0]
    for row in range(dim):
        last = Mat[row, m]
        if abs(last) != half:
            continue
        sign = 1 if last == half else -1
        e1 = sign * Mat[row, 0]
        x1_centered = -e1
        x1 = (x1_centered + half) % N
        d_candidate = ((x1 - a1) * pow(b1, -1, N)) % N
        candidates.append(d_candidate)
    return candidates


def verify_privkey(d, Qx, Qy, N):
    from ecdsa import SECP256k1

    G = SECP256k1.generator
    P = d * G
    return P.x() == Qx and P.y() == Qy