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.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import junit.framework.TestCase;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.generators.SLHDSAKeyPairGenerator;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.crypto.params.SLHDSAKeyGenerationParameters;
import org.bouncycastle.crypto.params.SLHDSAParameters;
import org.bouncycastle.crypto.params.SLHDSAPrivateKeyParameters;
import org.bouncycastle.crypto.params.SLHDSAPublicKeyParameters;
import org.bouncycastle.crypto.signers.SLHDSASigner;
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.test.TestResourceFinder;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.FixedSecureRandom;

public class SLHDSATest
    extends TestCase
{
    private static final Map<String, SLHDSAParameters> parametersMap = new HashMap<String, SLHDSAParameters>()
    {
        {
            put("SLH-DSA-SHA2-128s", SLHDSAParameters.sha2_128s);
            put("SLH-DSA-SHA2-128f", SLHDSAParameters.sha2_128f);
            put("SLH-DSA-SHA2-192s", SLHDSAParameters.sha2_192s);
            put("SLH-DSA-SHA2-192f", SLHDSAParameters.sha2_192f);
            put("SLH-DSA-SHA2-256s", SLHDSAParameters.sha2_256s);
            put("SLH-DSA-SHA2-256f", SLHDSAParameters.sha2_256f);

            put("SLH-DSA-SHAKE-128s", SLHDSAParameters.shake_128s);
            put("SLH-DSA-SHAKE-128f", SLHDSAParameters.shake_128f);
            put("SLH-DSA-SHAKE-192s", SLHDSAParameters.shake_192s);
            put("SLH-DSA-SHAKE-192f", SLHDSAParameters.shake_192f);
            put("SLH-DSA-SHAKE-256s", SLHDSAParameters.shake_256s);
            put("SLH-DSA-SHAKE-256f", SLHDSAParameters.shake_256f);
        }
    };

    SLHDSAParameters[] PARAMETER_SETS = new SLHDSAParameters[]
        {
            SLHDSAParameters.sha2_128f,
            SLHDSAParameters.sha2_128s,
            SLHDSAParameters.sha2_192f,
            SLHDSAParameters.sha2_192s,
            SLHDSAParameters.sha2_256f,
            SLHDSAParameters.sha2_256s,
            SLHDSAParameters.shake_128f,
            SLHDSAParameters.shake_128s,
            SLHDSAParameters.shake_192f,
            SLHDSAParameters.shake_192s,
            SLHDSAParameters.shake_256f,
            SLHDSAParameters.shake_256s,
        };

    public void testConsistency()
    {
        SecureRandom random = new SecureRandom();

        SLHDSAKeyPairGenerator kpg = new SLHDSAKeyPairGenerator();

        for (int idx = 0; idx != PARAMETER_SETS.length; idx++)
        {
            SLHDSAParameters parameters = PARAMETER_SETS[idx];
            kpg.init(new SLHDSAKeyGenerationParameters(random, parameters));

            {
                AsymmetricCipherKeyPair kp = kpg.generateKeyPair();

                SLHDSASigner signer = new SLHDSASigner();

                {
                    int msgLen = random.nextInt(257);
                    byte[] msg = new byte[msgLen];
                    random.nextBytes(msg);

                    // sign
                    signer.init(true, new ParametersWithRandom(kp.getPrivate(), random));
                    byte[] signature = signer.generateSignature(msg);

                    // verify
                    signer.init(false, kp.getPublic());
                    boolean shouldVerify = signer.verifySignature(msg, signature);

                    assertTrue(shouldVerify);
                }
            }
        }
    }

    public void testKeyGenSingleFile()
        throws IOException
    {
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/slhdsa/", "SLH-DSA-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[] skSeed = Hex.decode((String)buf.get("skSeed"));
                    byte[] skPrf = Hex.decode((String)buf.get("skPrf"));
                    byte[] pkSeed = Hex.decode((String)buf.get("pkSeed"));
                    byte[] pk = Hex.decode((String)buf.get("pk"));
                    byte[] sk = Hex.decode((String)buf.get("sk"));

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

                    SLHDSAKeyPairGenerator kpGen = new SLHDSAKeyPairGenerator();
                    SLHDSAKeyGenerationParameters genParam = new SLHDSAKeyGenerationParameters(new SecureRandom(), parameters);

                    //
                    // Generate keys and test.
                    //
                    kpGen.init(genParam);
                    AsymmetricCipherKeyPair kp = kpGen.internalGenerateKeyPair(skSeed, skPrf, pkSeed);

                    SLHDSAPublicKeyParameters pubParams = (SLHDSAPublicKeyParameters)PublicKeyFactory.createKey(
                        SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo((SLHDSAPublicKeyParameters)kp.getPublic()));
                    SLHDSAPrivateKeyParameters privParams = (SLHDSAPrivateKeyParameters)PrivateKeyFactory.createKey(
                        PrivateKeyInfoFactory.createPrivateKeyInfo((SLHDSAPrivateKeyParameters)kp.getPrivate()));

                    assertTrue("public key", Arrays.areEqual(pk, pubParams.getEncoded()));
                    assertTrue("secret key", Arrays.areEqual(sk, 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 testSigGenSingleFile()
        throws IOException
    {
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/slhdsa", "SLH-DSA-sigGen.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)
                {
                    boolean deterministic = !buf.containsKey("additionalRandomness");
                    byte[] sk = Hex.decode((String)buf.get("sk"));
//                    int messageLength = Integer.parseInt((String)buf.get("messageLength"));
                    byte[] message = Hex.decode((String)buf.get("message"));
                    byte[] signature = Hex.decode((String)buf.get("signature"));
                    byte[] rnd = null;

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

                    SLHDSAPrivateKeyParameters privParams = new SLHDSAPrivateKeyParameters(parameters, sk);

                    if (!deterministic)
                    {
                        rnd = Hex.decode((String)buf.get("additionalRandomness"));
                    }
                    else
                    {
                        rnd = privParams.getPublicSeed();
                    }

                    // sign
                    InternalSLHDSASigner signer = new InternalSLHDSASigner();

                    signer.init(true, privParams);
                    byte[] sigGenerated = signer.internalGenerateSignature(message, rnd);
                    assertTrue(Arrays.areEqual(sigGenerated, signature));
                }
                buf.clear();

                continue;
            }

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

    public void testSigVerSingleFile()
        throws IOException
    {
        String name = "SLH-DSA-sigVer.txt";
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/slhdsa", 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)
                {
                    boolean testPassed = TestUtils.parseBoolean((String)buf.get("testPassed"));
//                    boolean deterministic = !buf.containsKey("additionalRandomness");
                    String reason = (String)buf.get("reason");

                    byte[] pk = Hex.decode((String)buf.get("pk"));
                    byte[] message = Hex.decode((String)buf.get("message"));
                    byte[] signature = Hex.decode((String)buf.get("signature"));

//                    byte[] rnd = null;
//                    if (!deterministic)
//                    {
//                        rnd = Hex.decode((String)buf.get("additionalRandomness"));
//                    }

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

                    SLHDSAPublicKeyParameters pubParams = new SLHDSAPublicKeyParameters(parameters, pk);

                    InternalSLHDSASigner verifier = new InternalSLHDSASigner();
                    verifier.init(false, pubParams);
                    boolean ver = verifier.internalVerifySignature(message, signature);
                    assertEquals("expected " + testPassed + " " + reason, ver, testPassed);
                }
                buf.clear();

                continue;
            }

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

    public void testKeyGen()
        throws IOException
    {
        String[] files = new String[]{
            "keyGen_SLH-DSA-SHA2-128s.txt",
            "keyGen_SLH-DSA-SHA2-192f.txt",
            "keyGen_SLH-DSA-SHAKE-192s.txt",
            "keyGen_SLH-DSA-SHAKE-256f.txt",
        };

        SLHDSAParameters[] params = new SLHDSAParameters[]{
            SLHDSAParameters.sha2_128s,
            SLHDSAParameters.sha2_192f,
            SLHDSAParameters.shake_192s,
            SLHDSAParameters.shake_256f,
        };

        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource("pqc/crypto/slhdsa/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[] skSeed = Hex.decode((String)buf.get("skSeed"));
                        byte[] skPrf = Hex.decode((String)buf.get("skPrf"));
                        byte[] pkSeed = Hex.decode((String)buf.get("pkSeed"));
                        byte[] pk = Hex.decode((String)buf.get("pk"));
                        byte[] sk = Hex.decode((String)buf.get("sk"));

                        SLHDSAParameters parameters = params[fileIndex];

                        SLHDSAKeyPairGenerator kpGen = new SLHDSAKeyPairGenerator();
                        SLHDSAKeyGenerationParameters genParam = new SLHDSAKeyGenerationParameters(new SecureRandom(), parameters);
                        //
                        // Generate keys and test.
                        //
                        kpGen.init(genParam);
                        AsymmetricCipherKeyPair kp = kpGen.internalGenerateKeyPair(skSeed, skPrf, pkSeed);

                        SLHDSAPublicKeyParameters pubParams = (SLHDSAPublicKeyParameters)PublicKeyFactory.createKey(
                            SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo((SLHDSAPublicKeyParameters)kp.getPublic()));
                        SLHDSAPrivateKeyParameters privParams = (SLHDSAPrivateKeyParameters)PrivateKeyFactory.createKey(
                            PrivateKeyInfoFactory.createPrivateKeyInfo((SLHDSAPrivateKeyParameters)kp.getPrivate()));

                        assertTrue(name + ": public key", Arrays.areEqual(pk, pubParams.getEncoded()));
                        assertTrue(name + ": secret key", Arrays.areEqual(sk, 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 testSigGen()
        throws IOException
    {
        String[] files = new String[]{
            "sigGen_SLH-DSA-SHA2-192s.txt",
            "sigGen_SLH-DSA-SHA2-256f.txt",
            "sigGen_SLH-DSA-SHAKE-128f.txt",
            "sigGen_SLH-DSA-SHAKE-192s.txt",
            "sigGen_SLH-DSA-SHAKE-256f.txt",
        };

        SLHDSAParameters[] params = new SLHDSAParameters[]{
            SLHDSAParameters.sha2_192s,
            SLHDSAParameters.sha2_256f,
            SLHDSAParameters.shake_128f,
            SLHDSAParameters.shake_192s,
            SLHDSAParameters.shake_256f,
        };

        TestSampler sampler = new TestSampler();

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

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

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        if (sampler.skipTest(count++))
                        {
                            continue;
                        }

                        boolean deterministic = !buf.containsKey("additionalRandomness");
                        byte[] sk = Hex.decode((String)buf.get("sk"));
//                        int messageLength = Integer.parseInt((String)buf.get("messageLength"));
                        byte[] message = Hex.decode((String)buf.get("message"));
                        byte[] signature = Hex.decode((String)buf.get("signature"));
                        byte[] rnd = null;

                        SLHDSAParameters parameters = params[fileIndex];

                        SLHDSAPrivateKeyParameters privParams = new SLHDSAPrivateKeyParameters(parameters, sk);

                        if (!deterministic)
                        {
                            rnd = Hex.decode((String)buf.get("additionalRandomness"));
                        }
                        else
                        {
                            rnd = privParams.getPublicSeed();
                        }

                        // sign
                        InternalSLHDSASigner signer = new InternalSLHDSASigner();

                        signer.init(true, privParams);
                        byte[] sigGenerated = signer.internalGenerateSignature(message, rnd);
                        assertTrue(Arrays.areEqual(sigGenerated, signature));
                    }
                    buf.clear();

                    continue;
                }

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

    public void testSigVer()
        throws IOException
    {
        String[] files = new String[]{
            "sigVer_SLH-DSA-SHA2-192s.txt",
            "sigVer_SLH-DSA-SHA2-256f.txt",
            "sigVer_SLH-DSA-SHAKE-128f.txt",
            "sigVer_SLH-DSA-SHAKE-192s.txt",
            "sigVer_SLH-DSA-SHAKE-256f.txt",
        };

        SLHDSAParameters[] params = new SLHDSAParameters[]{
            SLHDSAParameters.sha2_192s,
            SLHDSAParameters.sha2_256f,
            SLHDSAParameters.shake_128f,
            SLHDSAParameters.shake_192s,
            SLHDSAParameters.shake_256f,
        };

        TestSampler sampler = new TestSampler();

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

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

                if (line.startsWith("#"))
                {
                    continue;
                }

                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        if (sampler.skipTest(count++))
                        {
                            continue;
                        }

                        boolean testPassed = TestUtils.parseBoolean((String)buf.get("testPassed"));
//                        boolean deterministic = !buf.containsKey("additionalRandomness");
                        String reason = (String)buf.get("reason");

                        byte[] pk = Hex.decode((String)buf.get("pk"));
//                        byte[] sk = Hex.decode((String)buf.get("sk"));
                        byte[] message = Hex.decode((String)buf.get("message"));
                        byte[] signature = Hex.decode((String)buf.get("signature"));

//                        byte[] rnd = null;
//                        if (!deterministic)
//                        {
//                            rnd = Hex.decode((String)buf.get("additionalRandomness"));
//                        }

                        SLHDSAParameters parameters = params[fileIndex];

                        SLHDSAPublicKeyParameters pubParams = new SLHDSAPublicKeyParameters(parameters, pk);
//                        SLHDSAPrivateKeyParameters privParams = new SLHDSAPrivateKeyParameters(parameters, sk);

                        InternalSLHDSASigner verifier = new InternalSLHDSASigner();
                        verifier.init(false, pubParams);
                        boolean ver = verifier.internalVerifySignature(message, signature);
                        assertEquals("expected " + testPassed + " " + reason, ver, testPassed);
                    }
                    buf.clear();

                    continue;
                }

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

    public void testBasicKeyGenerationSha2128sSimple()
    {
        byte[] skSeed = Hex.decode("2F896D61D9CD9038CA303394FADAA22A");
        byte[] skPrf = Hex.decode("24AC5EC1D86A989CA2196C3C8632419C");
        byte[] pkSeed = Hex.decode("1A05A42FE300E87B16AEE116CB2E2363");
        byte[] pk = Hex.decode("1A05A42FE300E87B16AEE116CB2E236358E2C3E62632C9DE03D08A535A0EB7E7");
        byte[] sk = Hex.decode("2F896D61D9CD9038CA303394FADAA22A24AC5EC1D86A989CA2196C3C8632419C1A05A42FE300E87B16AEE116CB2E236358E2C3E62632C9DE03D08A535A0EB7E7");

        SLHDSAParameters parameters = SLHDSAParameters.sha2_128s;

        SLHDSAKeyPairGenerator kpGen = new SLHDSAKeyPairGenerator();
        SLHDSAKeyGenerationParameters genParam = new SLHDSAKeyGenerationParameters(new SecureRandom(), parameters);
        //
        // Generate keys and test.
        //
        kpGen.init(genParam);
        AsymmetricCipherKeyPair kp = kpGen.internalGenerateKeyPair(skSeed, skPrf, pkSeed);

        SLHDSAPublicKeyParameters pubParams = (SLHDSAPublicKeyParameters)kp.getPublic();
        SLHDSAPrivateKeyParameters privParams = (SLHDSAPrivateKeyParameters)kp.getPrivate();

        assertTrue("public key", Arrays.areEqual(pk, pubParams.getEncoded()));
        assertTrue("secret key", Arrays.areEqual(sk, privParams.getEncoded()));
    }

    public void testBasicKeyGenerationShake256128fSimpleSign()
    {
        byte[] sk = Hex.decode("DADB023900B157BAEDFF38B4BDE4B308C83A26A11170274E7E35CD3935AEAF07119231DA3849A12477373395D264043DA6CECC80D20A2E15A3622ABFFC221FC8");
        byte[] message = Hex.decode("3048BDE7F28C0414CC318C90048F23AFECF079866C34858521192E1684F37F0BC5D2C8585E9BF753626F6E853779D41C15BDA83DEF79DBF8A11B82EAE066833AB6C409D8AC386C942D69FF482D26A1A4030F7C082E36CFCEAA7491CB2F25BD61B79BACD91DD72C91C5D673BE48866D33E6B20F9DC83BD5639D27B0D8CA326AA1");
        byte[] signature = Hex.decode("[encoded data, 34176 characters]");
        byte[] rnd = Hex.decode("934CEDC78C6F657E3BF6120E38EBB228");

        SLHDSAParameters parameters = SLHDSAParameters.shake_128f;

        SLHDSAPrivateKeyParameters privParams = new SLHDSAPrivateKeyParameters(parameters, sk);

        // sign
        InternalSLHDSASigner signer = new InternalSLHDSASigner();

        signer.init(true, privParams);
        byte[] sigGenerated = signer.internalGenerateSignature(message, rnd);
        assertTrue(Arrays.areEqual(sigGenerated, signature));
    }

    private static class InternalSLHDSASigner
        extends SLHDSASigner
    {
        public byte[] internalGenerateSignature(byte[] message, byte[] optRand)
        {
            return super.internalGenerateSignature(message, optRand);
        }

        public boolean internalVerifySignature(byte[] message, byte[] signature)
        {
            return super.internalVerifySignature(message, signature);
        }
    }
}
