import importlib.util
import os
import unittest

# Load mldsa_keygen directly to avoid electrum/__init__.py heavy imports
_spec = importlib.util.spec_from_file_location(
    'mldsa_keygen',
    os.path.join(os.path.dirname(__file__), '..', 'electrum', 'mldsa_keygen.py'))
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
mldsa_keygen_from_seed = _mod.mldsa_keygen_from_seed


class TestMldsaKeygen(unittest.TestCase):
    def test_deterministic(self):
        seed = bytes(range(32))
        pk1, sk1 = mldsa_keygen_from_seed(seed)
        pk2, sk2 = mldsa_keygen_from_seed(seed)
        self.assertEqual(pk1, pk2)
        self.assertEqual(sk1, sk2)

    def test_key_sizes(self):
        pk, sk = mldsa_keygen_from_seed(bytes(32))
        self.assertEqual(len(pk), 1952)
        self.assertEqual(len(sk), 4032)

    def test_different_seeds(self):
        pk1, _ = mldsa_keygen_from_seed(bytes(32))
        pk2, _ = mldsa_keygen_from_seed(bytes([1] + [0]*31))
        self.assertNotEqual(pk1, pk2)

    def test_pk_starts_with_rho(self):
        import hashlib
        seed = bytes(range(32))
        pk, sk = mldsa_keygen_from_seed(seed)
        h = hashlib.shake_256(seed + bytes([6, 5]))
        rho = h.digest(128)[:32]
        self.assertEqual(pk[:32], rho)

    def test_sk_format(self):
        import hashlib
        seed = bytes(range(32))
        pk, sk = mldsa_keygen_from_seed(seed)
        # sk[0:32] = rho = pk[0:32]
        self.assertEqual(sk[:32], pk[:32])
        # sk[64:128] = tr = SHAKE256(pk, 64)
        tr = hashlib.shake_256(pk).digest(64)
        self.assertEqual(sk[64:128], tr)

    def test_interop_sign_verify(self):
        from pqcrypto.sign.ml_dsa_65 import sign, verify
        seed = bytes(range(32))
        pk, sk = mldsa_keygen_from_seed(seed)
        sig = sign(sk, b"test message")
        verify(pk, b"test message", sig)  # should not raise

    def test_all_zero_seed_interop(self):
        from pqcrypto.sign.ml_dsa_65 import sign, verify
        pk, sk = mldsa_keygen_from_seed(bytes(32))
        sig = sign(sk, b"zero seed test")
        verify(pk, b"zero seed test", sig)

    def test_multiple_seeds_interop(self):
        from pqcrypto.sign.ml_dsa_65 import sign, verify
        for i in range(5):
            seed = bytes([i] * 32)
            pk, sk = mldsa_keygen_from_seed(seed)
            msg = f"seed {i}".encode()
            sig = sign(sk, msg)
            verify(pk, msg, sig)


if __name__ == '__main__':
    unittest.main()
