// vtest build: has_modern_openssl?
// regenerate go test vecs: v run testdata/gen.vsh [go-source-path]

module mldsa

import crypto.sha256
import encoding.hex
import json
import os

struct TestVec {
    kind       string
    seed       string
    msg        string
    pk_sha256  string
    sig_sha256 string
    context    string
}

const vecs_json = os.read_file(os.real_path(os.join_path(os.dir(@FILE), 'testdata', 'vectors.json'))) or {
    panic(err)
}
const test_vecs = json.decode([]TestVec, vecs_json) or { panic(err) }

fn parse_kind(s string) Kind {
    return match s {
        'ml_dsa_44' { Kind.ml_dsa_44 }
        'ml_dsa_65' { Kind.ml_dsa_65 }
        'ml_dsa_87' { Kind.ml_dsa_87 }
        else { panic('unknown kind: ${s}') }
    }
}

fn test_keygen_sign_verify() {
    assert test_vecs.len > 0, 'no test vectors loaded'

    for tv in test_vecs {
        kind := parse_kind(tv.kind)
        seed := hex.decode(tv.seed) or { panic(err) }
        msg := hex.decode(tv.msg) or { panic(err) }
        expected_pk_hash := hex.decode(tv.pk_sha256) or { panic(err) }
        expected_sig_hash := hex.decode(tv.sig_sha256) or { panic(err) }

        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()

        pk_hash := sha256.sum(pk.bytes())
        assert pk_hash[..] == expected_pk_hash, 'pk hash mismatch for ${tv.kind} seed=${tv.seed[..16]}...'

        sig := sk.sign(msg, deterministic: true, context: tv.context) or { panic(err) }
        sig_hash := sha256.sum(sig)
        assert sig_hash[..] == expected_sig_hash, 'sig hash mismatch for ${tv.kind} seed=${tv.seed[..16]}...'

        verified := pk.verify(msg, sig, context: tv.context) or { panic(err) }
        assert verified, 'verify returned false for ${tv.kind} seed=${tv.seed[..16]}...'
    }
}

fn test_verify_rejects_bad_signature() {
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x00}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()
        msg := 'deadbeef'.bytes()

        sig := sk.sign(msg, deterministic: true) or { panic(err) }

        mut bad_sig := sig.clone()
        bad_sig[10] ^= 0xff

        result := pk.verify(msg, bad_sig) or { false }
        assert result == false, 'verify should reject tampered sig for ${kind}'
    }
}

fn test_verify_rejects_wrong_message() {
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x01}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()
        msg := 'the beef is alive'.bytes()

        sig := sk.sign(msg, deterministic: true) or { panic(err) }

        result := pk.verify('I love strawberries'.bytes(), sig) or { false }
        assert result == false, 'verify should reject wrong message for ${kind}'
    }
}

fn test_verify_rejects_wrong_context() {
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x02}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()
        msg := 'very cool message'.bytes()

        sig := sk.sign(msg, deterministic: true, context: 'some context a') or { panic(err) }

        result := pk.verify(msg, sig, context: 'another context b') or { false }
        assert result == false, 'verify should reject wrong context for ${kind}'
    }
}

fn test_public_key_roundtrip() {
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x03}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()
        msg := 'pk roundtrip'.bytes()

        sig := sk.sign(msg, deterministic: true) or { panic(err) }

        pk2 := PublicKey.from_bytes(pk.bytes(), kind) or { panic(err) }
        assert pk.equal(&pk2), 'roundtripped public key not equal'

        verified := pk2.verify(msg, sig) or { panic(err) }
        assert verified, 'verify failed after public key roundtrip for ${kind}'
    }
}

fn test_prehash_sign_verify() {
    prehashes := [
        PreHash.sha2_256,
        .sha2_384,
        .sha2_512,
        .sha3_256,
        .sha3_512,
        .shake_128,
        .shake_256,
    ]
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x05}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        pk := sk.public_key()
        msg := 'prehash test message'.bytes()

        for ph in prehashes {
            sig := sk.sign(msg, deterministic: true, prehash: ph) or { panic(err) }
            verified := pk.verify(msg, sig, prehash: ph) or { panic(err) }
            assert verified, 'prehash verify failed for ${kind} ${ph}'

            // pure verify must reject a prehashed sig
            pure_result := pk.verify(msg, sig) or { false }
            assert pure_result == false, 'pure verify should reject prehash signature for ${kind} ${ph}'
        }
    }
}

fn test_field_to_montgomery_roundtrip() {
    for val in [u32(0), 1, 2, 100, 1000, q - 1] {
        m := field_to_montgomery(val) or { panic(err) }
        back := field_from_montgomery(m)
        assert back == val, 'roundtrip failed for ${val}: got ${back}'
    }
}

fn test_field_add_sub() {
    a := field_to_montgomery(100) or { panic(err) }
    b := field_to_montgomery(200) or { panic(err) }
    sum := field_add(a, b)
    assert field_from_montgomery(sum) == 300

    diff := field_sub(sum, b)
    assert field_from_montgomery(diff) == 100
}

fn test_field_mul() {
    a := field_to_montgomery(1000) or { panic(err) }
    b := field_to_montgomery(2000) or { panic(err) }
    prod := field_montgomery_mul(a, b)
    assert field_from_montgomery(prod) == (1000 * 2000) % q
}

fn test_ntt_inverse_ntt_roundtrip() {
    mut f := RingElement{}
    for i in 0 .. n {
        f[i] = field_to_montgomery(u32(i % 100)) or { panic(err) }
    }
    ntt_f := ntt(f)
    back := inverse_ntt(ntt_f)
    for i in 0 .. n {
        assert field_from_montgomery(back[i]) == field_from_montgomery(f[i]), 'NTT roundtrip failed at index ${i}'
    }
}

fn test_ntt_mul_is_polynomial_product() {
    // (1 + x)^2 ?= x^2 + 2x + 1
    mut a := RingElement{}
    a[0] = field_to_montgomery(1) or { panic(err) }
    a[1] = field_to_montgomery(1) or { panic(err) }

    a_ntt := ntt(a)
    prod_ntt := ntt_mul(a_ntt, a_ntt)
    prod := inverse_ntt(prod_ntt)

    assert field_from_montgomery(prod[0]) == 1, 'expected x^2'
    assert field_from_montgomery(prod[1]) == 2, 'expected 2x'
    assert field_from_montgomery(prod[2]) == 1, 'expected 1'

    for i in 3 .. n {
        assert field_from_montgomery(prod[i]) == 0, 'expected 0 at index ${i}, got ${field_from_montgomery(prod[i])}'
    }
}

fn test_power2_round() {
    for val in [u32(0), 1, 100, 1000, q / 2, q - 1] {
        r := field_to_montgomery(val) or { panic(err) }
        hi, lo := power2_round(r)
        reconstructed := field_add(field_to_montgomery(u32(hi) << d) or { panic(err) }, lo)
        assert field_from_montgomery(reconstructed) == val, 'power2_round failed for ${val}'
    }
}

fn test_private_key_roundtrip() {
    for kind in [Kind.ml_dsa_44, .ml_dsa_65, .ml_dsa_87] {
        seed := []u8{len: 32, init: 0x04}
        sk := PrivateKey.from_seed(seed, kind) or { panic(err) }
        sk2 := PrivateKey.from_seed(sk.seed(), kind) or { panic(err) }
        assert sk.equal(&sk2), 'roundtripped private key not equal for ${kind}'
    }
}