package org.bouncycastle.pqc.crypto.test;

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

import junit.framework.Assert;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.AsymmetricCipherKeyPairGenerator;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.EncapsulatedSecretExtractor;
import org.bouncycastle.crypto.EncapsulatedSecretGenerator;
import org.bouncycastle.crypto.SecretWithEncapsulation;
import org.bouncycastle.crypto.Signer;
import org.bouncycastle.crypto.params.AsymmetricKeyParameter;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.pqc.crypto.MessageSigner;
import org.bouncycastle.pqc.crypto.util.PrivateKeyFactory;
import org.bouncycastle.pqc.crypto.util.PrivateKeyInfoFactory;
import org.bouncycastle.pqc.crypto.util.PublicKeyFactory;
import org.bouncycastle.pqc.crypto.util.SubjectPublicKeyInfoFactory;
import org.bouncycastle.test.TestResourceFinder;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;

class TestUtils
{
    static boolean parseBoolean(String value)
    {
        return "true".equalsIgnoreCase(value);
    }

    public abstract static class SignerOperation
    {
        public abstract SecureRandom getSecureRandom(byte[] seed);

        public abstract AsymmetricCipherKeyPairGenerator getAsymmetricCipherKeyPairGenerator(int fileIndex, SecureRandom random);

        public abstract byte[] getPublicKeyEncoded(CipherParameters pubParams);

        public abstract byte[] getPrivateKeyEncoded(CipherParameters privParams);

        public abstract Signer getSigner();

        public abstract MessageSigner getMessageSigner();

        public CipherParameters setSignParameters(CipherParameters privParams, SecureRandom random)
        {
            return new ParametersWithRandom(privParams, random);
        }

        public CipherParameters setVerifyParameters(CipherParameters pubParams, CipherParameters privParams)
        {
            return pubParams;
        }
    }

    public interface KeyEncapsulationOperation
    {
        SecureRandom getSecureRandom(byte[] seed);

        AsymmetricCipherKeyPairGenerator getAsymmetricCipherKeyPairGenerator(int fileIndex, SecureRandom random);

        byte[] getPublicKeyEncoded(AsymmetricKeyParameter pubParams);

        byte[] getPrivateKeyEncoded(AsymmetricKeyParameter privParams);

        EncapsulatedSecretGenerator getKEMGenerator(SecureRandom random);

        EncapsulatedSecretExtractor getKEMExtractor(AsymmetricKeyParameter privParams);

        int getSessionKeySize();
    }

    /**
     * TEMPORARY (SQIsign port): keygen-only variant. Runs the same KAT-parsing
     * loop as the full {@link #testTestVector(boolean, boolean, boolean, String,
     * String[], SignerOperation)} but stops after the pk/sk equality check.
     * Used while the SQIsign sign / verify engine is being ported piece by
     * piece. Once the engine is complete, callers should switch to the full
     * variant and this method can be removed.
     */
    public static void testTestVectorKeygenOnly(boolean sampleOnly, boolean enableFactory, String homeDir, String[] files, SignerOperation operation)
        throws Exception
    {
        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];

            InputStream src = TestResourceFinder.findTestResource(homeDir, name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

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

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        String count = (String)buf.get("count");
                        if (sampler != null && sampler.skipTest(count))
                        {
                            continue;
                        }

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

                        SecureRandom random = operation.getSecureRandom(seed);

                        AsymmetricCipherKeyPairGenerator kpGen = operation.getAsymmetricCipherKeyPairGenerator(fileIndex, random);

                        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
                        AsymmetricKeyParameter pubParams;
                        CipherParameters privParams;
                        if (enableFactory)
                        {
                            pubParams = PublicKeyFactory.createKey(
                                SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(kp.getPublic()));
                            privParams = PrivateKeyFactory.createKey(
                                PrivateKeyInfoFactory.createPrivateKeyInfo(kp.getPrivate()));
                        }
                        else
                        {
                            pubParams = kp.getPublic();
                            privParams = kp.getPrivate();
                        }

                        Assert.assertTrue(name + ": public key", Arrays.areEqual(pk, operation.getPublicKeyEncoded(pubParams)));
                        Assert.assertTrue(name + ": secret key", Arrays.areEqual(sk, operation.getPrivateKeyEncoded(privParams)));
                        System.out.println("Count " + count + " keygen pass");
                    }
                    buf.clear();
                    continue;
                }

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

    public static void testTestVector(boolean sampleOnly, boolean enableFactory, boolean isSigner, String homeDir, String[] files, SignerOperation operation)
        throws Exception
    {
        for (int fileIndex = 0; fileIndex != files.length; fileIndex++)
        {
            String name = files[fileIndex];

            InputStream src = TestResourceFinder.findTestResource(homeDir, name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

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

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (buf.size() > 0)
                    {
                        String count = (String)buf.get("count");
                        if (sampler != null && sampler.skipTest(count))
                        {
                            continue;
                        }

                        byte[] seed = Hex.decode((String)buf.get("seed"));
                        byte[] pk = Hex.decode((String)buf.get("pk"));
                        byte[] sk = Hex.decode((String)buf.get("sk"));
                        byte[] message = Hex.decode((String)buf.get("msg"));
                        byte[] signature = Hex.decode((String)(buf.get("sm") == null ? buf.get("sig") : buf.get("sm")));

                        SecureRandom random = operation.getSecureRandom(seed);

                        AsymmetricCipherKeyPairGenerator kpGen = operation.getAsymmetricCipherKeyPairGenerator(fileIndex, random);

                        //
                        // Generate keys and test.
                        //
                        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
                        CipherParameters pubParams, privParams;
                        if (enableFactory)
                        {
                            pubParams = PublicKeyFactory.createKey(
                                SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(kp.getPublic()));
                            privParams = PrivateKeyFactory.createKey(
                                PrivateKeyInfoFactory.createPrivateKeyInfo(kp.getPrivate()));
                        }
                        else
                        {
                            pubParams = kp.getPublic();
                            privParams = kp.getPrivate();
                        }

                        Assert.assertTrue(name + ": public key", Arrays.areEqual(pk, operation.getPublicKeyEncoded(pubParams)));
                        Assert.assertTrue(name + ": secret key", Arrays.areEqual(sk, operation.getPrivateKeyEncoded(privParams)));

                        byte[] sigGenerated;
                        privParams = operation.setSignParameters(privParams, random);
                        pubParams = operation.setVerifyParameters(pubParams, privParams);
                        if (isSigner)
                        {
                            Signer signer = operation.getSigner();
                            signer.init(true, privParams);
                            signer.update(message, 0, message.length);
                            sigGenerated = signer.generateSignature();
                        }
                        else
                        {
                            MessageSigner signer = operation.getMessageSigner();
                            signer.init(true, privParams);
                            sigGenerated = signer.generateSignature(message);
                        }
//                        for (int i = 0; i < sigGenerated.length; i++)
//                        {
//                            if (sigGenerated[i] != signature[i])
//                            {
//                                System.out.println(i + " " + sigGenerated[i] + " " + signature[i]);
//                            }
//                        }
                        Assert.assertTrue(Arrays.areEqual(sigGenerated, signature));

                        if (isSigner)
                        {
                            Signer signer = operation.getSigner();
                            signer.init(false, pubParams);
                            signer.update(message, 0, message.length);
                            Assert.assertTrue(signer.verifySignature(sigGenerated));
                        }
                        else
                        {
                            MessageSigner signer = operation.getMessageSigner();
                            signer.init(false, pubParams);
                            Assert.assertTrue(signer.verifySignature(message, sigGenerated));
                        }
                        //System.out.println("Count " + count + " pass");
                    }
                    buf.clear();
                    continue;
                }

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

    public static void testTestVector(boolean alwaysFull, boolean enableFactory, String homeDir, String[] files, KeyEncapsulationOperation operation)
        throws Exception
    {
        for (int fileIndex = 0; fileIndex < files.length; fileIndex++)
        {
            String name = files[fileIndex];
            InputStream src = TestResourceFinder.findTestResource(homeDir, name);
            BufferedReader bin = new BufferedReader(new InputStreamReader(src));

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

                if (line.startsWith("#"))
                {
                    continue;
                }
                if (line.length() == 0)
                {
                    if (!buf.isEmpty())
                    {
                        String count = (String)buf.get("count");
                        if (sampler != null && sampler.skipTest(count))
                        {
                            continue;
                        }
                        System.out.println("test case: " + count);

                        byte[] seed = Hex.decode((String)buf.get("seed")); // seed for bike secure random
                        byte[] pk = Hex.decode((String)buf.get("pk"));     // public key
                        byte[] sk = Hex.decode((String)buf.get("sk"));     // private key
                        byte[] ct = Hex.decode((String)buf.get("ct"));     // ciphertext
                        byte[] ss = Hex.decode((String)buf.get("ss"));     // session key

                        SecureRandom random = operation.getSecureRandom(seed);

                        AsymmetricCipherKeyPairGenerator kpGen = operation.getAsymmetricCipherKeyPairGenerator(fileIndex, random);

                        //
                        // Generate keys and test.
                        //
                        AsymmetricCipherKeyPair kp = kpGen.generateKeyPair();
                        AsymmetricKeyParameter pubParams = kp.getPublic();
                        ;
                        AsymmetricKeyParameter privParams = kp.getPrivate();
                        if (enableFactory)
                        {
                            pubParams = PublicKeyFactory.createKey(
                                SubjectPublicKeyInfoFactory.createSubjectPublicKeyInfo(pubParams));
                            privParams = PrivateKeyFactory.createKey(
                                PrivateKeyInfoFactory.createPrivateKeyInfo(privParams));
                        }

                        Assert.assertTrue(name + " " + count + ": public key", Arrays.areEqual(pk, operation.getPublicKeyEncoded(pubParams)));
                        Assert.assertTrue(name + " " + count + ": secret key", Arrays.areEqual(sk, operation.getPrivateKeyEncoded(privParams)));

                        // KEM Encapsulation
                        EncapsulatedSecretGenerator kemGenerator = operation.getKEMGenerator(random);
                        SecretWithEncapsulation secretWithEnc = kemGenerator.generateEncapsulated(pubParams);

                        byte[] cipherText = secretWithEnc.getEncapsulation();
                        Assert.assertTrue(name + " " + count + ": cipherText", Arrays.areEqual(ct, cipherText));

                        byte[] encapSecret = secretWithEnc.getSecret();
                        Assert.assertTrue(name + " " + count + ": encapSecret", Arrays.areEqual(ss, encapSecret));

                        // KEM Decapsulation
                        EncapsulatedSecretExtractor kemExtractor = operation.getKEMExtractor(privParams);

                        byte[] decapSecret = kemExtractor.extractSecret(cipherText);
                        Assert.assertTrue(name + " " + count + ": decapSecret", Arrays.areEqual(ss, decapSecret));

                        //Assert.assertTrue(operation.getSessionKeySize() == secret.length * 8);
                    }
                    buf.clear();
                    continue;
                }
                int a = line.indexOf("=");
                if (a > -1)
                {
                    buf.put(line.substring(0, a).trim(), line.substring(a + 1).trim());
                }
            }
            // System.out.println("Testing successful!");
        }
    }
}
