/*
 * 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.core.channel;

import java.nio.Buffer;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.security.cert.Certificate;
import java.util.ArrayList;
import java.util.List;
import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import org.eclipse.milo.opcua.stack.core.StatusCodes;
import org.eclipse.milo.opcua.stack.core.UaException;
import org.eclipse.milo.opcua.stack.core.channel.headers.AsymmetricSecurityHeader;
import org.eclipse.milo.opcua.stack.core.channel.headers.SecureMessageHeader;
import org.eclipse.milo.opcua.stack.core.channel.headers.SequenceHeader;
import org.eclipse.milo.opcua.stack.core.channel.headers.SymmetricSecurityHeader;
import org.eclipse.milo.opcua.stack.core.channel.messages.MessageType;
import org.eclipse.milo.opcua.stack.core.security.SecurityAlgorithm;
import org.eclipse.milo.opcua.stack.core.util.BufferUtil;
import org.eclipse.milo.opcua.stack.core.util.LongSequence;
import org.eclipse.milo.opcua.stack.core.util.SignatureUtil;

import static org.eclipse.milo.opcua.stack.core.channel.headers.SecureMessageHeader.SECURE_MESSAGE_HEADER_SIZE;
import static org.eclipse.milo.opcua.stack.core.channel.headers.SequenceHeader.SEQUENCE_HEADER_SIZE;

public final class ChunkEncoder {

    private final AsymmetricEncoder asymmetricEncoder = new AsymmetricEncoder();
    private final SymmetricEncoder symmetricEncoder = new SymmetricEncoder();

    // Wrap after UInt32.MAX - 1024
    private final LongSequence sequenceNumber = new LongSequence(1L, 4294966271L);

    private final ChannelParameters parameters;

    public ChunkEncoder(ChannelParameters parameters) {
        this.parameters = parameters;
    }

    public EncodedMessage encodeAsymmetric(
        SecureChannel channel,
        long requestId,
        ByteBuf messageBuffer,
        MessageType messageType
    ) throws MessageEncodeException {

        return encode(asymmetricEncoder, channel, requestId, messageBuffer, messageType);
    }

    public EncodedMessage encodeSymmetric(
        SecureChannel channel,
        long requestId,
        ByteBuf messageBuffer,
        MessageType messageType
    ) throws MessageEncodeException {

        return encode(symmetricEncoder, channel, requestId, messageBuffer, messageType);
    }

    private EncodedMessage encode(
        AbstractEncoder encoder,
        SecureChannel channel,
        long requestId,
        ByteBuf messageBuffer,
        MessageType messageType
    ) throws MessageEncodeException {

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

        try {
            return encoder.encode(chunks, channel, requestId, messageBuffer, messageType);
        } catch (UaException e) {
            chunks.forEach(ReferenceCountUtil::safeRelease);

            throw new MessageEncodeException(e);
        }
    }

    private abstract class AbstractEncoder {

        EncodedMessage encode(
            List<ByteBuf> chunks,
            SecureChannel channel,
            long requestId,
            ByteBuf messageBuffer,
            MessageType messageType
        ) throws UaException {

            boolean encrypted = isEncryptionEnabled(channel);

            int securityHeaderSize = getSecurityHeaderSize(channel);
            int cipherTextBlockSize = getCipherTextBlockSize(channel);
            int plainTextBlockSize = getPlainTextBlockSize(channel);
            int signatureSize = getSignatureSize(channel);

            int maxChunkSize = parameters.getLocalSendBufferSize();
            int paddingOverhead = encrypted ? (cipherTextBlockSize > 256 ? 2 : 1) : 0;

            int maxCipherTextSize = maxChunkSize - SECURE_MESSAGE_HEADER_SIZE - securityHeaderSize;
            int maxCipherTextBlocks = maxCipherTextSize / cipherTextBlockSize;
            int maxPlainTextSize = maxCipherTextBlocks * plainTextBlockSize;
            int maxBodySize = maxPlainTextSize - SEQUENCE_HEADER_SIZE - paddingOverhead - signatureSize;

            assert (maxPlainTextSize + securityHeaderSize + SECURE_MESSAGE_HEADER_SIZE <= maxChunkSize);

            while (messageBuffer.readableBytes() > 0) {
                int bodySize = Math.min(messageBuffer.readableBytes(), maxBodySize);

                int paddingSize;
                if (encrypted) {
                    int plainTextSize = SEQUENCE_HEADER_SIZE + bodySize + paddingOverhead + signatureSize;
                    int remaining = plainTextSize % plainTextBlockSize;
                    paddingSize = remaining > 0 ? plainTextBlockSize - remaining : 0;
                } else {
                    paddingSize = 0;
                }

                int plainTextContentSize = SEQUENCE_HEADER_SIZE + bodySize +
                    signatureSize + paddingSize + paddingOverhead;

                assert (plainTextContentSize % plainTextBlockSize == 0);

                int chunkSize = SecureMessageHeader.SECURE_MESSAGE_HEADER_SIZE + securityHeaderSize +
                    (plainTextContentSize / plainTextBlockSize) * cipherTextBlockSize;

                assert (chunkSize <= maxChunkSize);

                ByteBuf chunkBuffer = BufferUtil.pooledBuffer(chunkSize);

                chunks.add(chunkBuffer);

                int remoteMaxChunkCount = parameters.getRemoteMaxChunkCount();
                if (remoteMaxChunkCount > 0 && chunks.size() > remoteMaxChunkCount) {
                    throw new UaException(
                        StatusCodes.Bad_EncodingLimitsExceeded,
                        "remote chunk count exceeded: " + remoteMaxChunkCount);
                }

                /* Message Header */
                SecureMessageHeader messageHeader = new SecureMessageHeader(
                    messageType,
                    messageBuffer.readableBytes() > bodySize ? 'C' : 'F',
                    chunkSize,
                    channel.getChannelId()
                );

                SecureMessageHeader.encode(messageHeader, chunkBuffer);

                /* Security Header */
                encodeSecurityHeader(channel, chunkBuffer);

                /* Sequence Header */
                SequenceHeader sequenceHeader = new SequenceHeader(
                    sequenceNumber.getAndIncrement(),
                    requestId
                );

                SequenceHeader.encode(sequenceHeader, chunkBuffer);

                /* Message Body */
                chunkBuffer.writeBytes(messageBuffer, bodySize);

                /* Padding and Signature */
                if (encrypted) {
                    writePadding(cipherTextBlockSize, paddingSize, chunkBuffer);
                }

                if (isSigningEnabled(channel)) {
                    ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer(0, chunkBuffer.writerIndex());

                    byte[] signature = signChunk(channel, chunkNioBuffer);

                    chunkBuffer.writeBytes(signature);
                }

                /* Encryption */
                if (encrypted) {
                    chunkBuffer.readerIndex(SECURE_MESSAGE_HEADER_SIZE + securityHeaderSize);

                    assert (chunkBuffer.readableBytes() % plainTextBlockSize == 0);

                    try {
                        int blockCount = chunkBuffer.readableBytes() / plainTextBlockSize;

                        ByteBuffer chunkNioBuffer = chunkBuffer.nioBuffer(
                            chunkBuffer.readerIndex(), blockCount * cipherTextBlockSize);

                        ByteBuf copyBuffer = chunkBuffer.copy();
                        ByteBuffer plainTextNioBuffer = copyBuffer.nioBuffer();

                        Cipher cipher = getCipher(channel);

                        if (isAsymmetric()) {
                            for (int blockNumber = 0; blockNumber < blockCount; blockNumber++) {
                                int position = blockNumber * plainTextBlockSize;
                                int limit = (blockNumber + 1) * plainTextBlockSize;
                                ((Buffer) plainTextNioBuffer).position(position);
                                ((Buffer) plainTextNioBuffer).limit(limit);

                                int bytesWritten = cipher.doFinal(plainTextNioBuffer, chunkNioBuffer);

                                assert (bytesWritten == cipherTextBlockSize);
                            }
                        } else {
                            cipher.doFinal(plainTextNioBuffer, chunkNioBuffer);
                        }

                        copyBuffer.release();
                    } catch (GeneralSecurityException e) {
                        throw new UaException(StatusCodes.Bad_SecurityChecksFailed, e);
                    }
                }

                chunkBuffer.readerIndex(0).writerIndex(chunkSize);
            }

            return new EncodedMessage(chunks, requestId);
        }

        private void writePadding(int cipherTextBlockSize, int paddingSize, ByteBuf buffer) {
            buffer.writeByte(paddingSize);

            for (int i = 0; i < paddingSize; i++) {
                buffer.writeByte(paddingSize);
            }

            if (cipherTextBlockSize > 256) {
                // Add the extra padding byte containing
                // the MSB of the 2-byte padding length.
                int paddingLengthMsb = (paddingSize >> 8) & 0xFF;
                buffer.writeByte(paddingLengthMsb);
            }
        }

        protected abstract byte[] signChunk(SecureChannel channel, ByteBuffer chunkNioBuffer) throws UaException;

        protected abstract void encodeSecurityHeader(SecureChannel channel, ByteBuf buffer) throws UaException;

        protected abstract Cipher getCipher(SecureChannel channel) throws UaException;

        protected abstract int getSecurityHeaderSize(SecureChannel channel) throws UaException;

        protected abstract int getCipherTextBlockSize(SecureChannel channel);

        protected abstract int getPlainTextBlockSize(SecureChannel channel);

        protected abstract int getSignatureSize(SecureChannel channel);

        protected abstract boolean isAsymmetric();

        protected abstract boolean isEncryptionEnabled(SecureChannel channel);

        protected abstract boolean isSigningEnabled(SecureChannel channel);

    }

    /**
     * A fully encoded message, in one more more chunks ready to send.
     */
    public static class EncodedMessage {

        private final List<ByteBuf> messageChunks;
        private final long requestId;

        public EncodedMessage(List<ByteBuf> messageChunks, long requestId) {
            this.messageChunks = messageChunks;
            this.requestId = requestId;
        }

        public List<ByteBuf> getMessageChunks() {
            return messageChunks;
        }

        public long getRequestId() {
            return requestId;
        }

    }

    private final class AsymmetricEncoder extends AbstractEncoder {

        @Override
        public byte[] signChunk(SecureChannel channel, ByteBuffer chunkNioBuffer) throws UaException {
            return SignatureUtil.sign(
                channel.getSecurityPolicy().getAsymmetricSignatureAlgorithm(),
                channel.getKeyPair().getPrivate(),
                chunkNioBuffer
            );
        }

        @Override
        public Cipher getCipher(SecureChannel channel) throws UaException {
            Certificate remoteCertificate = channel.getRemoteCertificate();

            assert (remoteCertificate != null);

            try {
                String transformation = channel.getSecurityPolicy()
                    .getAsymmetricEncryptionAlgorithm().getTransformation();
                Cipher cipher = Cipher.getInstance(transformation);
                cipher.init(Cipher.ENCRYPT_MODE, remoteCertificate.getPublicKey());
                return cipher;
            } catch (GeneralSecurityException e) {
                throw new UaException(StatusCodes.Bad_SecurityChecksFailed, e);
            }
        }

        @Override
        public void encodeSecurityHeader(SecureChannel channel, ByteBuf buffer) throws UaException {
            AsymmetricSecurityHeader header = new AsymmetricSecurityHeader(
                channel.getSecurityPolicy().getUri(),
                channel.getLocalCertificateChainBytes(),
                channel.getRemoteCertificateThumbprint()
            );

            AsymmetricSecurityHeader.encode(header, buffer);
        }

        @Override
        public int getSecurityHeaderSize(SecureChannel channel) throws UaException {
            String securityPolicyUri = channel.getSecurityPolicy().getUri();
            byte[] localCertificateChainBytes = channel.getLocalCertificateChainBytes().bytes();
            byte[] remoteCertificateThumbprint = channel.getRemoteCertificateThumbprint().bytes();

            return 12 + securityPolicyUri.length() +
                (localCertificateChainBytes != null ? localCertificateChainBytes.length : 0) +
                (remoteCertificateThumbprint != null ? remoteCertificateThumbprint.length : 0);
        }

        @Override
        public int getCipherTextBlockSize(SecureChannel channel) {
            return channel.getRemoteAsymmetricCipherTextBlockSize();
        }

        @Override
        public int getPlainTextBlockSize(SecureChannel channel) {
            return channel.getRemoteAsymmetricPlainTextBlockSize();
        }

        @Override
        public int getSignatureSize(SecureChannel channel) {
            return channel.getLocalAsymmetricSignatureSize();
        }

        @Override
        protected boolean isAsymmetric() {
            return true;
        }

        @Override
        public boolean isEncryptionEnabled(SecureChannel channel) {
            return channel.isAsymmetricEncryptionEnabled();
        }

        @Override
        public boolean isSigningEnabled(SecureChannel channel) {
            return channel.isAsymmetricSigningEnabled();
        }

    }

    private final class SymmetricEncoder extends AbstractEncoder {

        private volatile ChannelSecurity.SecurityKeys securityKeys;
        private volatile Cipher cipher = null;
        private volatile long cipherId = -1;

        @Override
        public void encodeSecurityHeader(SecureChannel channel, ByteBuf buffer) throws UaException {
            ChannelSecurity channelSecurity = channel.getChannelSecurity();
            long tokenId = channelSecurity != null ? channelSecurity.getCurrentToken().getTokenId().longValue() : 0L;

            SymmetricSecurityHeader.encode(new SymmetricSecurityHeader(tokenId), buffer);

            securityKeys = channelSecurity != null ? channelSecurity.getCurrentKeys() : null;

            if (cipherId != tokenId && channel.isSymmetricEncryptionEnabled()) {
                cipher = initCipher(channel);
                cipherId = tokenId;
            }
        }

        @Override
        public byte[] signChunk(SecureChannel channel, ByteBuffer chunkNioBuffer) throws UaException {
            SecurityAlgorithm signatureAlgorithm = channel.getSecurityPolicy().getSymmetricSignatureAlgorithm();
            byte[] signatureKey = channel.getEncryptionKeys(securityKeys).getSignatureKey();

            return SignatureUtil.hmac(
                signatureAlgorithm,
                signatureKey,
                chunkNioBuffer
            );
        }

        @Override
        public Cipher getCipher(SecureChannel channel) {
            assert cipher != null;
            return cipher;
        }

        @Override
        public int getSecurityHeaderSize(SecureChannel channel) {
            return SymmetricSecurityHeader.SYMMETRIC_SECURITY_HEADER_SIZE;
        }

        @Override
        public int getCipherTextBlockSize(SecureChannel channel) {
            return channel.getSymmetricBlockSize();
        }

        @Override
        public int getPlainTextBlockSize(SecureChannel channel) {
            return channel.getSymmetricBlockSize();
        }

        @Override
        public int getSignatureSize(SecureChannel channel) {
            return channel.getSymmetricSignatureSize();
        }

        @Override
        protected boolean isAsymmetric() {
            return false;
        }

        @Override
        public boolean isEncryptionEnabled(SecureChannel channel) {
            return channel.isSymmetricEncryptionEnabled();
        }

        @Override
        public boolean isSigningEnabled(SecureChannel channel) {
            return channel.isSymmetricSigningEnabled();
        }

        private Cipher initCipher(SecureChannel channel) throws UaException {
            try {
                String transformation = channel.getSecurityPolicy()
                    .getSymmetricEncryptionAlgorithm().getTransformation();
                ChannelSecurity.SecretKeys secretKeys = channel.getEncryptionKeys(securityKeys);

                SecretKeySpec keySpec = new SecretKeySpec(secretKeys.getEncryptionKey(), "AES");
                IvParameterSpec ivSpec = new IvParameterSpec(secretKeys.getInitializationVector());

                Cipher cipher = Cipher.getInstance(transformation);
                cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);

                assert (cipher.getBlockSize() == channel.getSymmetricBlockSize());

                return cipher;
            } catch (GeneralSecurityException e) {
                throw new UaException(StatusCodes.Bad_SecurityChecksFailed, e);
            }
        }

    }

}
