"""Tests for electrum.mldsa — ML-DSA-65 wrapper."""
import importlib.util
import os

_ELECTRUM_DIR = os.path.join(os.path.dirname(__file__), '..', 'electrum')

def _load(name, filename):
    spec = importlib.util.spec_from_file_location(name, os.path.join(_ELECTRUM_DIR, filename))
    mod = importlib.util.module_from_spec(spec)
    import sys
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod

# Load modules in dependency order, bypassing __init__.py
_load('electrum.mldsa_keygen', 'mldsa_keygen.py')
mldsa = _load('electrum.mldsa', 'mldsa.py')


def test_random_keygen_sizes():
    pk, sk = mldsa.mldsa_generate_keypair()
    assert len(pk) == 1952, f"pk should be 1952 bytes, got {len(pk)}"
    assert len(sk) == 4032, f"sk should be 4032 bytes, got {len(sk)}"


def test_seeded_keygen_deterministic():
    seed = bytes(range(32))
    pk1, sk1 = mldsa.mldsa_keygen_from_seed(seed)
    pk2, sk2 = mldsa.mldsa_keygen_from_seed(seed)
    assert pk1 == pk2, "seeded keygen should be deterministic (pk)"
    assert sk1 == sk2, "seeded keygen should be deterministic (sk)"
    assert len(pk1) == 1952
    assert len(sk1) == 4032


def test_seeded_keygen_different_seeds():
    pk1, sk1 = mldsa.mldsa_keygen_from_seed(b'\x00' * 32)
    pk2, sk2 = mldsa.mldsa_keygen_from_seed(b'\x01' * 32)
    assert pk1 != pk2, "different seeds should produce different keys"


def test_sign_verify():
    pk, sk = mldsa.mldsa_generate_keypair()
    msg = b"test message for ML-DSA-65"
    sig = mldsa.mldsa_sign(sk, msg)
    assert len(sig) == 3309, f"signature should be 3309 bytes, got {len(sig)}"
    assert mldsa.mldsa_verify(pk, msg, sig) is True


def test_sign_verify_seeded():
    seed = bytes(range(32))
    pk, sk = mldsa.mldsa_keygen_from_seed(seed)
    msg = b"seeded key signing test"
    sig = mldsa.mldsa_sign(sk, msg)
    assert mldsa.mldsa_verify(pk, msg, sig) is True


def test_wrong_message_fails():
    pk, sk = mldsa.mldsa_generate_keypair()
    sig = mldsa.mldsa_sign(sk, b"correct message")
    assert mldsa.mldsa_verify(pk, b"wrong message", sig) is False


def test_wrong_key_fails():
    pk1, sk1 = mldsa.mldsa_generate_keypair()
    pk2, sk2 = mldsa.mldsa_generate_keypair()
    sig = mldsa.mldsa_sign(sk1, b"test")
    assert mldsa.mldsa_verify(pk2, b"test", sig) is False


if __name__ == "__main__":
    tests = [f for f in dir() if f.startswith('test_')]
    for t in sorted(tests):
        print(f"  {t} ... ", end="", flush=True)
        globals()[t]()
        print("OK")
    print(f"\nAll {len(tests)} mldsa tests passed.")
