package org.bouncycastle.pqc.crypto.test;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.SecureRandom;
import java.util.HashMap;
import java.util.Map;

import junit.framework.Assert;
import junit.framework.TestCase;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.SecretWithEncapsulation;
import org.bouncycastle.crypto.generators.MLKEMKeyPairGenerator;
import org.bouncycastle.crypto.kems.MLKEMExtractor;
import org.bouncycastle.crypto.kems.MLKEMGenerator;
import org.bouncycastle.crypto.params.MLKEMKeyGenerationParameters;
import org.bouncycastle.crypto.params.MLKEMParameters;
import org.bouncycastle.crypto.params.MLKEMPrivateKeyParameters;
import org.bouncycastle.crypto.params.MLKEMPublicKeyParameters;
import org.bouncycastle.crypto.util.DEROtherInfo;
import org.bouncycastle.crypto.util.OtherInfoGenerator;
import org.bouncycastle.crypto.util.PrivateKeyFactory;
import org.bouncycastle.crypto.util.PrivateKeyInfoFactory;
import org.bouncycastle.crypto.util.PublicKeyFactory;
import org.bouncycastle.crypto.util.SubjectPublicKeyInfoFactory;
import org.bouncycastle.internal.asn1.oiw.OIWObjectIdentifiers;
import org.bouncycastle.test.TestResourceFinder;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.FixedSecureRandom;

public class MLKEMTest
    extends TestCase
{
    private static byte[] okayKey = Hex.decode("29F853ECE08908E99D4B7195B89C296037B786E3000CA95607D80288A62EBDB53E5AD7A4BC716AE3BC0200788D7FBA9E49D5AAAA0A46E7A6C829A7B64869C458CACDB9043D350BC920D29E1D1B8E5ED57D83132AA6119C33770A56EB64511579FF2404F4CAA6CA8026BC647EA1692143139D4B5526006C38E87C1AC88CB615B87910E75371F97376938533B9720810B63D7354CE6B726184A392D90AF5F5821179C0ED4A519C4530818788CD877027369B5C8855486C2E8DE9BD7A2835A5D0248B772DFBE26BFA4049D0BC51EDCB9C51FAAFF638CB45399B8186C12E8119D5CBBCADE116FE1ACB097C8280343A3E666F6C5B89DB1070E698185D290429107712A46A3D057233B90FBC9A475CA73D3FF2743E54CEB255AF1809209D042E54D4C099D6C3B46868CBC64F00C5B352831C28C28260E40233CA62F5DC814C0809A9653D63CB756032BD9322544A199EF6E696AB026BA2701CE4D2A13850275BDA905A05130B2CB1E7410F46523DB54646E94C508CF548190A507174CA8B481BDC14C27B9267D8F428C0F73EBA769F35398824B2C4D6C171A1CA7A9DE0A315FA75ECD477B10969956B34F5188BA290390755066BFA8FB4E9BAE097A59243CE0DC57089E022C8D52A19D03F7EB23141B51A44BC1B1C0BA52A70AA1B8A61654CC54079A3AABC4D86E3C046880CDA9BCCE3A609E6A1149DFB6931E91D743BCBCDF6A963175FB48520FFD22885D0BCC59393F487133A5A0E115A02CD6049C06C8FF4D8A221E20D3F23C905501180F878C4D52C65C5A7FD745592CC38F5266918808B730138CEE47B0914544A653DDACB729581AC7296A7A5C90009DB63EDFAA45B6086BEA186983487A0E530D069AB99A10DC2F59894A77D15A7933AF23ECD283962E778BCAB985B2BA7DA84C0E10C165500A1A93427B976516E867A94095F30D406C209126789ADBE325D652ACD09B4B104A431ED4691F7A5CD15279A36E939A80C91C6468BF73811F103CA405C93CAC493245516ABCAC7B3F90CD2F030BE6461A18818C3983B03106D2D343D88877C41B9C329CB76B7DB2720D0157A04141459032AF4017D09D47B472F4CF643436ABF8D29AC9F7E27C7D9EB181F5AC93833D930C05D3E");
    private static byte[] tooLargeKey = Hex.decode("00615CA0D7AD5C0C84A7F77A9ACA3C32E17732C42EDC46249490021B9822D06361705ABFFC28730CE741B144A27A5748DE280CD253575EA6AF87C12E9ED6029825AB2F14C482E86DBB020257A25A33F38524858DBE4305A6B13859382615C6CCF7A997085384F8F0401F63127F5358D37236D978377365023FF659B3470195C43824441DA37641931041039093B185145C8644846B57EBF3862DC41AB7E54ABD0730BCC25B369C8B86E2165C73691F1863887071A77CAFFE39B537026991153A92D90E0568729F4C568895B965B8A1030A93DBA93672C2B9C5331491F442069550E4555F085B94F9F1611A028776361A6696671207B18803BE44395C7A560BD675B94ED0C0FFC97253E58FE2FBBBAACA802BB98F5E2B3DD810AE2CA91C00088B7D3175A2A65468A872917400163337BE503BEC2A24D798CF85A6653470AA5F0635447C94E73CA0FD09A2BD104432768F239B9633DC943BD052A8934D38D408BA38B57947699DF3783871CC23A46AA871BAA096A7DE77CF4320B2AF26A40F656351F90DFEB11088E85E6E181F42475324240288F13B5DE310E07C698914B8B9355416A69219D92315B9C3421205B0698D6AE3161AD18ECC325C3BC4BDE37002C4AA128C4C4FD4EAB040831264684ACA88A0D9C4616824085CE9A3A0C37C1148B41EA6C242B77E277C4A867113D9B3920354BE976C68ABD89E76819E7C3573F79A85D9A1778DA81D871838BBB43870CB786525267FF009B21A3275B7958936BABB7020D7F832362AA95E7761B9A82FD813266C709D4AC778632298F5800079D88A447C7A6B004843831775499C62A553BCF302FB0556CCA10577F24C81A651232BC9A5E97D37FC046498BCCFC69F7E4CBF890CC4DF5AA5E76C14371C38C5A348D1303C154C11899B7E9AA34D806275D6AB0E1E7B3A0BE6B48FC21475C4B5ED3B6AAFC23ADFBB7E33E501B7545B71446BABB6680A7316C39994DA351E36D228572080DFB47E0FBC12D24A842847109700309CF47B0F338389DAC6EFF534FA2251DA7773DD6845112843B95739D8F2109A91B078BAAC3612230D77569BA22138F9AB3FA02DFACA54C943920D254550C863261C771C32F6BC643A77D9C279F08AAC3E51F90DFEB11088E85E6E181F42475324240288F13B5DE310E07C698914B8B9355416A69219D92315B9C3421205B0698D6AE3161AD18ECC325C3BC4BDE37002C4AA128C4C4FD4EAB040831264684ACA88A0D9C4616824085CE9A3A0C37C1148B41EA6C242B77E277C4A867113D9B3920354BE976C68ABD89E76819E7C3573F79A85D9A1778DA81D871838BBB43870CB786525267FF009B21A3275B7958936BABB7020D7F832362AA95E7761B9A82FD813266C709D4AC778632298F5800079D88A447C7A6B004843831775499C62A553BCF302FB0556CCA10577F24C81A651232BC9A5E97D37FC046498BCCFC69F7E4CBF890CC4DF5AA5E76C14371C38C5A348D1303C154C11899B7E9AA34D806275D6AB0E1E7B3A0BE6B48FC21475C4B5ED3B6AAFC23ADFBB7E33E501B7545B71446BABB6680A7316C39994DA351E36D228572080DFB47E0FBC12D24A842847109700309CF47B0F338389DAC6EFF534FA2251DA7773DD6845112843B95739D8F2109A91B078BAAC3612230D77569BA22138F9AB3FA02DFACA54C943920D254550C863261C771C32F6BC643A77D9C279F08AAC3E");

    private static byte[] faultyPrivateKey = Hex.decode("F57A2F125726FF2B75F19321D55A7AA33142A1BA9B8C1A1DDF9068D4759BFDFAB942A0AA90710033846B195BA44FC1762CA26C8AC0639BC1ADC2DC0E49CC6559830AF9463401EA181FEC83BD69594E3911108A929E1C2AD0F65AD60C5DE509BE8A966FE1F3452097C0BC186F6DA467F84A4A64A7B657D06D96E1A411611D39A3726752C30E48244CD675B7AC7F993C2EF5CCC5C598123024AFD3403DAA4A1541E1A8EE797B06BA5A9A01C0D2340B0122A379FB8F1EF73D2E46117F6B6FF23468C1D43413D4B0760A49B7D6BFE0784D1A646209C719C9E66147467E940163D03A0D4106C019E743F864A809F9BFD2E8C40AD185A428B1AAEC7510ECCC039A62E2E27472F77B451A8C5BFC1AF9D3C632B006CEE13A9B6B8F3280889A0515CC54BD2D676BB930B0C440771D965E4CA92942F9705D4BA535D43D18F94957BA285C0B5BF0F40080B1B42E82446A64C3B2707ADB41954F803B72884A13EBABC2801385928FF3890607B3C60FC16155A51C20BCB3BF3ACD07559F4BA67CE349AD5BD86106D01C19E222EFF33768782BBFF10D07D42045C69F68126241C985D1E3C1E39195DC133B8E6913E17359A331706C882D82527BDC90A0C7B223B4C30537CCC986E17C09513B1FF9A026A80A9358CAFA97A0C6D65BDAF471880C5B7A69882DAA2AE5C9223024AF3AA38FDFB1960F2527F7C03AAEB24E071BB4AE7CC170B4C9A302CF6ECA4FFBC96B3A801FE643CA65605A3CE5BC0944C1E234BA9EB9C7294543064309E266108E21110C5251648BB0AB449832115DF6D5538004581B371E6C779AD3949B2E4BA041E218EB853D172341C426B0DD1C3CB1B3283F27C716C288C1CA3C6C4BAE3C8B7AFC480565B13FA13A0CED8B2C0BBC83D26C0EC81C9EA25B9663710E51A811D8E398F4BB761DD862BF30502873C04A3B412BE20C8C2819312245288AC6DF74AE1B85BA217ABE08F465B12CAFF664B3693CAA5688C05F2A2070C307DAFA7F37BBAB87F11708F04E10175A74D6ABD7D09D8A0862AC487AED60BE508B28A52B1FC394A176C6A89F68BCC2217B4B463EF631C9821B3D71216CDE75A7D8763665956670ECBAE9B13A73899ED68264E65298B0B8256B713FE62470DDD3443E103AA679B8A6B3A8D690A9A655BFD0AA0A1030C5C57A1E6419114555605D1197F45C9258C617B28983C4257B67F688087A77A0990E6D11982589BF1FD424CDF73B156896FA534530416D1AB75915B3460AA295216C362C48A2B9944D94F4A00AEB5A9991979BD9B7859C881B5B3254F59079DC2DB209CE73AAC00D945713D46A62D9013F909A7E88B925704669B5356F851A99F77181215D0A73351871C853847B7BB3715E4892FDFC9466A4C7EAA83F0A9468FBCA6684786C61607A6FA1875C6043092B3DD6D5127C92559A9157F6DC1CD999B86A096A5D1524A578CA36C667C3F092DA3AA889319669469648C4BA920B53843686A18A8246B3B5A659C936EA0D0B01147B5832214488139224F9595FDB9151AE3621E1DCAAA03213B0C36CEF674FCF0517E313CF34E2A894FC063647408D239AF66C40ABF49E88268A1BBC2CC4CABF45C11A6379CEA9D326F685ADF021969942AD4A2A324B6C9033720A55F280A01648BA614C2DE6AB110AA71F56379910CEED337A64599222888E18598078E6A2B7D4497024A0CDD0B382C371592A65BA538DC025B82361C9AF7334A8D6C259A57F3EBC4B21F01F0BE6346B2CB25EB10A41D5AA2B28255DA750BAA89C2C0982A5C2B5C2292B98E91C922848DC2036D96001BE1208C7A1C8BCF29F4CA463DFB964D07443DBBA9CF4DC7A90342DAE70C2703A9FBE324C08AA2B7C86A94E6A3D7760BD40366E1336A56D57C29225363F6A5F8E91216ED2574037B028B3A69E3C0511B08DD6E07D0BEC5A6DF0A2111AAD9AA58EFAAB4596E25099930A56809CE31A987D2A9C52702437A5A929A44478172470F75DABDAB1C6B66772B0720E2B53CF0459384A6DB94A18B621AD271881B7DBAD8AE4C330965D6CD367B96B830256233113C03191186A05B3F4D90A349CC8816642A0A55DD3D95DF8CBC999C05A48E79282BC434C7AA7A08C3BAA8C4FA686C78F8A2F99D102F75023C3FA87188429817170348A40C1350EAC679BE5158B9CF5950D4871F9EC687621196D3B420F0B6360A3A721B2B29E4988F315A52253369340BF54CB70A6F569192FA55361EE3482F39F344B5E5217A624461B9B6FA830E17A7A45D0E59D2EECE64AADF692860235678DCC2B9818B3DF268AB12B07C94C627D2C5FB929D435CA2260A6AAED");

    private static final Map<String, MLKEMParameters> parametersMap = new HashMap<String, MLKEMParameters>()
    {
        {
            put("ML-KEM-512", MLKEMParameters.ml_kem_512);
            put("ML-KEM-768", MLKEMParameters.ml_kem_768);
            put("ML-KEM-1024", MLKEMParameters.ml_kem_1024);
        }
    };

    private final SecureRandom RANDOM = new SecureRandom();

    public void testKeys()
    {
        new MLKEMPublicKeyParameters(MLKEMParameters.ml_kem_512, okayKey);

        try
        {
            new MLKEMPublicKeyParameters(MLKEMParameters.ml_kem_512, tooLargeKey);
            fail("no exception for invalid public key");
        }
        catch (IllegalArgumentException e)
        {
            assertEquals("'encoding' has invalid length", e.getMessage());
        }

        try
        {
            new MLKEMPrivateKeyParameters(MLKEMParameters.ml_kem_512, faultyPrivateKey);
            fail("no exception for invalid private key");
        }
        catch (IllegalArgumentException e)
        {
            assertEquals("'encoding' fails hash check", e.getMessage());
        }
    }

    public void testInputs()
        throws Exception
    {
        MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
        MLKEMKeyGenerationParameters genParam = new MLKEMKeyGenerationParameters(new SecureRandom(), MLKEMParameters.ml_kem_512);

        //
        // Generate keys and test.
        //
        kpGen.init(genParam);
        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
        MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

        MLKEMExtractor kemExt = new MLKEMExtractor(privateKey);

        try
        {
            kemExt.extractSecret(new byte[2000]);
            Assert.fail();
        }
        catch (IllegalArgumentException e)
        {
            assertEquals("encapsulation wrong length", e.getMessage());
        }

        try
        {
            kemExt.extractSecret(new byte[600]);
            Assert.fail();
        }
        catch (IllegalArgumentException e)
        {
            assertEquals("encapsulation wrong length", e.getMessage());
        }
    }

    public void testKeyGen() throws IOException
    {
        MLKEMParameters[] params = new MLKEMParameters[]{
            MLKEMParameters.ml_kem_512,
            MLKEMParameters.ml_kem_768,
            MLKEMParameters.ml_kem_1024,
        };

        String[] files = new String[]{
            "keyGen_ML-KEM-512.txt",
            "keyGen_ML-KEM-768.txt",
            "keyGen_ML-KEM-1024.txt",
        };

        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/kyber/acvp", name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

            String line = null;
            HashMap<String, String> buf = new HashMap<String, String>();
            while ((line = bin.readLine()) != null)
            {
                line = line.trim();

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        byte[] z = Hex.decode((String)buf.get("z"));
                        byte[] d = Hex.decode((String)buf.get("d"));
                        byte[] ek = Hex.decode((String)buf.get("ek"));
                        byte[] dk = Hex.decode((String)buf.get("dk"));

                        MLKEMParameters parameters = params[fileIndex];

                        MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
                        MLKEMKeyGenerationParameters genParam = new MLKEMKeyGenerationParameters(new FixedSecureRandom(Arrays.concatenate(d, z)), parameters);

                        //
                        // Generate keys and test.
                        //
                        kpGen.init(genParam);
                        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
                        MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters)kp.getPublic();
                        MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

                        MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(publicKey));
                        MLKEMPrivateKeyParameters privParams = (MLKEMPrivateKeyParameters)PrivateKeyFactory.createKey(
                            PrivateKeyInfoFactory.createPrivateKeyInfo(privateKey));

                        assertTrue(name + ": public key", Arrays.areEqual(ek, pubParams.getEncoded()));
                        assertTrue(name + ": secret key", Arrays.areEqual(dk, privParams.getEncoded()));

                    }
                    buf.clear();

                    continue;
                }

                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
        }
    }

    public void testKeyGenCombinedVectorSet() throws IOException
    {
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/mlkem", "ML-KEM-keyGen.txt");
        BufferedReader bin = new BufferedReader(new InputStreamReader(src));

        String line = null;
        HashMap<String, String> buf = new HashMap<String, String>();
        while ((line = bin.readLine()) != null)
        {
            line = line.trim();

            if (line.startsWith("#"))
            {
                continue;
            }
            if (line.length() == 0)
            {
                if (buf.size() > 0)
                {
                    byte[] z = Hex.decode((String)buf.get("z"));
                    byte[] d = Hex.decode((String)buf.get("d"));
                    byte[] ek = Hex.decode((String)buf.get("ek"));
                    byte[] dk = Hex.decode((String)buf.get("dk"));

                    MLKEMParameters parameters = (MLKEMParameters)parametersMap.get((String)buf.get("parameterSet"));

                    MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
                    MLKEMKeyGenerationParameters genParam = new MLKEMKeyGenerationParameters(new FixedSecureRandom(Arrays.concatenate(d, z)), parameters);

                    //
                    // Generate keys and test.
                    //
                    kpGen.init(genParam);
                    AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
                    MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters)kp.getPublic();
                    MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

                    MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
                        SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(publicKey));
                    MLKEMPrivateKeyParameters privParams = (MLKEMPrivateKeyParameters)PrivateKeyFactory.createKey(
                        PrivateKeyInfoFactory.createPrivateKeyInfo(privateKey));

                    assertTrue("public key", Arrays.areEqual(ek, pubParams.getEncoded()));
                    assertTrue("secret key", Arrays.areEqual(dk, privParams.getEncoded()));
                }
                buf.clear();

                continue;
            }

            int a = line.indexOf("=");
            if (a > -1)
            {
                buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
            }
        }
    }

    public void testEncapDecap_encapsulation() throws IOException
    {
        MLKEMParameters[] params = new MLKEMParameters[]{
            MLKEMParameters.ml_kem_512,
            MLKEMParameters.ml_kem_768,
            MLKEMParameters.ml_kem_1024,
        };

        String[] files = new String[]{
            "encapDecap_encapsulation_ML-KEM-512.txt",
            "encapDecap_encapsulation_ML-KEM-768.txt",
            "encapDecap_encapsulation_ML-KEM-1024.txt",
        };

        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/kyber/acvp", name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

            String line = null;
            HashMap<String, String> buf = new HashMap<String, String>();
            while ((line = bin.readLine()) != null)
            {
                line = line.trim();

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        byte[] m = Hex.decode((String)buf.get("m"));
                        byte[] c = Hex.decode((String)buf.get("c"));
                        byte[] k = Hex.decode((String)buf.get("k"));
//                        String reason = (String)buf.get("reason");
                        byte[] ek = Hex.decode((String)buf.get("ek"));
//                        byte[] dk = Hex.decode((String)buf.get("dk"));

                        MLKEMParameters parameters = params[fileIndex];

                        MLKEMPublicKeyParameters pubKey = new MLKEMPublicKeyParameters(parameters, ek);
//                        MLKEMPrivateKeyParameters privKey = new MLKEMPrivateKeyParameters(parameters,  dk);

                        MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubKey));
//                        MLKEMPrivateKeyParameters privParams = (MLKEMPrivateKeyParameters)PrivateKeyFactory.createKey(
//                            PrivateKeyInfoFactory.createPrivateKeyInfo(privKey));

                        // KEM Enc
                        SecretWithEncapsulation secWenc = MLKEMGenerator.internalGenerateEncapsulated(pubParams, m);
                        byte[] generated_cipher_text = secWenc.getEncapsulation();

                        byte[] secret = secWenc.getSecret();
                        assertTrue(name  + ": c", Arrays.areEqual(c,  generated_cipher_text));
                        assertTrue(name  + ": k", Arrays.areEqual(k, 0, secret.length, secret, 0, secret.length));

                    }
                    buf.clear();

                    continue;
                }

                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
        }
    }

    public void testEncapDecap_decapsulation() throws IOException
    {
        MLKEMParameters[] params = new MLKEMParameters[]{
            MLKEMParameters.ml_kem_512,
            MLKEMParameters.ml_kem_768,
            MLKEMParameters.ml_kem_1024,
        };

        String[] files = new String[]{
            "encapDecap_decapsulation_ML-KEM-512.txt",
            "encapDecap_decapsulation_ML-KEM-768.txt",
            "encapDecap_decapsulation_ML-KEM-1024.txt",
        };

        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/kyber/acvp", name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

            String line = null;
            HashMap<String, String> buf = new HashMap<String, String>();
            while ((line = bin.readLine()) != null)
            {
                line = line.trim();

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        byte[] c = Hex.decode((String)buf.get("c"));
                        byte[] k = Hex.decode((String)buf.get("k"));
//                        String reason = (String)buf.get("reason");
//                        byte[] ek = Hex.decode((String)buf.get("ek"));
                        byte[] dk = Hex.decode((String)buf.get("dk"));

                        MLKEMParameters parameters = params[fileIndex];

//                        MLKEMPublicKeyParameters pubKey = new MLKEMPublicKeyParameters(parameters, ek);
                        MLKEMPrivateKeyParameters privKey = new MLKEMPrivateKeyParameters(parameters,  dk);

//                        MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
//                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubKey));
                        MLKEMPrivateKeyParameters privParams = (MLKEMPrivateKeyParameters)PrivateKeyFactory.createKey(
                            PrivateKeyInfoFactory.createPrivateKeyInfo(privKey));

                        MLKEMExtractor extractor = new MLKEMExtractor(privParams);

                        byte[] dec_key = extractor.extractSecret(c);

                        assertTrue(name  + ": dk", Arrays.areEqual(dec_key,  k));
                    }
                    buf.clear();

                    continue;
                }

                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
        }
    }

    public void testEncapDecapCombinedVectorSet() throws IOException
    {
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/mlkem", "ML-KEM-encapDecap.txt");
        BufferedReader bin = new BufferedReader(new InputStreamReader(src));

        String line = null;
        HashMap<String, String> buf = new HashMap<String, String>();
        while ((line = bin.readLine()) != null)
        {
            line = line.trim();

            if (line.startsWith("#"))
            {
                continue;
            }
            if (line.length() == 0)
            {
                if (buf.size() > 0)
                {
                    byte[] c = Hex.decode((String)buf.get("c"));
                    byte[] k = Hex.decode((String)buf.get("k"));

                    MLKEMParameters parameters = (MLKEMParameters)parametersMap.get((String)buf.get("parameterSet"));

                    String function = (String)buf.get("function");
                    if ("encapsulation".equals(function))
                    {
                        byte[] m = Hex.decode((String)buf.get("m"));
                        byte[] ek = Hex.decode((String)buf.get("ek"));

                        MLKEMPublicKeyParameters pubKey = new MLKEMPublicKeyParameters(parameters, ek);

                        MLKEMPublicKeyParameters pubParams = (MLKEMPublicKeyParameters)PublicKeyFactory.createKey(
                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubKey));

                        SecretWithEncapsulation secWenc = MLKEMGenerator.internalGenerateEncapsulated(pubParams, m);
                        byte[] generated_cipher_text = secWenc.getEncapsulation();
                        byte[] secret = secWenc.getSecret();

                        assertTrue("encap: c", Arrays.areEqual(c,  generated_cipher_text));
                        assertTrue("encap: k", Arrays.areEqual(k, 0, secret.length, secret, 0, secret.length));
                    }
                    else
                    {
                        byte[] dk = Hex.decode((String)buf.get("dk"));
                        
                        MLKEMPrivateKeyParameters privKey = new MLKEMPrivateKeyParameters(parameters,  dk);

                        MLKEMPrivateKeyParameters privParams = (MLKEMPrivateKeyParameters)PrivateKeyFactory.createKey(
                            PrivateKeyInfoFactory.createPrivateKeyInfo(privKey));

                        MLKEMExtractor extractor = new MLKEMExtractor(privParams);
                        byte[] dec_key = extractor.extractSecret(c);

                        assertTrue("decap: dk", Arrays.areEqual(dec_key,  k));
                    }
                }
                buf.clear();

                continue;
            }

            int a = line.indexOf("=");
            if (a > -1)
            {
                buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
            }
        }
    }

    public void testModulus() throws IOException
    {
        MLKEMParameters[] params = new MLKEMParameters[]{
            MLKEMParameters.ml_kem_512,
            MLKEMParameters.ml_kem_768,
            MLKEMParameters.ml_kem_1024,
        };

        String[] files = new String[]{
            "ML-KEM-512.txt",
            "ML-KEM-768.txt",
            "ML-KEM-1024.txt",
        };

        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/kyber/modulus", name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

            String line = null;
            while ((line = bin.readLine()) != null)
            {
                line = line.trim();
                byte[] key = Hex.decode(line);
                MLKEMParameters parameters = params[fileIndex];

                try
                {
                    new MLKEMPublicKeyParameters(parameters, key);
                    fail();
                }
                catch (IllegalArgumentException e)
                {
                    assertEquals("Modulus check failed for ML-KEM public key", e.getMessage());
                }
            }
        }
    }

    public void testPrivInfoGeneration()
        throws IOException
    {
        OtherInfoGenerator.PartyU partyU = new OtherInfoGenerator.PartyU(MLKEMParameters.ml_kem_512,
            new AlgorithmIdentifier(OIWObjectIdentifiers.idSHA1), Hex.decode("beef"), Hex.decode("cafe"), RANDOM);

        byte[] partA = partyU.getSuppPrivInfoPartA();

        OtherInfoGenerator.PartyV partyV = new OtherInfoGenerator.PartyV(MLKEMParameters.ml_kem_512,
            new AlgorithmIdentifier(OIWObjectIdentifiers.idSHA1), Hex.decode("beef"), Hex.decode("cafe"), RANDOM);

        byte[] partB = partyV.getSuppPrivInfoPartB(partA);

        DEROtherInfo otherInfoU = partyU.generate(partB);

        DEROtherInfo otherInfoV = partyV.generate();

        assertTrue(Arrays.areEqual(otherInfoU.getEncoded(), otherInfoV.getEncoded()));
    }
    
    public void testMLKEM()
    {
        byte[] z = Hex.decode("99E3246884181F8E1DD44E0C7629093330221FD67D9B7D6E1510B2DBAD8762F7");
        byte[] d = Hex.decode("49AC8B99BB1E6A8EA818261F8BE68BDEAA52897E7EC6C40B530BC760AB77DCE3");
        String expectedPubKey = "A04184D4BC7B532A0F70A54D7757CDE6175A6843B861CB2BC4830C0012554CFC5D2C8A2027AA3CD967130E9B96241B11C4320C7649CC23A71BAFE691AFC08E680BCEF42907000718E4EACE8DA28214197BE1C269DA9CB541E1A3CE97CFADF9C6058780FE6793DBFA8218A2760B802B8DA2AA271A38772523A76736A7A31B9D3037AD21CEBB11A472B8792EB17558B940E70883F264592C689B240BB43D5408BF446432F412F4B9A5F6865CC252A43CF40A320391555591D67561FDD05353AB6B019B3A08A73353D51B6113AB2FA51D975648EE254AF89A230504A236A4658257740BDCBBE1708AB022C3C588A410DB3B9C308A06275BDF5B4859D3A2617A295E1A22F90198BAD0166F4A943417C5B831736CB2C8580ABFDE5714B586ABEEC0A175A08BC710C7A2895DE93AC438061BF7765D0D21CD418167CAF89D1EFC3448BCBB96D69B3E010C82D15CAB6CACC6799D3639669A5B21A633C865F8593B5B7BC800262BB837A924A6C5440E4FC73B41B23092C3912F4C6BEBB4C7B4C62908B03775666C22220DF9C88823E344C7308332345C8B795D34E8C051F21F5A21C214B69841358709B1C305B32CC2C3806AE9CCD3819FFF4507FE520FBFC27199BC23BE6B9B2D2AC1717579AC769279E2A7AAC68A371A47BA3A7DBE016F14E1A727333663C4A5CD1A0F8836CF7B5C49AC51485CA60345C990E06888720003731322C5B8CD5E6907FDA1157F468FD3FC20FA8175EEC95C291A262BA8C5BE990872418930852339D88A19B37FEFA3CFE82175C224407CA414BAEB37923B4D2D83134AE154E490A9B45A0563B06C953C3301450A2176A07C614A74E3478E48509F9A60AE945A8EBC7815121D90A3B0E07091A096CF02C57B25BCA58126AD0C629CE166A7EDB4B33221A0D3F72B85D562EC698B7D0A913D73806F1C5C87B38EC003CB303A3DC51B4B35356A67826D6EDAA8FEB93B98493B2D1C11B676A6AD9506A1AAAE13A824C7C08D1C6C2C4DBA9642C76EA7F6C8264B64A23CCCA9A74635FCBF03E00F1B5722B214376790793B2C4F0A13B5C40760B4218E1D2594DCB30A70D9C1782A5DD30576FA4144BFC8416EDA8118FC6472F56A979586F33BB070FB0F1B0B10BC4897EBE01BCA3893D4E16ADB25093A7417D0708C83A26322E22E6330091E30152BF823597C04CCF4CFC7331578F43A2726CCB428289A90C863259DD180C5FF142BEF41C7717094BE07856DA2B140FA67710967356AA47DFBC8D255B4722AB86D439B7E0A6090251D2D4C1ED5F20BBE6807BF65A90B7CB2EC0102AF02809DC9AC7D0A3ABC69C18365BCFF59185F33996887746185906C0191AED4407E139446459BE29C6822717644353D24AB6339156A9C424909F0A9025BB74720779BE43F16D81C8CC666E99710D8C68BB5CC4E12F314E925A551F09CC59003A1F88103C254BB978D75F394D3540E31E771CDA36E39EC54A62B5832664D821A72F1E6AFBBA27F84295B2694C498498E812BC8E9378FE541CEC5891B25062901CB7212E3CDC46179EC5BCEC10BC0B9311DE05074290687FD6A5392671654284CD9C8CC3EBA80EB3B662EB53EB75116704A1FEB5C2D056338532868DDF24EB8992AB8565D9E490CADF14804360DAA90718EAB616BAB0765D33987B47EFB6599C5563235E61E4BE670E97955AB292D9732CB8930948AC82DF230AC72297A23679D6B94C17F1359483254FEDC2F05819F0D069A443B78E3FC6C3EF4714B05A3FCA81CBBA60242A7060CD885D8F39981BB18092B23DAA59FD9578388688A09BBA079BC809A54843A60385E2310BBCBCC0213CE3DFAAB33B47F9D6305BC95C6107813C585C4B657BF30542833B14949F573C0612AD524BAAE69590C1277B86C286571BF66B3CFF46A3858C09906A794DF4A06E9D4B0A2E43F10F72A6C6C47E5646E2C799B71C33ED2F01EEB45938EB7A4E2E2908C53558A540D350369FA189C616943F7981D7618CF02A5B0A2BCC422E857D1A47871253D08293C1C179BCDC0437069107418205FDB9856623B8CA6B694C96C084B17F13BB6DF12B2CFBBC2B0E0C34B00D0FCD0AECFB27924F6984E747BE2A09D83A8664590A8077331491A4F7D720843F23E652C6FA840308DB4020337AAD37967034A9FB523B67CA70330F02D9EA20C1E84CB8E5757C9E1896B60581441ED618AA5B26DA56C0A5A73C4DCFD755E610B4FC81FF84E21";
        String expectedPrivKey = "[encoded data, 6336 characters]";

        MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
        kpGen.init(new MLKEMKeyGenerationParameters(new FixedSecureRandom(Arrays.concatenate(d, z)), MLKEMParameters.ml_kem_1024));

        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
        MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters)kp.getPublic();
        MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

        assertTrue(Arrays.areEqual(Hex.decode(expectedPubKey), publicKey.getEncoded()));
        assertTrue(Arrays.areEqual(Hex.decode(expectedPrivKey), privateKey.getEncoded()));

        byte[] message = Hex.decode("59C5154C04AE43AAFF32700F081700389D54BEC4C37C088B1C53F66212B12C72");

        SecretWithEncapsulation secretEncap = MLKEMGenerator.internalGenerateEncapsulated(publicKey, message);

        String expectedSharedSecret = "5CF38F578AC4AE95FBFED574B3D8EBF7CB1DC9074F22277360E36D775347C058";

        assertTrue(Arrays.areEqual(Hex.decode(expectedSharedSecret), secretEncap.getSecret()));

        String expectedCipherText = "8B9FE419250C5FB0463C8181FCF7CEC777136B738E015EBA31067AA4A8C378BBAC0121B88214F1AEB866E4F33C277099E09B4BF7E21CDDA30B5B32C18B0E9660C30601D85DAEC07AAF4B343EC5516FA501DD63088B999FB9A414C6CA593806C08CD4C775139BF0F0BF3676D773EDD56E616A13830D5F5FE35E515DBC84E43AAD0167D57E60A9DE30886ACD3F7F2006CAC26A7A07B4DADBEDFBED7F305764386AAD726D5B2BF14A376BAD8B4896688491733FB34E6EDEA10BFD5E448541CB6E69E3D87DF190AFA7FF62577775BAACEA444A6128A20200251D8FA759DC60FDA6A9730CFFE4997FE7EBCDD1644AE2D55290A4074CDD2CE53C18D22BC33671E68727A9B5A2FEAFB114A8045D96A56981E200A09661375987625ACC233EDE817AF1DEEAA21C7C4377423E73C5AF9BFF58A49DE6DAFD07A3E3BABD891F62BBA41D1856B8BC502CC86EE115A3598431E2B54AB0C5EACC3CE6A03090925C1FD5A251B00576763A963994A7A23EE12EBFC1B994F93C6144178F0BEF88245CE77CD32EF651826A6090AF561A5864DEC2A51D846F1F48F88B4B55F58C2373E0F67BDC95DC23A43E8546232A7B234E49F5226A3A63BDBCED7240FC81C2DB68AAEB2671A2FD231997BF8839C63A7F41F15E7242821D42E80BBC0F43FA9E353DE8B25ED8FFC242EB512C6A5260919AAE89A11176532BCCC762A520A37AEC4E7209AA81CEE0DD4ADD932C47EB8100BE98AA1DEEA9EA698115ADCED950A6C536D19AEB325CEA8C5245C0A2281533FB90809DC2BE90567EBE6AE229FE09B44DA2182585EA694D8A9AB33EBC24B44E09BD510F34B4140E1FB41162F9415F2D9106A0CEA00A26ED0920021F4E5BCFB3DABF5850DAB22B2E889D9611FBE06D0C899708EB5E5FAD2FBBE0D5C0BDE080F8E760EDFA037D55DA77F0F39591BF5B050C905FA538B7228E238A290DF340778DCBD6BE40A3B1DD455FB27ADBE176AEF6CC295BEA570BDC221BA14002E3B113B0EF237452FBC9F1AEC42E0D2B33F19832DB0A6171CAEB0B30EEAD3A54B704B761C7D4AFEA8F6AFC15156666A081C43AEB2E04FEECEF8AABA4049BD78B120B9ABA86A60342A0CF806411C473C26C4BE1540E3312388BCBC8523BA73F40EA28D5564274F3661D7ACAA0F1E8D0F28DCF6B501329963E6857FDB2AAE873A7D9D6C14821F6C0B6AA50AC449075CD6F2A256C5A05959DAB5A5912CC8E8F8B9F59941BFCCE6A28CBA74A20382B1FD3382D056547D5BC5EF4AAE62F96F038C595A4F901D6AE790F8978292AD1CC3A1E800B71A5BBE84533646655E3752FBD6B02B97B204E75D28A34C2F990FB8E8CD31CE6E683FA7E67DA03367E8D47DC626F060FBA2D0425004CAC2A61D982D2E3D85008624B45DB022CF51BA265B5E974712A9372EECAC0EA272B2FC56EBED0D32105521BA2C4A8FE0C678CE4E45902C7BA9D510BD47B2B5F931DD732F27DE9B42FD4AA39EAC765283A9965EE97C0D88E23EFA6F718242C67770B87BF8832858C1D13FC520870BD34F2B9C6FBFD1A528B744F814C93F4F4E87108316FE2AB06E02292DEA7FCF6FEFB17BF5AA7376A4A9BDB7C49BF709EB1E05D60EF14CD85A75239B97BCA9A6A3CC1B28F28979D612431BAAC1ACEE5EF62776B4D51B7EB0F63DF507760097223CA903E16E02DEB7FCABFBEC26DAEDC0ED4CC55726BDC31D1775112EF3C35D1DF928C6EB7830D8CA6570CB5CE348E3F26DDE864F20E5BE7B99E264EBC0E9D8DE9C6E4B7FE3CFBE673833CF7E8B3081529062CB6815C7C0766822B3B31E56BA1FC73FE3DED4B5D435BFCE2F2997C1D4B9CE293220DD461103BE084BF12076372668A69836769C1F6D8C32E2C7BC2E7D66714C814793A2970C90DD94DF14C89C60DD35B52A14778E137E750CE83AC3AAB667FCBDCBA38B7FA6D1C6BF7B99D957078176D9779A09F84B75FBC2A11769EF65532B09ACA4C9A3766B4A1FC717F94648FB8B8D9363E54F1C4201C075C18B1EAE098B83598089585ED9DC06B96E2D1C96DC738086EBBC26C3193B64139E1FC1DFB22A17893506EF7B35792B4EB00196693686EB5DEB3CEB436DD16D2D92A0FD31F468AF8662040F5257BFA0F14991C0D560999EEF775178D14955ADF091DD797AC1FDCEC7776055271C0F130562D0B0A6749B159DD0DB9AC69271AC719B83B683CE8B32342AC4AB257B0F8083C8CC86338AFA4D386C9848F413ED0";

        assertTrue(Arrays.areEqual(Hex.decode(expectedCipherText), secretEncap.getEncapsulation()));

        MLKEMExtractor kemExtract = new MLKEMExtractor(privateKey);

        byte[] decryptedSharedSecret = kemExtract.extractSecret(secretEncap.getEncapsulation());

        assertTrue(Arrays.areEqual(Hex.decode(expectedSharedSecret), decryptedSharedSecret));
    }

    public void testRNG()
    {
        String temp = "061550234D158C5EC95595FE04EF7A25767F2E24CC2BC479D09D86DC9ABCFDE7056A8C266F9EF97ED08541DBD2E1FFA1";
        byte[] seed = Hex.decode(temp);
        NISTSecureRandom r = new NISTSecureRandom(seed, null);
        byte[] testBytes = new byte[48];
        r.nextBytes(testBytes);

        String randBytesString = "7C9935A0B07694AA0C6D10E4DB6B1ADD2FD81A25CCB148032DCD739936737F2DB505D7CFAD1B497499323C8686325E47";
        byte[] randBytes = Hex.decode(randBytesString);

        assertTrue(Arrays.areEqual(randBytes, testBytes));
    }

    public void testParameters()
        throws Exception
    {
        assertEquals(256, MLKEMParameters.ml_kem_512.getSessionKeySize());
        assertEquals(256, MLKEMParameters.ml_kem_768.getSessionKeySize());
        assertEquals(256, MLKEMParameters.ml_kem_1024.getSessionKeySize());
    }

    public void testMLKEMRandom()
    {
        MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
        kpGen.init(new MLKEMKeyGenerationParameters(RANDOM, MLKEMParameters.ml_kem_1024));

        MLKEMGenerator kemGen = new MLKEMGenerator(RANDOM);

        for (int i = 0; i != 1000; i++)
        {
            AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
            MLKEMPublicKeyParameters publicKey = (MLKEMPublicKeyParameters)kp.getPublic();
            MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

            SecretWithEncapsulation secretEncap = kemGen.generateEncapsulated(publicKey);

            MLKEMExtractor kemExtract = new MLKEMExtractor(privateKey);

            byte[] decryptedSharedSecret = kemExtract.extractSecret(secretEncap.getEncapsulation());

            assertTrue(Arrays.areEqual(secretEncap.getSecret(), decryptedSharedSecret));
        }
    }

    public void testWithPreferredFormat()
    {
        MLKEMParameters[] parameters = new MLKEMParameters[]{
            MLKEMParameters.ml_kem_512,
            MLKEMParameters.ml_kem_768,
            MLKEMParameters.ml_kem_1024,
        };

        for (int i = 0; i < parameters.length; ++i)
        {
            MLKEMKeyPairGenerator kpGen = new MLKEMKeyPairGenerator();
            kpGen.init(new MLKEMKeyGenerationParameters(RANDOM, parameters[i]));

            AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
            MLKEMPrivateKeyParameters privateKey = (MLKEMPrivateKeyParameters)kp.getPrivate();

            implWithPreferredFormat(privateKey, MLKEMPrivateKeyParameters.SEED_ONLY);
            implWithPreferredFormat(privateKey, MLKEMPrivateKeyParameters.EXPANDED_KEY);
            implWithPreferredFormat(privateKey, MLKEMPrivateKeyParameters.BOTH);
        }
    }

    private static void implWithPreferredFormat(MLKEMPrivateKeyParameters privateKey, int format)
    {
        int originalFormat = privateKey.getPreferredFormat();
        byte[] originalSeed = privateKey.getSeed();
        byte[] originalEncoding = privateKey.getEncoded();

        MLKEMPrivateKeyParameters updatedPrivateKey = privateKey.withPreferredFormat(format);

        int updatedFormat = updatedPrivateKey.getPreferredFormat();
        byte[] updatedSeed = updatedPrivateKey.getSeed();
        byte[] updatedEncoding = updatedPrivateKey.getEncoded();

        assertEquals(format, updatedFormat);
        assertTrue(Arrays.areEqual(originalSeed, updatedSeed));
        assertTrue(Arrays.areEqual(originalEncoding, updatedEncoding));

        if (format == originalFormat)
        {
            assertSame(privateKey, updatedPrivateKey);
        }
    }
}
