// 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.
module mldsa

import crypto.internal.subtle
import crypto.sha3

// algo. 22: pkEncode (s. 7.2)
@[direct_array_access]
fn pk_encode(rho []u8, t1 [][]u16, p Params) []u8 {
	mut pk := rho.clone()
	for i in 0 .. p.k {
		w := t1[i]
		mut j := 0
		for j < n {
			c0 := w[j]
			c1 := w[j + 1]
			c2 := w[j + 2]
			c3 := w[j + 3]
			pk << u8(c0)
			pk << u8((c0 >> 8) | (c1 << 2))
			pk << u8((c1 >> 6) | (c2 << 4))
			pk << u8((c2 >> 4) | (c3 << 6))
			pk << u8(c3 >> 2)
			j += 4
		}
	}
	return pk
}

// algo. 23: pkDecode (s. 7.2)
@[direct_array_access]
fn pk_decode(pk []u8, p Params) !([]u8, [][]u16) {
	expected := pub_key_size(p)
	if pk.len != expected {
		return error('invalid public key length')
	}
	rho := pk[..32].clone()
	// avoid cloning on each iteration
	mut data := unsafe { pk[32..] }
	mut t1 := [][]u16{len: p.k, init: []u16{len: n}}
	for r in 0 .. p.k {
		mut j := 0
		for j < n {
			b0 := data[0]
			b1 := data[1]
			b2 := data[2]
			b3 := data[3]
			b4 := data[4]
			t1[r][j] = u16(b0) | (u16(b1 & 0x03) << 8)
			t1[r][j + 1] = u16(b1 >> 2) | (u16(b2 & 0x0f) << 6)
			t1[r][j + 2] = u16(b2 >> 4) | (u16(b3 & 0x3f) << 4)
			t1[r][j + 3] = u16(b3 >> 6) | (u16(b4) << 2)
			data = unsafe { data[5..] }
			j += 4
		}
	}
	return rho, t1
}

fn compute_pk_hash(pk []u8) [64]u8 {
	return slice_to_64(sha3.shake256(pk, 64))
}

@[direct_array_access]
fn compute_t1_hat(t1 [][]u16) []NttElement {
	mut result := []NttElement{len: t1.len}
	for i in 0 .. t1.len {
		mut w := RingElement{}
		for j in 0 .. n {
			// panics if (t1[i][j] << d) >= q, c.f field.v
			// only called from pk_decode, which produces 10bit vals (<= 1023)
			// so 1023 << 13 = 8_380_416 = q - 1
			z := field_to_montgomery(u32(t1[i][j]) << d) or { panic(err) }
			w[j] = z
		}
		result[i] = ntt(w)
	}
	return result
}

// algo. 35: Power2Round (s. 7.4)
fn power2_round(r FieldElement) (u16, FieldElement) {
	rr_ := field_from_montgomery(r)
	r1 := (rr_ + (1 << 12) - 1) >> d
	r0 := field_sub_to_montgomery(rr_, r1 << d)
	return u16(r1), r0
}

// algo. 37: HighBits (s. 7.4)
@[direct_array_access]
fn high_bits(r RingElement, p Params) [256]u8 {
	mut w := [256]u8{}
	match p.gamma2 {
		32 {
			for i in 0 .. n {
				w[i] = high_bits_32(field_from_montgomery(r[i]))
			}
		}
		88 {
			for i in 0 .. n {
				w[i] = high_bits_88(field_from_montgomery(r[i]))
			}
		}
		else {
			panic('mldsa: unsupported gamma2') // unreachable
		}
	}

	return w
}

fn high_bits_32(x u32) u8 {
	mut r1 := (x + 127) >> 7 // approx div by 2*gamma2
	r1 = (r1 * 1025 + (1 << 21)) >> 22
	r1 &= 0xf
	return u8(r1)
}

fn high_bits_88(x u32) u8 {
	mut r1 := (x + 127) >> 7 // approx div by 2*gamma2
	r1 = (r1 * 11275 + (1 << 23)) >> 24
	r1 = ct_select_eq(r1, 44, 0, r1)
	return u8(r1)
}

// algo. 36: Decompose, gamma2 = (q-1)/32 (s. 7.4)
fn decompose_32(r FieldElement) (u8, i32) {
	x := field_from_montgomery(r)
	r1 := high_bits_32(x)
	r0 := i32(x) - i32(r1) * 2 * i32((q - 1) / 32)
	r0_adj := ct_select_leq(i32(q / 2 + 1), r0, r0 - i32(q), r0)
	return r1, r0_adj
}

// algo. 36: Decompose, gamma2 = (q-1)/88 (s. 7.4)
fn decompose_88(r FieldElement) (u8, i32) {
	x := field_from_montgomery(r)
	r1 := high_bits_88(x)
	r0 := i32(x) - i32(r1) * 2 * i32((q - 1) / 88)
	r0_adj := ct_select_leq(i32(q / 2 + 1), r0, r0 - i32(q), r0)
	return r1, r0_adj
}

@[direct_array_access]
fn low_bits_exceed_bound(w RingElement, bound u32, p Params) bool {
	match p.gamma2 {
		32 {
			for i in 0 .. n {
				_, r0 := decompose_32(w[i])
				if ct_abs(r0) >= bound {
					return true
				}
			}
		}
		88 {
			for i in 0 .. n {
				_, r0 := decompose_88(w[i])
				if ct_abs(r0) >= bound {
					return true
				}
			}
		}
		else {
			panic('mldsa: unsupported gamma2') // unreachable
		}
	}

	return false
}

// algo. 40: UseHint (s. 7.4)
@[direct_array_access]
fn use_hint(r RingElement, h [256]u8, p Params) [256]u8 {
	mut w := [256]u8{}
	match p.gamma2 {
		32 {
			for i in 0 .. n {
				w[i] = use_hint_32(r[i], h[i])
			}
		}
		88 {
			for i in 0 .. n {
				w[i] = use_hint_88(r[i], h[i])
			}
		}
		else {
			panic('mldsa: unsupported gamma2') // unreachable
		}
	}

	return w
}

fn use_hint_32(r FieldElement, hint u8) u8 {
	mut r1, r0 := decompose_32(r)
	if hint == 1 {
		if r0 > 0 {
			r1 = (r1 + 1) % 16
		} else {
			r1 = (r1 - 1) % 16
		}
	}
	return r1
}

fn use_hint_88(r FieldElement, hint u8) u8 {
	mut r1, r0 := decompose_88(r)
	if hint == 1 {
		if r0 > 0 {
			if r1 == 43 {
				r1 = 0
			} else {
				r1++
			}
		} else {
			if r1 == 0 {
				r1 = 43
			} else {
				r1--
			}
		}
	}
	return r1
}

// algo. 39: MakeHint (s. 7.4)
@[direct_array_access]
fn make_hint(ct0 RingElement, w RingElement, cs2 RingElement, p Params) ([256]u8, int) {
	mut h := [256]u8{}
	mut count := 0
	match p.gamma2 {
		32 {
			for i in 0 .. n {
				h[i] = make_hint_32(ct0[i], w[i], cs2[i])
				count += int(h[i])
			}
		}
		88 {
			for i in 0 .. n {
				h[i] = make_hint_88(ct0[i], w[i], cs2[i])
				count += int(h[i])
			}
		}
		else {
			panic('mldsa: unsupported gamma2') // unreachable
		}
	}

	return h, count
}

fn make_hint_32(ct0 FieldElement, w FieldElement, cs2 FieldElement) u8 {
	r_plus_z := field_sub(w, cs2)
	v1 := high_bits_32(field_from_montgomery(r_plus_z))
	r1 := high_bits_32(field_from_montgomery(field_add(r_plus_z, ct0)))
	return u8(1 - subtle.constant_time_byte_eq(v1, r1))
}

fn make_hint_88(ct0 FieldElement, w FieldElement, cs2 FieldElement) u8 {
	r_plus_z := field_sub(w, cs2)
	v1 := high_bits_88(field_from_montgomery(r_plus_z))
	r1 := high_bits_88(field_from_montgomery(field_add(r_plus_z, ct0)))
	return u8(1 - subtle.constant_time_byte_eq(v1, r1))
}

fn w1_encode_len(p Params) int {
	return match p.gamma2 {
		32 { 4 * n / 8 }
		88 { 6 * n / 8 }
		else { panic('mldsa: unsupported gamma2') } // unreachable
	}
}

// algo. 28: w1Encode (s. 7.2)
@[direct_array_access]
fn w1_encode(w [256]u8, p Params, mut buf []u8) {
	match p.gamma2 {
		32 {
			for i := 0; i < n; i += 2 {
				buf[i / 2] = w[i] | (w[i + 1] << 4)
			}
		}
		88 {
			for i := 0; i < n; i += 4 {
				buf[3 * i / 4] = w[i] | (w[i + 1] << 6)
				buf[3 * i / 4 + 1] = (w[i + 1] >> 2) | (w[i + 2] << 4)
				buf[3 * i / 4 + 2] = (w[i + 2] >> 4) | (w[i + 3] << 2)
			}
		}
		else {
			panic('mldsa: unsupported gamma2') // unreachable
		}
	}
}

// algo. 26: sigEncode (s. 7.2)
fn sig_encode(ch []u8, z []RingElement, h [][256]u8, p Params) []u8 {
	mut sig := ch.clone()
	for i in 0 .. z.len {
		sig << bit_pack(z[i], p)
	}
	sig << hint_encode(h, p)
	return sig
}

// algo. 27: sigDecode (s. 7.2)
@[direct_array_access]
fn sig_decode(sig []u8, p Params) !([]u8, []RingElement, [][256]u8) {
	expected := sig_size(p)
	if sig.len != expected {
		return error('invalid signature length')
	}
	ch_len := p.lambda / 4
	ch := sig[..ch_len].clone()
	mut offset := ch_len
	mut z := []RingElement{len: p.l}
	for i in 0 .. p.l {
		length := (p.gamma1 + 1) * n / 8
		z[i] = bit_unpack(sig[offset..offset + length], p)
		offset += length
	}
	h := hint_decode(sig[offset..], p)!
	return ch, z, h
}

// algo. 17: BitPack (s. 7.1)
fn bit_pack(r RingElement, p Params) []u8 {
	match p.gamma1 {
		17 { return bit_pack_18(r) }
		19 { return bit_pack_20(r) }
		else { panic('mldsa: unsupported gamma1') } // unreachable
	}
}

@[direct_array_access]
fn bit_pack_18(r RingElement) []u8 {
	mut v := []u8{len: 18 * n / 8}
	mut pos := 0
	for i := 0; i < n; i += 4 {
		w0 := u32(1 << 17) - u32(field_centered_mod(r[i]))
		w1 := u32(1 << 17) - u32(field_centered_mod(r[i + 1]))
		w2 := u32(1 << 17) - u32(field_centered_mod(r[i + 2]))
		w3 := u32(1 << 17) - u32(field_centered_mod(r[i + 3]))
		v[pos] = u8(w0)
		v[pos + 1] = u8(w0 >> 8)
		v[pos + 2] = u8(w0 >> 16)
		v[pos + 2] |= u8(w1 << 2)
		v[pos + 3] = u8(w1 >> 6)
		v[pos + 4] = u8(w1 >> 14)
		v[pos + 4] |= u8(w2 << 4)
		v[pos + 5] = u8(w2 >> 4)
		v[pos + 6] = u8(w2 >> 12)
		v[pos + 6] |= u8(w3 << 6)
		v[pos + 7] = u8(w3 >> 2)
		v[pos + 8] = u8(w3 >> 10)
		pos += 9
	}
	return v
}

@[direct_array_access]
fn bit_pack_20(r RingElement) []u8 {
	mut v := []u8{len: 20 * n / 8}
	mut pos := 0
	for i := 0; i < n; i += 2 {
		w0 := u32(1 << 19) - u32(field_centered_mod(r[i]))
		w1 := u32(1 << 19) - u32(field_centered_mod(r[i + 1]))
		v[pos] = u8(w0)
		v[pos + 1] = u8(w0 >> 8)
		v[pos + 2] = u8(w0 >> 16)
		v[pos + 2] |= u8(w1 << 4)
		v[pos + 3] = u8(w1 >> 4)
		v[pos + 4] = u8(w1 >> 12)
		pos += 5
	}
	return v
}

// algo. 19: BitUnpack (s. 7.1)
fn bit_unpack(v []u8, p Params) RingElement {
	match p.gamma1 {
		17 { return bit_unpack_18(v) }
		19 { return bit_unpack_20(v) }
		else { panic('mldsa: unsupported gamma1') } // unreachable
	}
}

@[direct_array_access]
fn bit_unpack_18(v []u8) RingElement {
	mut r := RingElement{}
	mut pos := 0
	for i := 0; i < n; i += 4 {
		w0 := u32(v[pos]) | (u32(v[pos + 1]) << 8) | (u32(v[pos + 2]) << 16)
		r[i] = field_sub_to_montgomery(u32(1 << 17), w0 & 0x3ffff)
		w1 := (u32(v[pos + 2]) >> 2) | (u32(v[pos + 3]) << 6) | (u32(v[pos + 4]) << 14)
		r[i + 1] = field_sub_to_montgomery(u32(1 << 17), w1 & 0x3ffff)
		w2 := (u32(v[pos + 4]) >> 4) | (u32(v[pos + 5]) << 4) | (u32(v[pos + 6]) << 12)
		r[i + 2] = field_sub_to_montgomery(u32(1 << 17), w2 & 0x3ffff)
		w3 := (u32(v[pos + 6]) >> 6) | (u32(v[pos + 7]) << 2) | (u32(v[pos + 8]) << 10)
		r[i + 3] = field_sub_to_montgomery(u32(1 << 17), w3 & 0x3ffff)
		pos += 9
	}
	return r
}

@[direct_array_access]
fn bit_unpack_20(v []u8) RingElement {
	mut r := RingElement{}
	mut pos := 0
	for i := 0; i < n; i += 2 {
		w0 := u32(v[pos]) | (u32(v[pos + 1]) << 8) | (u32(v[pos + 2]) << 16)
		r[i] = field_sub_to_montgomery(u32(1 << 19), w0 & 0xfffff)
		w1 := (u32(v[pos + 2]) >> 4) | (u32(v[pos + 3]) << 4) | (u32(v[pos + 4]) << 12)
		r[i + 1] = field_sub_to_montgomery(u32(1 << 19), w1 & 0xfffff)
		pos += 5
	}
	return r
}

// algo. 20: HintBitPack (s. 7.2)
@[direct_array_access]
fn hint_encode(h [][256]u8, p Params) []u8 {
	mut out := []u8{len: p.omega + p.k}
	mut idx := 0
	for i in 0 .. p.k {
		for j in 0 .. n {
			if h[i][j] != 0 {
				out[idx] = u8(j)
				idx++
			}
		}
		out[p.omega + i] = u8(idx)
	}
	return out
}

// algo. 21: HintBitUnpack (s. 7.2)
@[direct_array_access]
fn hint_decode(y []u8, p Params) ![][256]u8 {
	if y.len != p.omega + p.k {
		return error('invalid hint length')
	}
	mut h := [][256]u8{len: p.k, init: [256]u8{}}
	mut idx := u8(0)
	for i in 0 .. p.k {
		limit := y[p.omega + i]
		if limit < idx || limit > u8(p.omega) {
			return error('invalid hint limits')
		}
		first := idx
		for idx < limit {
			if idx > first && y[idx - 1] >= y[idx] {
				return error('invalid hint index order')
			}
			h[i][y[idx]] = 1
			idx++
		}
	}
	for i := int(idx); i < p.omega; i++ {
		if y[i] != 0 {
			return error('invalid hint padding')
		}
	}
	return h
}

@[direct_array_access]
fn bit_pack_slow(r RingElement, a int, b int) []u8 {
	bitlen := bits_len(u32(a + b))
	mut out := []u8{len: n * bitlen / 8}
	mut acc := u32(0)
	mut acc_bits := u32(0)
	mut vi := 0
	for i in 0 .. n {
		w := u32(int(b) - int(field_centered_mod(r[i])))
		acc |= w << acc_bits
		acc_bits += u32(bitlen)
		for acc_bits >= 8 {
			out[vi] = u8(acc)
			vi++
			acc >>= 8
			acc_bits -= 8
		}
	}
	if acc_bits > 0 {
		out[vi] = u8(acc)
	}
	return out
}

@[direct_array_access]
fn bit_unpack_slow(v []u8, a int, b int) !RingElement {
	bitlen := bits_len(u32(a + b))
	if v.len != n * bitlen / 8 {
		return error('mldsa: invalid input length for bit_unpack_slow')
	}

	mask := u32((1 << bitlen) - 1)
	max_value := u32(a + b)

	mut r := RingElement{}
	mut acc := u32(0)
	mut acc_bits := u32(0)
	mut vi := 0

	for i in 0 .. n {
		for acc_bits < u32(bitlen) {
			if vi < v.len {
				acc |= u32(v[vi]) << acc_bits
				vi++
				acc_bits += 8
			}
		}
		w := acc & mask
		if w > max_value {
			return error('mldsa: coefficient out of range')
		}
		r[i] = field_sub_to_montgomery(u32(b), w)
		acc >>= u32(bitlen)
		acc_bits -= u32(bitlen)
	}

	return r
}

fn bits_len(x u32) int {
	if x == 0 {
		return 0
	}
	mut v := x
	mut n_ := 0
	for v > 0 {
		v >>= 1
		n_++
	}
	return n_
}

// algo. 24: skEncode (s. 7.2)
fn sk_encode(rho []u8, capital_k [32]u8, tr [64]u8, s1 []NttElement, s2 []NttElement, t0 []NttElement, p Params) []u8 {
	mut out := []u8{}
	out << rho[..32]
	out << capital_k[..]
	out << tr[..]
	eta := int(p.eta)
	for i in 0 .. p.l {
		out << bit_pack_slow(inverse_ntt(s1[i]), eta, eta)
	}
	for i in 0 .. p.k {
		out << bit_pack_slow(inverse_ntt(s2[i]), eta, eta)
	}
	for i in 0 .. p.k {
		out << bit_pack_slow(inverse_ntt(t0[i]), 4095, 4096)
	}
	return out
}

// algo. 25: skDecode (s. 7.2)
fn sk_decode(sk []u8, p Params) !([]u8, [32]u8, [64]u8, []RingElement, []RingElement, []RingElement) {
	k, l, eta := p.k, p.l, p.eta
	if sk.len != priv_key_size(p) {
		return error('mldsa: invalid private key size')
	}
	rho := sk[..32].clone()
	capital_k := slice_to_32(sk[32..64])
	tr := slice_to_64(sk[64..128])
	mut offset := 128

	eta_len := n * bits_len(u32(eta * 2)) / 8
	mut s1 := []RingElement{len: l}
	for i in 0 .. l {
		s1[i] = bit_unpack_slow(sk[offset..offset + eta_len], eta, eta)!
		offset += eta_len
	}

	mut s2 := []RingElement{len: k}
	for i in 0 .. k {
		s2[i] = bit_unpack_slow(sk[offset..offset + eta_len], eta, eta)!
		offset += eta_len
	}

	t0_len := n * 13 / 8
	mut t0 := []RingElement{len: k}
	for i in 0 .. k {
		t0[i] = bit_unpack_slow(sk[offset..offset + t0_len], 4095, 4096)!
		offset += t0_len
	}

	return rho, capital_k, tr, s1, s2, t0
}