/*
 * Copyright (c) 2019 the Eclipse Milo Authors
 *
 * This program and the accompanying materials are made
 * available under the terms of the Eclipse Public License 2.0
 * which is available at https://www.eclipse.org/legal/epl-2.0/
 *
 * SPDX-License-Identifier: EPL-2.0
 */

package org.eclipse.milo.opcua.stack;

import java.security.Security;
import java.util.ArrayList;
import java.util.List;

import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.eclipse.milo.opcua.stack.client.transport.uasc.ClientSecureChannel;
import org.eclipse.milo.opcua.stack.core.channel.ChannelParameters;
import org.eclipse.milo.opcua.stack.core.channel.ChunkDecoder;
import org.eclipse.milo.opcua.stack.core.channel.ChunkEncoder;
import org.eclipse.milo.opcua.stack.core.channel.EncodingLimits;
import org.eclipse.milo.opcua.stack.core.channel.MessageEncodeException;
import org.eclipse.milo.opcua.stack.core.channel.SecureChannel;
import org.eclipse.milo.opcua.stack.core.channel.ServerSecureChannel;
import org.eclipse.milo.opcua.stack.core.channel.messages.MessageType;
import org.eclipse.milo.opcua.stack.core.security.SecurityPolicy;
import org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.UInteger;
import org.eclipse.milo.opcua.stack.core.types.enumerated.MessageSecurityMode;
import org.eclipse.milo.opcua.stack.core.util.BufferUtil;
import org.eclipse.milo.opcua.stack.core.util.LongSequence;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import static org.eclipse.milo.opcua.stack.core.channel.EncodingLimits.DEFAULT_MAX_CHUNK_SIZE;
import static org.eclipse.milo.opcua.stack.core.channel.EncodingLimits.DEFAULT_MAX_MESSAGE_SIZE;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.fail;

public class ChunkSerializationTest extends SecureChannelFixture {

    static {
        // Required for SecurityPolicy.Aes256_Sha256_RsaPss
        Security.addProvider(new BouncyCastleProvider());
    }

    private Logger logger = LoggerFactory.getLogger(getClass());

    private ChannelParameters smallParameters = new ChannelParameters(
        32 * 8196,
        8196,
        8196,
        64,
        32 * 8196,
        8196,
        8196,
        64
    );

    private ChannelParameters defaultParameters = new ChannelParameters(
        DEFAULT_MAX_MESSAGE_SIZE,
        DEFAULT_MAX_CHUNK_SIZE,
        DEFAULT_MAX_CHUNK_SIZE,
        0,
        DEFAULT_MAX_MESSAGE_SIZE,
        DEFAULT_MAX_CHUNK_SIZE,
        DEFAULT_MAX_CHUNK_SIZE,
        0
    );

    private ChannelParameters unlimitedChunkCountParameters = new ChannelParameters(
        EncodingLimits.DEFAULT_MAX_MESSAGE_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        0,
        EncodingLimits.DEFAULT_MAX_MESSAGE_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        0
    );

    private ChannelParameters unlimitedMessageSizeParameters = new ChannelParameters(
        0,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        0,
        0,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        EncodingLimits.DEFAULT_MAX_CHUNK_SIZE,
        0
    );


    @DataProvider
    public Object[][] getAsymmetricSecurityParameters() {
        return new Object[][]{
            {SecurityPolicy.None, MessageSecurityMode.None, 128},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.Sign, 128},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.SignAndEncrypt, 128},
            {SecurityPolicy.Basic256, MessageSecurityMode.Sign, 128},
            {SecurityPolicy.Basic256, MessageSecurityMode.SignAndEncrypt, 128},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.Sign, 128},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.SignAndEncrypt, 128},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.Sign, 128},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.SignAndEncrypt, 128},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.Sign, 128},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.SignAndEncrypt, 128},

            {SecurityPolicy.None, MessageSecurityMode.None, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.Sign, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.SignAndEncrypt, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic256, MessageSecurityMode.Sign, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic256, MessageSecurityMode.SignAndEncrypt, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.Sign, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.SignAndEncrypt, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.Sign, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.SignAndEncrypt, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.Sign, DEFAULT_MAX_CHUNK_SIZE},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.SignAndEncrypt, DEFAULT_MAX_CHUNK_SIZE},
        };
    }

    @Test
    public void testAsymmetric4096() throws Exception {
        ChannelParameters parameters = defaultParameters;

        ChunkEncoder encoder = new ChunkEncoder(parameters);

        ChunkDecoder decoder = new ChunkDecoder(
            parameters,
            EncodingLimits.DEFAULT
        );

        SecureChannel[] channels = generateChannels4096();

        ClientSecureChannel clientChannel = (ClientSecureChannel) channels[0];
        ServerSecureChannel serverChannel = (ServerSecureChannel) channels[1];

        clientChannel
            .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE)
            .setIfAbsent(new LongSequence(1L, UInteger.MAX_VALUE));

        LongSequence requestId = clientChannel
            .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE).get();

        for (int messageSize = 0; messageSize < 512; messageSize++) {
            byte[] messageBytes = new byte[messageSize];
            for (int i = 0; i < messageBytes.length; i++) {
                messageBytes[i] = (byte) i;
            }

            ByteBuf messageBuffer = BufferUtil.pooledBuffer().writeBytes(messageBytes);

            List<ByteBuf> chunkBuffers = new ArrayList<>();

            try {
                ChunkEncoder.EncodedMessage message = encoder.encodeAsymmetric(
                    clientChannel,
                    requestId.getAndIncrement(),
                    messageBuffer,
                    MessageType.OpenSecureChannel
                );

                chunkBuffers.addAll(message.getMessageChunks());
            } catch (MessageEncodeException e) {
                fail("encoding error", e);
            }

            try {
                ChunkDecoder.DecodedMessage decodedMessage =
                    decoder.decodeAsymmetric(serverChannel, chunkBuffers);

                ByteBuf message = decodedMessage.getMessage();

                messageBuffer.readerIndex(0);
                assertEquals(message, messageBuffer);

                ReferenceCountUtil.release(message);
                ReferenceCountUtil.release(messageBuffer);
            } catch (Throwable t) {
                fail("decoding error", t);
            }
        }
    }

    @Test
    public void testSymmetric4096() throws Exception {
        ChannelParameters parameters = defaultParameters;

        ChunkEncoder encoder = new ChunkEncoder(parameters);

        ChunkDecoder decoder = new ChunkDecoder(
            parameters,
            EncodingLimits.DEFAULT
        );

        SecureChannel[] channels = generateChannels4096();

        ClientSecureChannel clientChannel = (ClientSecureChannel) channels[0];
        ServerSecureChannel serverChannel = (ServerSecureChannel) channels[1];

        clientChannel
            .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE)
            .setIfAbsent(new LongSequence(1L, UInteger.MAX_VALUE));

        LongSequence requestId = clientChannel
            .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE).get();

        for (int messageSize = 0; messageSize < 1024; messageSize++) {
            byte[] messageBytes = new byte[messageSize];
            for (int i = 0; i < messageBytes.length; i++) {
                messageBytes[i] = (byte) i;
            }

            ByteBuf messageBuffer = BufferUtil.pooledBuffer().writeBytes(messageBytes);

            List<ByteBuf> chunkBuffers = new ArrayList<>();

            try {
                ChunkEncoder.EncodedMessage message = encoder.encodeSymmetric(
                    clientChannel,
                    requestId.getAndIncrement(),
                    messageBuffer,
                    MessageType.OpenSecureChannel
                );

                chunkBuffers.addAll(message.getMessageChunks());
            } catch (MessageEncodeException e) {
                fail("encoding error", e);
            }

            try {
                ChunkDecoder.DecodedMessage decodedMessage =
                    decoder.decodeSymmetric(serverChannel, chunkBuffers);

                ByteBuf message = decodedMessage.getMessage();

                messageBuffer.readerIndex(0);
                assertEquals(message, messageBuffer);

                ReferenceCountUtil.release(message);
                ReferenceCountUtil.release(messageBuffer);
            } catch (Throwable t) {
                fail("decoding error", t);
            }
        }
    }


    @Test(dataProvider = "getAsymmetricSecurityParameters")
    public void testAsymmetricMessage(SecurityPolicy securityPolicy,
                                      MessageSecurityMode messageSecurity,
                                      int messageSize) throws Exception {

        logger.info("Asymmetric chunk serialization, securityPolicy={}, messageSecurityMode={}, messageSize={}",
            securityPolicy, messageSecurity, messageSize);

        ChannelParameters[] channelParameters = {
            smallParameters,
            defaultParameters,
            unlimitedChunkCountParameters,
            unlimitedMessageSizeParameters
        };

        for (ChannelParameters parameters : channelParameters) {
            ChunkEncoder encoder = new ChunkEncoder(parameters);

            ChunkDecoder decoder = new ChunkDecoder(
                parameters,
                EncodingLimits.DEFAULT
            );

            SecureChannel[] channels = generateChannels(securityPolicy, messageSecurity);
            ClientSecureChannel clientChannel = (ClientSecureChannel) channels[0];
            ServerSecureChannel serverChannel = (ServerSecureChannel) channels[1];

            clientChannel
                .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE)
                .setIfAbsent(new LongSequence(1L, UInteger.MAX_VALUE));

            LongSequence requestId = clientChannel
                .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE).get();

            byte[] messageBytes = new byte[messageSize];
            for (int i = 0; i < messageBytes.length; i++) {
                messageBytes[i] = (byte) i;
            }

            ByteBuf messageBuffer = BufferUtil.pooledBuffer().writeBytes(messageBytes);

            List<ByteBuf> chunkBuffers = new ArrayList<>();

            try {
                ChunkEncoder.EncodedMessage message = encoder.encodeAsymmetric(
                    clientChannel,
                    requestId.getAndIncrement(),
                    messageBuffer,
                    MessageType.OpenSecureChannel
                );

                chunkBuffers.addAll(message.getMessageChunks());
            } catch (MessageEncodeException e) {
                fail("encoding error", e);
            }

            try {
                ChunkDecoder.DecodedMessage decodedMessage =
                    decoder.decodeAsymmetric(serverChannel, chunkBuffers);

                ByteBuf message = decodedMessage.getMessage();

                messageBuffer.readerIndex(0);
                assertEquals(message, messageBuffer);

                ReferenceCountUtil.release(message);
                ReferenceCountUtil.release(messageBuffer);
            } catch (Throwable t) {
                fail("decoding error", t);
            }
        }
    }

    @DataProvider
    public Object[][] getSymmetricSecurityParameters() {
        return new Object[][]{
            {SecurityPolicy.None, MessageSecurityMode.None},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.Sign},
            {SecurityPolicy.Basic128Rsa15, MessageSecurityMode.SignAndEncrypt},
            {SecurityPolicy.Basic256, MessageSecurityMode.Sign},
            {SecurityPolicy.Basic256, MessageSecurityMode.SignAndEncrypt},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.Sign},
            {SecurityPolicy.Basic256Sha256, MessageSecurityMode.SignAndEncrypt},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.Sign},
            {SecurityPolicy.Aes128_Sha256_RsaOaep, MessageSecurityMode.SignAndEncrypt},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.Sign},
            {SecurityPolicy.Aes256_Sha256_RsaPss, MessageSecurityMode.SignAndEncrypt}
        };
    }

    @Test(dataProvider = "getSymmetricSecurityParameters")
    public void testSymmetricMessage(SecurityPolicy securityPolicy,
                                     MessageSecurityMode messageSecurity) throws Exception {

        logger.info(
            "Symmetric chunk serialization, " +
                "securityPolicy={}, messageSecurityMode={}",
            securityPolicy, messageSecurity);


        ChannelParameters[] channelParameters = {
            smallParameters,
            defaultParameters,
            unlimitedChunkCountParameters,
            unlimitedMessageSizeParameters
        };

        for (ChannelParameters parameters : channelParameters) {
            int[] messageSizes = new int[]{
                128,
                parameters.getRemoteMaxMessageSize()
            };

            for (int messageSize : messageSizes) {
                ChunkEncoder encoder = new ChunkEncoder(parameters);

                ChunkDecoder decoder = new ChunkDecoder(
                    parameters,
                    EncodingLimits.DEFAULT
                );

                SecureChannel[] channels = generateChannels(securityPolicy, messageSecurity);
                ClientSecureChannel clientChannel = (ClientSecureChannel) channels[0];
                ServerSecureChannel serverChannel = (ServerSecureChannel) channels[1];

                clientChannel
                    .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE)
                    .setIfAbsent(new LongSequence(1L, UInteger.MAX_VALUE));

                LongSequence requestId = clientChannel
                    .attr(ClientSecureChannel.KEY_REQUEST_ID_SEQUENCE).get();

                byte[] messageBytes = new byte[messageSize];
                for (int i = 0; i < messageBytes.length; i++) {
                    messageBytes[i] = (byte) i;
                }

                ByteBuf messageBuffer = BufferUtil.pooledBuffer().writeBytes(messageBytes);

                List<ByteBuf> chunkBuffers = new ArrayList<>();

                try {
                    ChunkEncoder.EncodedMessage message = encoder.encodeSymmetric(
                        clientChannel,
                        requestId.getAndIncrement(),
                        messageBuffer,
                        MessageType.SecureMessage
                    );

                    chunkBuffers.addAll(message.getMessageChunks());
                } catch (MessageEncodeException e) {
                    fail("encoding error", e);
                }

                try {
                    ChunkDecoder.DecodedMessage decodedMessage =
                        decoder.decodeSymmetric(serverChannel, chunkBuffers);

                    ByteBuf message = decodedMessage.getMessage();

                    messageBuffer.readerIndex(0);
                    assertEquals(message, messageBuffer);

                    ReferenceCountUtil.release(messageBuffer);
                    ReferenceCountUtil.release(message);
                } catch (Throwable t) {
                    fail("decoding error", t);
                }
            }
        }
    }

}
