// Copyright 2025 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Ported to V from Go's crypto/internal/fips140/mldsa.

// ML-DSA (Module-Lattice-Based Digital Signature Algorithm) per FIPS 204
// https://nvlpubs.nist.gov/nistpubs/fips/nist.fips.204.pdf

module mldsa

import crypto.rand
import crypto.sha3
import crypto.internal.subtle

@[direct_array_access]
fn slice_to_32(s []u8) [32]u8 {
	mut a := [32]u8{}
	for i in 0 .. 32 {
		a[i] = s[i]
	}
	return a
}

@[direct_array_access]
fn slice_to_64(s []u8) [64]u8 {
	mut a := [64]u8{}
	for i in 0 .. 64 {
		a[i] = s[i]
	}
	return a
}

pub struct PrivateKey {
	seed [32]u8
	pk   PublicKey
	s1   []NttElement // len = l
	s2   []NttElement // len = k
	t0   []NttElement // len = k
	k    [32]u8
}

pub struct PublicKey {
	raw []u8
	p   Params
	a   []NttElement // k*l matrix in NTT domain
	t1  []NttElement // len = k, NTT(t1 * 2^d)
	tr  [64]u8
}

// algo. 1: ML-DSA.KeyGen (s. 5.1)
pub fn PrivateKey.generate(kind Kind) !PrivateKey {
	return new_private_key(slice_to_32(rand.read(32)!), kind.params())
}

pub fn PrivateKey.from_seed(seed []u8, kind Kind) !PrivateKey {
	if seed.len != 32 {
		return error('invalid seed length')
	}
	return new_private_key(slice_to_32(seed), kind.params())
}

// from FIPS 204 semi-expanded encoding. seed() and equal() are
// meaningless on the result — use from_seed when possible.
pub fn PrivateKey.from_bytes(raw []u8, kind Kind) !PrivateKey {
	return new_private_key_from_bytes(raw, kind.params())
}

pub fn PublicKey.from_bytes(raw []u8, kind Kind) !PublicKey {
	return new_public_key(raw, kind.params())
}

pub fn (sk &PrivateKey) public_key() &PublicKey {
	return &sk.pk
}

pub fn (sk &PrivateKey) seed() []u8 {
	mut s := []u8{len: 32}
	for i in 0 .. 32 {
		s[i] = sk.seed[i]
	}
	return s
}

pub fn (sk &PrivateKey) bytes() []u8 {
	return sk_encode(sk.pk.raw[..32], sk.k, sk.pk.tr, sk.s1, sk.s2, sk.t0, sk.pk.p)
}

// seed-based constant-time comparison. not meaningful for from_bytes keys.
pub fn (sk &PrivateKey) equal(other &PrivateKey) bool {
	mut a := []u8{len: 32}
	mut b := []u8{len: 32}
	for i in 0 .. 32 {
		a[i] = sk.seed[i]
		b[i] = other.seed[i]
	}
	return sk.pk.p == other.pk.p && subtle.constant_time_compare(a, b) == 1
}

// constant-time comparison of the serialized key material. slower but works for from_bytes keys.
pub fn (sk &PrivateKey) equal_bytes(other &PrivateKey) bool {
	return sk.pk.p == other.pk.p && subtle.constant_time_compare(sk.bytes(), other.bytes()) == 1
}

// algo. 2/4: ML-DSA.Sign / HashML-DSA.Sign (s. 5.2, 5.4.1)
pub fn (sk &PrivateKey) sign(msg []u8, opts SignerOpts) ![]u8 {
	if opts.context.len > 255 {
		return error('context too long')
	}
	mu := if opts.prehash != .none {
		compute_mu_prehash(sk.pk.tr[..], msg, opts.context, opts.prehash)
	} else {
		compute_mu(sk.pk.tr[..], msg, opts.context)
	}
	if opts.deterministic {
		return sign_internal(sk, mu, [32]u8{})
	}
	return sign_internal(sk, mu, slice_to_32(rand.read(32)!))
}

// sign_mu signs a precomputed mu value with explicit randomness.
// mu must be 64 bytes. rnd must be 32 bytes (use all zeros for deterministic signing).
pub fn (sk &PrivateKey) sign_mu(mu []u8, rnd []u8) ![]u8 {
	if mu.len != 64 {
		return error('mu must be 64 bytes')
	}
	if rnd.len != 32 {
		return error('rnd must be 32 bytes')
	}
	return sign_internal(sk, slice_to_64(mu), slice_to_32(rnd))
}

pub fn (pk &PublicKey) bytes() []u8 {
	return pk.raw.clone()
}

// tr returns the 64-byte transcript hash (H(pk)) used in mu computation.
pub fn (pk &PublicKey) tr() []u8 {
	return pk.tr[..]
}

pub fn (pk &PublicKey) equal(other &PublicKey) bool {
	return pk.p == other.p && subtle.constant_time_compare(pk.raw, other.raw) == 1
}

// algo. 3/5: ML-DSA.Verify / HashML-DSA.Verify (s. 5.3, 5.4.1)
pub fn (pk &PublicKey) verify(msg []u8, sig []u8, opts SignerOpts) !bool {
	if opts.context.len > 255 {
		return error('context too long')
	}
	mu := if opts.prehash != .none {
		compute_mu_prehash(pk.tr[..], msg, opts.context, opts.prehash)
	} else {
		compute_mu(pk.tr[..], msg, opts.context)
	}
	return verify_internal(pk, mu, sig)
}

pub fn (pk &PublicKey) verify_mu(mu []u8, sig []u8) !bool {
	if mu.len != 64 {
		return error('mu must be exactly 64 bytes')
	}
	return verify_internal(pk, slice_to_64(mu), sig)
}

// algo. 6: ML-DSA.KeyGen_internal (s. 6.1)
fn new_private_key(seed [32]u8, p Params) PrivateKey {
	k, l := p.k, p.l

	// expand seed into rho, rho', K
	mut xi := sha3.new_shake256()
	xi.write(seed[..])
	xi.write([u8(k), u8(l)])
	rho := xi.read(32)
	rho_s := xi.read(64)
	k_bytes := xi.read(32)

	a := compute_matrix_a(rho, p)

	mut s1 := []NttElement{len: l}
	for r in 0 .. l {
		s1[r] = ntt(sample_bounded_poly(rho_s, u8(r), p))
	}
	mut s2 := []NttElement{len: k}
	for r in 0 .. k {
		s2[r] = ntt(sample_bounded_poly(rho_s, u8(l + r), p))
	}

	// t_hat = A_hat * s1_hat + s2_hat
	mut t_hat := []NttElement{len: k}
	for i in 0 .. k {
		t_hat[i] = s2[i]
		for j in 0 .. l {
			t_hat[i] = poly_add_ntt(t_hat[i], ntt_mul(a[i * l + j], s1[j]))
		}
	}

	mut t1 := [][]u16{len: k, init: []u16{len: n}}
	mut t0 := []NttElement{len: k}
	for i in 0 .. k {
		t_i := inverse_ntt(t_hat[i])
		mut w := RingElement{}
		for j in 0 .. n {
			t1[i][j], w[j] = power2_round(t_i[j])
		}
		t0[i] = ntt(w)
	}

	pk_bytes := pk_encode(rho, t1, p)
	tr := compute_pk_hash(pk_bytes)
	t1_hat := compute_t1_hat(t1)

	k_arr := slice_to_32(k_bytes)

	return PrivateKey{
		seed: seed
		pk:   PublicKey{
			raw: pk_bytes
			p:   p
			a:   a
			t1:  t1_hat
			tr:  tr
		}
		s1:   s1
		s2:   s2
		t0:   t0
		k:    k_arr
	}
}

fn new_private_key_from_bytes(sk []u8, p Params) !PrivateKey {
	k, l := p.k, p.l

	rho, capital_k, tr, s1_ring, s2_ring, t0_ring := sk_decode(sk, p)!

	a := compute_matrix_a(rho, p)

	mut s1 := []NttElement{len: l}
	for r in 0 .. l {
		s1[r] = ntt(s1_ring[r])
	}
	mut s2 := []NttElement{len: k}
	for r in 0 .. k {
		s2[r] = ntt(s2_ring[r])
	}
	mut t0 := []NttElement{len: k}
	for r in 0 .. k {
		t0[r] = ntt(t0_ring[r])
	}

	// recompute t1 from rho, s1, s2 to verify consistency
	mut t1 := [][]u16{len: k, init: []u16{len: n}}
	for i in 0 .. k {
		mut t_hat := s2[i]
		for j in 0 .. l {
			t_hat = poly_add_ntt(t_hat, ntt_mul(a[i * l + j], s1[j]))
		}
		t_i := inverse_ntt(t_hat)
		for j in 0 .. n {
			r1, r0 := power2_round(t_i[j])
			t1[i][j] = r1
			if r0 != t0_ring[i][j] {
				return error('mldsa: private key inconsistent with t0')
			}
		}
	}

	pk_bytes := pk_encode(rho, t1, p)
	computed_tr := compute_pk_hash(pk_bytes)
	if computed_tr != tr {
		return error('mldsa: private key inconsistent with public key hash')
	}
	t1_hat := compute_t1_hat(t1)

	// use random bytes for seed since the semi-expanded format doesn't contain it
	seed := slice_to_32(rand.read(32)!)

	return PrivateKey{
		seed: seed
		pk:   PublicKey{
			raw: pk_bytes
			p:   p
			a:   a
			t1:  t1_hat
			tr:  tr
		}
		s1:   s1
		s2:   s2
		t0:   t0
		k:    capital_k
	}
}

fn new_public_key(raw []u8, p Params) !PublicKey {
	k, l := p.k, p.l

	rho, t1 := pk_decode(raw, p)!
	a := compute_matrix_a(rho, p)
	tr := compute_pk_hash(raw)
	t1_hat := compute_t1_hat(t1)

	return PublicKey{
		raw: raw.clone()
		p:   p
		a:   a[..k * l].clone()
		t1:  t1_hat[..k].clone()
		tr:  tr
	}
}

// algo. 2, lines 10-11: M' = 0x00 || |ctx| || ctx || M; mu = H(tr || M', 64)
// compute_mu computes mu = H(tr || M', 64) where M' = 0x00 || |ctx| || ctx || msg.
pub fn compute_mu(tr []u8, msg []u8, context string) [64]u8 {
	mut h := sha3.new_shake256()
	h.write(tr)
	h.write([u8(0)]) // pure mode domain sep
	h.write([u8(context.len)])
	h.write(context.bytes())
	h.write(msg)
	return slice_to_64(h.read(64))
}

// algo. 7: ML-DSA.Sign_internal (s. 6.2)
@[direct_array_access]
fn sign_internal(sk &PrivateKey, mu [64]u8, random [32]u8) ![]u8 {
	p := sk.pk.p
	k, l := p.k, p.l
	a := sk.pk.a
	s1 := sk.s1
	s2 := sk.s2
	t0 := sk.t0

	beta := u32(p.tau * p.eta)
	gamma1 := u32(1) << p.gamma1
	gamma1_beta := gamma1 - beta
	gamma2 := (q - 1) / u32(p.gamma2)
	gamma2_beta := gamma2 - beta

	// line 7: rho'' = H(K || rnd || mu, 64)
	mut h_nonce := sha3.new_shake256()
	h_nonce.write(sk.k[..])
	h_nonce.write(random[..])
	h_nonce.write(mu[..])
	nonce := h_nonce.read(64)

	mut kappa := 0

	mut y := []RingElement{len: l}
	mut y_hat := []NttElement{len: l}
	mut w := []RingElement{len: k}
	mut cs1 := []RingElement{len: l}
	mut cs2 := []RingElement{len: k}
	mut z := []RingElement{len: l}
	mut ct0 := []RingElement{len: k}
	mut h := [][256]u8{len: k, init: [256]u8{}}
	mut w1_buf := []u8{len: w1_encode_len(p)}

	// lines 10-32: rejection sampling loop (bounded by max_sign_attempts)
	for _ in 0 .. max_sign_attempts {
		// line 11: y = ExpandMask(rho'', kappa) (algo. 34)
		for r in 0 .. l {
			counter := [u8(kappa & 0xff), u8(kappa >> 8)]
			kappa++

			mut h_y := sha3.new_shake256()
			h_y.write(nonce)
			h_y.write(counter)
			v_bytes := h_y.read((p.gamma1 + 1) * n / 8)
			y[r] = bit_unpack(v_bytes, p)
		}

		// line 12: w = NTT^-1(A_hat * NTT(y))
		for i in 0 .. l {
			y_hat[i] = ntt(y[i])
		}
		for i in 0 .. k {
			mut w_hat := NttElement{}
			for j in 0 .. l {
				w_hat = poly_add_ntt(w_hat, ntt_mul(a[i * l + j], y_hat[j]))
			}
			w[i] = inverse_ntt(w_hat)
		}

		// line 13-14: w1 = HighBits(w); c_tilde = H(mu || w1Encode(w1), lambda/4)
		mut h_ch := sha3.new_shake256()
		h_ch.write(mu[..])
		for i in 0 .. k {
			w1_encode(high_bits(w[i], p), p, mut w1_buf)
			h_ch.write(w1_buf)
		}
		ch := h_ch.read(p.lambda / 4)

		// line 15-16: c = SampleInBall(c_tilde); c_hat = NTT(c)
		c := ntt(sample_in_ball(ch, p))

		// lines 17-20: cs1 = NTT^-1(c_hat * s1_hat); z = y + cs1
		for i in 0 .. l {
			cs1[i] = inverse_ntt(ntt_mul(c, s1[i]))
		}
		for i in 0 .. k {
			cs2[i] = inverse_ntt(ntt_mul(c, s2[i]))
		}

		// line 23: ||z||_inf >= gamma1 - beta
		mut reject := false
		for i in 0 .. l {
			z[i] = poly_add_ring(y[i], cs1[i])
			if coefficients_exceed_bound(z[i], gamma1_beta) {
				reject = true
				break
			}
		}
		if reject {
			continue
		}

		// line 23: ||r0||_inf >= gamma2 - beta
		reject = false
		for i in 0 .. k {
			r0 := poly_sub_ring(w[i], cs2[i])
			if low_bits_exceed_bound(r0, gamma2_beta, p) {
				reject = true
				break
			}
		}
		if reject {
			continue
		}

		// line 25, 28: ct0 = NTT^-1(c_hat * t0_hat); ||ct0||_inf >= gamma2
		reject = false
		for i in 0 .. k {
			ct0[i] = inverse_ntt(ntt_mul(c, t0[i]))
			if coefficients_exceed_bound(ct0[i], gamma2) {
				reject = true
				break
			}
		}
		if reject {
			continue
		}

		// line 26, 28: h = MakeHint(-ct0, w - cs2 + ct0); count(h) > omega
		mut count1s := 0
		for i in 0 .. k {
			hint_result, count := make_hint(ct0[i], w[i], cs2[i], p)
			h[i] = hint_result
			count1s += count
		}
		if count1s > p.omega {
			continue
		}

		return sig_encode(ch, z, h, p) // line 33: sigEncode(c_tilde, z, h)
	}
	return error('signing failed: rejection sampling did not converge after ${max_sign_attempts} attempts')
}

// algo. 8: ML-DSA.Verify_internal (s. 6.3)
@[direct_array_access]
fn verify_internal(pk &PublicKey, mu [64]u8, sig []u8) !bool {
	p := pk.p
	k, l := p.k, p.l
	t1 := pk.t1
	a := pk.a

	beta := u32(p.tau * p.eta)
	gamma1 := u32(1) << p.gamma1
	gamma1_beta := gamma1 - beta

	ch, z, h := sig_decode(sig, p) or { return false }

	c := ntt(sample_in_ball(ch, p))

	// line 9: w'_approx = NTT^-1(A_hat * NTT(z) - NTT(c) * NTT(t1 * 2^d))
	mut z_hat := []NttElement{len: l}
	for i in 0 .. l {
		z_hat[i] = ntt(z[i])
	}
	mut w := []RingElement{len: k}
	for i in 0 .. k {
		mut w_hat := NttElement{}
		for j in 0 .. l {
			w_hat = poly_add_ntt(w_hat, ntt_mul(a[i * l + j], z_hat[j]))
		}
		w_hat = poly_sub_ntt(w_hat, ntt_mul(c, t1[i]))
		w[i] = inverse_ntt(w_hat)
	}

	// line 10: w'1 = UseHint(h, w'_approx)
	mut w1 := [][256]u8{len: k, init: [256]u8{}}
	for i in 0 .. k {
		w1[i] = use_hint(w[i], h[i], p)
	}

	// line 12: c_tilde' = H(mu || w1Encode(w'1), lambda/4)
	mut h_ch := sha3.new_shake256()
	h_ch.write(mu[..])
	mut w1_buf := []u8{len: w1_encode_len(p)}
	for i in 0 .. k {
		w1_encode(w1[i], p, mut w1_buf)
		h_ch.write(w1_buf)
	}
	computed_ch := h_ch.read(p.lambda / 4)

	// line 13: ||z||_inf < gamma1 - beta and c_tilde == c_tilde'
	for i in 0 .. l {
		if coefficients_exceed_bound(z[i], gamma1_beta) {
			return false
		}
	}

	if subtle.constant_time_compare(ch, computed_ch) != 1 {
		return false
	}

	return true
}