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.TestCase;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.pqc.crypto.uov.UOVKeyPairGenerator;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.pqc.crypto.uov.UOVKeyGenerationParameters;
import org.bouncycastle.pqc.crypto.uov.UOVParameters;
import org.bouncycastle.pqc.crypto.uov.UOVPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.uov.UOVPublicKeyParameters;
import org.bouncycastle.pqc.crypto.uov.UOVSigner;
import org.bouncycastle.test.TestResourceFinder;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.util.test.FixedSecureRandom;

/**
 * Tests for the UOV signer ported from the pqov reference implementation
 * (https://github.com/pqov/pqov). Covers all three encoding variants:
 * classic (uncompressed), PKC (compressed pk), PKC-SKC (compressed pk and
 * seed-only sk). The expected pk/sk/sig vectors under
 * bc-test-data/pqc/crypto/uov/ref/ were generated by
 * pqov-reference/dump_vec.c with a fixed deterministic salt.
 */
public class UOVTest
    extends TestCase
{
    private static final UOVParameters[] PARAMETER_SETS = new UOVParameters[]{
        UOVParameters.uov_Is,
        UOVParameters.uov_Is_pkc,
        UOVParameters.uov_Is_pkc_skc,
        UOVParameters.uov_Ip,
        UOVParameters.uov_Ip_pkc,
        UOVParameters.uov_Ip_pkc_skc,
        UOVParameters.uov_III,
        UOVParameters.uov_III_pkc,
        UOVParameters.uov_III_pkc_skc,
        UOVParameters.uov_V,
        UOVParameters.uov_V_pkc,
        UOVParameters.uov_V_pkc_skc,
    };

    public void testParameterSizes()
    {
        // pqov README — classic sizes per security level.
        check(UOVParameters.uov_Is, 412160, 348704, 96);
        check(UOVParameters.uov_Ip, 278432, 237896, 128);
        check(UOVParameters.uov_III, 1225440, 1044320, 200);
        check(UOVParameters.uov_V, 2869440, 2436704, 260);

        // pqov README — PKC sizes (compressed pk: pk_seed + P3).
        check(UOVParameters.uov_Is_pkc, 66576, 348704, 96);
        check(UOVParameters.uov_Ip_pkc, 43576, 237896, 128);
        check(UOVParameters.uov_III_pkc, 189232, 1044320, 200);
        check(UOVParameters.uov_V_pkc, 446992, 2436704, 260);

        // pqov README — PKC-SKC sizes (sk = sk_seed only).
        check(UOVParameters.uov_Is_pkc_skc, 66576, 32, 96);
        check(UOVParameters.uov_Ip_pkc_skc, 43576, 32, 128);
        check(UOVParameters.uov_III_pkc_skc, 189232, 32, 200);
        check(UOVParameters.uov_V_pkc_skc, 446992, 32, 260);
    }

    private void check(UOVParameters p, int expectedPk, int expectedSk, int expectedSig)
    {
        assertEquals(p.getName() + " pk", expectedPk, p.getPublicKeyBytes());
        assertEquals(p.getName() + " sk", expectedSk, p.getSecretKeyBytes());
        assertEquals(p.getName() + " sig", expectedSig, p.getSignatureBytes());
    }

    public void testReferenceVectors()
        throws IOException
    {
        // All four parameter sets × three encoding variants — pk, sk, and sig
        // must match the pqov reference byte-for-byte under a fixed
        // deterministic salt.
        readAndAssertVector("uov_Is_vec0.txt", UOVParameters.uov_Is);
        readAndAssertVector("uov_Is_pkc_vec0.txt", UOVParameters.uov_Is_pkc);
        readAndAssertVector("uov_Is_pkc_skc_vec0.txt", UOVParameters.uov_Is_pkc_skc);

        readAndAssertVector("uov_Ip_vec0.txt", UOVParameters.uov_Ip);
        readAndAssertVector("uov_Ip_vec1.txt", UOVParameters.uov_Ip);
        readAndAssertVector("uov_Ip_pkc_vec0.txt", UOVParameters.uov_Ip_pkc);
        readAndAssertVector("uov_Ip_pkc_skc_vec0.txt", UOVParameters.uov_Ip_pkc_skc);

        readAndAssertVector("uov_III_vec0.txt", UOVParameters.uov_III);
        readAndAssertVector("uov_III_pkc_vec0.txt", UOVParameters.uov_III_pkc);
        readAndAssertVector("uov_III_pkc_skc_vec0.txt", UOVParameters.uov_III_pkc_skc);

        readAndAssertVector("uov_V_vec0.txt", UOVParameters.uov_V);
        readAndAssertVector("uov_V_pkc_vec0.txt", UOVParameters.uov_V_pkc);
        readAndAssertVector("uov_V_pkc_skc_vec0.txt", UOVParameters.uov_V_pkc_skc);
    }

    public void testConsistencyRandomVerify()
    {
        // Round-trip across all 12 (param × variant) combinations with random
        // inputs; tampered signatures must fail.
        SecureRandom random = new SecureRandom();
        for (int idx = 0; idx != PARAMETER_SETS.length; idx++)
        {
            UOVParameters params = PARAMETER_SETS[idx];
            UOVKeyPairGenerator kpg = new UOVKeyPairGenerator();
            kpg.init(new UOVKeyGenerationParameters(random, params));
            AsymmetricCipherKeyPair kp = kpg.generateKeyPair();
            UOVPrivateKeyParameters priv = (UOVPrivateKeyParameters)kp.getPrivate();
            UOVPublicKeyParameters pub = (UOVPublicKeyParameters)kp.getPublic();

            assertEquals(params.getName(), params.getPublicKeyBytes(), pub.getEncoded().length);
            assertEquals(params.getName(), params.getSecretKeyBytes(), priv.getEncoded().length);

            byte[] msg = new byte[37];
            random.nextBytes(msg);

            UOVSigner signer = new UOVSigner();
            signer.init(true, new ParametersWithRandom(priv, random));
            byte[] sig = signer.generateSignature(msg);
            assertEquals(params.getName(), params.getSignatureBytes(), sig.length);

            UOVSigner verifier = new UOVSigner();
            verifier.init(false, pub);
            assertTrue(params.getName() + ": signature must verify", verifier.verifySignature(msg, sig));

            byte[] tampered = Arrays.clone(sig);
            tampered[0] ^= 0x01;
            assertFalse(params.getName() + ": tampered signature must fail",
                verifier.verifySignature(msg, tampered));
        }
    }

    public void testPrivateKeyRoundTrip()
    {
        SecureRandom random = new SecureRandom();
        for (int idx = 0; idx != PARAMETER_SETS.length; idx++)
        {
            UOVParameters params = PARAMETER_SETS[idx];
            UOVKeyPairGenerator kpg = new UOVKeyPairGenerator();
            kpg.init(new UOVKeyGenerationParameters(random, params));
            AsymmetricCipherKeyPair kp = kpg.generateKeyPair();

            UOVPrivateKeyParameters priv = (UOVPrivateKeyParameters)kp.getPrivate();
            byte[] enc = priv.getEncoded();
            UOVPrivateKeyParameters reread = new UOVPrivateKeyParameters(params, enc);
            assertTrue(params.getName(), Arrays.areEqual(enc, reread.getEncoded()));
            assertTrue(params.getName(), Arrays.areEqual(priv.getSeed(), reread.getSeed()));

            UOVPublicKeyParameters pub = (UOVPublicKeyParameters)kp.getPublic();
            byte[] pubEnc = pub.getEncoded();
            UOVPublicKeyParameters pubReread = new UOVPublicKeyParameters(params, pubEnc);
            assertTrue(params.getName(), Arrays.areEqual(pubEnc, pubReread.getEncoded()));
        }
    }

    private void readAndAssertVector(String resource, UOVParameters params)
        throws IOException
    {
        Map<String, String> kv = readFixture(resource);
        byte[] skSeed = Hex.decode(kv.get("sk_seed"));
        byte[] salt = Hex.decode(kv.get("salt"));
        byte[] msg = Hex.decode(kv.get("msg"));
        byte[] expectedPk = Hex.decode(kv.get("pk"));
        byte[] expectedSk = Hex.decode(kv.get("sk"));
        byte[] expectedSig = Hex.decode(kv.get("sig"));

        // UOVKeyPairGenerator pulls UOVParameters.SK_SEED_BYTES (32) bytes
        // from the supplied SecureRandom for the sk_seed; UOVSigner pulls
        // UOVParameters.SALT_BYTES from the signer-side random. Both inputs
        // are injected via FixedSecureRandom for KAT reproducibility.
        UOVKeyPairGenerator kpg = new UOVKeyPairGenerator();
        kpg.init(new UOVKeyGenerationParameters(new FixedSecureRandom(skSeed), params));
        AsymmetricCipherKeyPair kp = kpg.generateKeyPair();
        UOVPublicKeyParameters pub = (UOVPublicKeyParameters)kp.getPublic();
        UOVPrivateKeyParameters priv = (UOVPrivateKeyParameters)kp.getPrivate();

        assertTrue(resource + " pk mismatch", Arrays.areEqual(expectedPk, pub.getEncoded()));
        assertTrue(resource + " sk mismatch", Arrays.areEqual(expectedSk, priv.getEncoded()));

        UOVSigner signer = new UOVSigner();
        signer.init(true, new ParametersWithRandom(priv, new FixedSecureRandom(salt)));
        byte[] sig = signer.generateSignature(msg);
        assertTrue(resource + " sig mismatch", Arrays.areEqual(expectedSig, sig));

        UOVSigner verifier = new UOVSigner();
        verifier.init(false, pub);
        assertTrue(resource + " verify(own)", verifier.verifySignature(msg, sig));
        assertTrue(resource + " verify(ref)", verifier.verifySignature(msg, expectedSig));
    }

    private static Map<String, String> readFixture(String name)
        throws IOException
    {
        InputStream src = TestResourceFinder.findTestResource("pqc/crypto/uov/ref", name);
        BufferedReader r = new BufferedReader(new InputStreamReader(src));
        Map<String, String> out = new HashMap<String, String>();
        String line;
        while ((line = r.readLine()) != null)
        {
            int eq = line.indexOf('=');
            if (eq < 0)
            {
                continue;
            }
            out.put(line.substring(0, eq).trim(), line.substring(eq + 1).trim());
        }
        r.close();
        return out;
    }
}
