// 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.sha3

// algo. 30: RejNTTPoly (s. 7.3)
@[direct_array_access]
fn sample_ntt(rho []u8, s u8, r u8) NttElement {
    mut g := sha3.new_shake128()
    g.write(rho)
    g.write([s, r])

    mut a := NttElement{}
    mut j := 0
    mut buf := g.read(168)
    mut off := 0
    for j < n {
        if off + 2 >= buf.len {
            buf = g.read(168)
            off = 0
        }
        v := u32(buf[off]) | (u32(buf[off + 1]) << 8) | (u32(buf[off + 2]) << 16)
        off += 3
        candidate := v & 0x7fffff
        if candidate < q {
            a[j] = field_to_montgomery(candidate) or { continue }
            j++
        }
    }
    return a
}

// algo. 31: RejBoundedPoly (s. 7.3)
@[direct_array_access]
fn sample_bounded_poly(rho []u8, r u8, p Params) RingElement {
    mut h := sha3.new_shake256()
    h.write(rho)
    h.write([r, 0])

    mut a := RingElement{}
    mut j := 0
    mut buf := h.read(136)
    mut off := 0
    for {
        if off >= buf.len {
            buf = h.read(136)
            off = 0
        }
        z0 := buf[off] & 0x0f
        z1 := buf[off] >> 4
        off++

        coeff, ok := coeff_from_half_byte(z0, p)
        if ok {
            a[j] = coeff
            j++
        }
        if j >= n {
            break
        }

        coeff2, ok2 := coeff_from_half_byte(z1, p)
        if ok2 {
            a[j] = coeff2
            j++
        }
        if j >= n {
            break
        }
    }
    return a
}

// algo. 29: SampleInBall (s. 7.3)
@[direct_array_access]
fn sample_in_ball(rho []u8, p Params) RingElement {
    mut h := sha3.new_shake256()
    h.write(rho)
    s := h.read(8)

    // pre-read ~2x expected bytes to avoid per-byte allocations
    mut buf := h.read(p.tau * 2)
    mut off := 0

    mut c := RingElement{}
    for i := 256 - p.tau; i < 256; i++ {
        for {
            if off >= buf.len {
                buf = h.read(256)
                off = 0
            }
            j := buf[off]
            off++
            if j <= u8(i) {
                c[i] = c[j]
                bit_idx := i + p.tau - 256
                bit := (s[bit_idx / 8] >> (bit_idx % 8)) & 1
                if bit == 0 {
                    c[j] = mont_one
                } else {
                    c[j] = mont_minus_one
                }
                break
            }
        }
    }
    return c
}

// algo. 15: CoeffFromHalfByte (s. 7.1)
fn coeff_from_half_byte(b u8, p Params) (FieldElement, bool) {
    match p.eta {
        2 {
            if b > 14 {
                return FieldElement(0), false
            }
            quotient := (u32(b) * 0x3334) >> 16 // barrett b % 5
            remainder := u32(b) - quotient * 5
            return field_sub_to_montgomery(2, remainder), true
        }
        4 {
            if b > 8 {
                return FieldElement(0), false
            }
            return field_sub_to_montgomery(4, u32(b)), true
        }
        else {
            panic('mldsa: unsupported eta') // unreachable
        }
    }
}

// algo. 32: ExpandA (s. 7.3)
fn compute_matrix_a(rho []u8, p Params) []NttElement {
    mut a := []NttElement{len: p.k * p.l}
    for r in 0 .. p.k {
        for s in 0 .. p.l {
            a[r * p.l + s] = sample_ntt(rho, u8(s), u8(r))
        }
    }
    return a
}