/*
 * 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.server.transport.uasc;

import java.io.IOException;
import java.security.KeyPair;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.CompositeByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.ByteToMessageDecoder;
import io.netty.util.AttributeKey;
import io.netty.util.ReferenceCountUtil;
import io.netty.util.Timeout;
import org.eclipse.milo.opcua.stack.core.Stack;
import org.eclipse.milo.opcua.stack.core.StatusCodes;
import org.eclipse.milo.opcua.stack.core.UaException;
import org.eclipse.milo.opcua.stack.core.UaSerializationException;
import org.eclipse.milo.opcua.stack.core.channel.ChannelSecurity;
import org.eclipse.milo.opcua.stack.core.channel.ChannelSecurity.SecurityKeys;
import org.eclipse.milo.opcua.stack.core.channel.ChunkDecoder;
import org.eclipse.milo.opcua.stack.core.channel.ChunkEncoder.EncodedMessage;
import org.eclipse.milo.opcua.stack.core.channel.ExceptionHandler;
import org.eclipse.milo.opcua.stack.core.channel.MessageAbortException;
import org.eclipse.milo.opcua.stack.core.channel.MessageDecodeException;
import org.eclipse.milo.opcua.stack.core.channel.MessageEncodeException;
import org.eclipse.milo.opcua.stack.core.channel.SerializationQueue;
import org.eclipse.milo.opcua.stack.core.channel.ServerSecureChannel;
import org.eclipse.milo.opcua.stack.core.channel.headers.AsymmetricSecurityHeader;
import org.eclipse.milo.opcua.stack.core.channel.headers.HeaderDecoder;
import org.eclipse.milo.opcua.stack.core.channel.messages.ErrorMessage;
import org.eclipse.milo.opcua.stack.core.channel.messages.MessageType;
import org.eclipse.milo.opcua.stack.core.security.CertificateManager;
import org.eclipse.milo.opcua.stack.core.security.CertificateValidator;
import org.eclipse.milo.opcua.stack.core.security.SecurityPolicy;
import org.eclipse.milo.opcua.stack.core.transport.TransportProfile;
import org.eclipse.milo.opcua.stack.core.types.builtin.ByteString;
import org.eclipse.milo.opcua.stack.core.types.builtin.DateTime;
import org.eclipse.milo.opcua.stack.core.types.builtin.StatusCode;
import org.eclipse.milo.opcua.stack.core.types.enumerated.SecurityTokenRequestType;
import org.eclipse.milo.opcua.stack.core.types.structured.ChannelSecurityToken;
import org.eclipse.milo.opcua.stack.core.types.structured.EndpointDescription;
import org.eclipse.milo.opcua.stack.core.types.structured.OpenSecureChannelRequest;
import org.eclipse.milo.opcua.stack.core.types.structured.OpenSecureChannelResponse;
import org.eclipse.milo.opcua.stack.core.types.structured.ResponseHeader;
import org.eclipse.milo.opcua.stack.core.util.BufferUtil;
import org.eclipse.milo.opcua.stack.core.util.EndpointUtil;
import org.eclipse.milo.opcua.stack.core.util.NonceUtil;
import org.eclipse.milo.opcua.stack.server.UaStackServer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static org.eclipse.milo.opcua.stack.core.types.builtin.unsigned.Unsigned.uint;
import static org.eclipse.milo.opcua.stack.core.util.NonceUtil.generateNonce;

public class UascServerAsymmetricHandler extends ByteToMessageDecoder implements HeaderDecoder {

    static final AttributeKey<EndpointDescription> ENDPOINT_KEY = AttributeKey.valueOf("endpoint");

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

    private ServerSecureChannel secureChannel;
    private Timeout secureChannelTimeout;

    private boolean symmetricHandlerAdded = false;

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

    private final AtomicReference<AsymmetricSecurityHeader> headerRef = new AtomicReference<>();

    private final int maxChunkCount;
    private final int maxChunkSize;

    private final UaStackServer stackServer;
    private final TransportProfile transportProfile;
    private final SerializationQueue serializationQueue;

    UascServerAsymmetricHandler(
        UaStackServer stackServer,
        TransportProfile transportProfile,
        SerializationQueue serializationQueue
    ) {

        this.stackServer = stackServer;
        this.transportProfile = transportProfile;
        this.serializationQueue = serializationQueue;

        maxChunkCount = serializationQueue.getParameters().getLocalMaxChunkCount();
        maxChunkSize = serializationQueue.getParameters().getLocalReceiveBufferSize();
    }

    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> out) throws Exception {
        if (buffer.readableBytes() >= HEADER_LENGTH) {
            int messageLength = getMessageLength(buffer, maxChunkSize);

            if (buffer.readableBytes() >= messageLength) {
                MessageType messageType = MessageType.fromMediumInt(
                    buffer.getMediumLE(buffer.readerIndex())
                );

                switch (messageType) {
                    case OpenSecureChannel:
                        onOpenSecureChannel(ctx, buffer.readSlice(messageLength));
                        break;

                    case CloseSecureChannel:
                        logger.debug("Received CloseSecureChannelRequest");

                        buffer.skipBytes(messageLength);

                        if (secureChannelTimeout != null) {
                            secureChannelTimeout.cancel();
                            secureChannelTimeout = null;
                        }

                        ctx.close();
                        break;

                    default:
                        throw new UaException(
                            StatusCodes.Bad_TcpMessageTypeInvalid,
                            "unexpected MessageType: " + messageType
                        );
                }
            }
        }
    }

    private void onOpenSecureChannel(ChannelHandlerContext ctx, ByteBuf buffer) throws UaException {
        buffer.skipBytes(3); // Skip messageType

        char chunkType = (char) buffer.readByte();

        if (chunkType == 'A') {
            chunkBuffers.forEach(ByteBuf::release);
            chunkBuffers.clear();
            headerRef.set(null);
        } else {
            buffer.skipBytes(4); // Skip messageSize

            final long secureChannelId = buffer.readUnsignedIntLE();

            final AsymmetricSecurityHeader header = AsymmetricSecurityHeader.decode(
                buffer,
                stackServer.getConfig().getEncodingLimits()
            );

            if (!headerRef.compareAndSet(null, header)) {
                if (!header.equals(headerRef.get())) {
                    throw new UaException(StatusCodes.Bad_SecurityChecksFailed,
                        "subsequent AsymmetricSecurityHeader did not match");
                }
            }

            if (secureChannelId != 0) {
                if (secureChannel == null) {
                    throw new UaException(StatusCodes.Bad_TcpSecureChannelUnknown,
                        "unknown secure channel id: " + secureChannelId);
                }

                if (secureChannelId != secureChannel.getChannelId()) {
                    throw new UaException(StatusCodes.Bad_TcpSecureChannelUnknown,
                        "unknown secure channel id: " + secureChannelId);
                }
            }

            if (secureChannel == null) {
                secureChannel = new ServerSecureChannel();
                secureChannel.setChannelId(stackServer.getNextChannelId());

                String securityPolicyUri = header.getSecurityPolicyUri();
                SecurityPolicy securityPolicy = SecurityPolicy.fromUri(securityPolicyUri);

                secureChannel.setSecurityPolicy(securityPolicy);

                if (securityPolicy != SecurityPolicy.None) {
                    secureChannel.setRemoteCertificate(header.getSenderCertificate().bytesOrEmpty());

                    CertificateValidator certificateValidator = stackServer.getConfig().getCertificateValidator();

                    certificateValidator.validateCertificateChain(secureChannel.getRemoteCertificateChain());

                    CertificateManager certificateManager = stackServer.getConfig().getCertificateManager();

                    Optional<X509Certificate[]> localCertificateChain = certificateManager
                        .getCertificateChain(header.getReceiverThumbprint());

                    Optional<KeyPair> keyPair = certificateManager
                        .getKeyPair(header.getReceiverThumbprint());

                    if (localCertificateChain.isPresent() && keyPair.isPresent()) {
                        X509Certificate[] chain = localCertificateChain.get();

                        secureChannel.setLocalCertificate(chain[0]);
                        secureChannel.setLocalCertificateChain(chain);
                        secureChannel.setKeyPair(keyPair.get());
                    } else {
                        throw new UaException(
                            StatusCodes.Bad_SecurityChecksFailed,
                            "no certificate for provided thumbprint");
                    }
                }
            }

            int chunkSize = buffer.readerIndex(0).readableBytes();

            if (chunkSize > maxChunkSize) {
                throw new UaException(StatusCodes.Bad_TcpMessageTooLarge,
                    String.format("max chunk size exceeded (%s)", maxChunkSize));
            }

            chunkBuffers.add(buffer.retain());

            if (maxChunkCount > 0 && chunkBuffers.size() > maxChunkCount) {
                throw new UaException(StatusCodes.Bad_TcpMessageTooLarge,
                    String.format("max chunk count exceeded (%s)", maxChunkCount));
            }

            if (chunkType == 'F') {
                final List<ByteBuf> buffersToDecode = chunkBuffers;

                chunkBuffers = new ArrayList<>();
                headerRef.set(null);

                serializationQueue.decode((binaryDecoder, chunkDecoder) -> {
                    ByteBuf message;
                    long requestId;

                    try {
                        ChunkDecoder.DecodedMessage decodedMessage =
                            chunkDecoder.decodeAsymmetric(secureChannel, buffersToDecode);

                        message = decodedMessage.getMessage();
                        requestId = decodedMessage.getRequestId();
                    } catch (MessageAbortException e) {
                        logger.warn(
                            "Received message abort chunk; error={}, reason={}",
                            e.getStatusCode(), e.getMessage()
                        );
                        return;
                    } catch (MessageDecodeException e) {
                        logger.error("Error decoding asymmetric message", e);

                        ctx.close();
                        return;
                    }

                    try {
                        OpenSecureChannelRequest request = (OpenSecureChannelRequest) binaryDecoder
                            .setBuffer(message)
                            .readMessage(null);

                        logger.debug(
                            "Received OpenSecureChannelRequest ({}, id={}).",
                            request.getRequestType(), secureChannelId
                        );

                        sendOpenSecureChannelResponse(ctx, requestId, request);
                    } catch (Throwable t) {
                        logger.error("Error decoding OpenSecureChannelRequest", t);

                        ctx.close();
                    } finally {
                        message.release();
                        buffersToDecode.clear();
                    }
                });
            }
        }
    }

    private void sendOpenSecureChannelResponse(
        ChannelHandlerContext ctx,
        long requestId,
        OpenSecureChannelRequest request
    ) {

        serializationQueue.encode((binaryEncoder, chunkEncoder) -> {
            ByteBuf messageBuffer = BufferUtil.pooledBuffer();

            try {
                OpenSecureChannelResponse response = openSecureChannel(ctx, request);

                binaryEncoder.setBuffer(messageBuffer);
                binaryEncoder.writeMessage(null, response);

                checkMessageSize(messageBuffer);

                EncodedMessage encodedMessage = chunkEncoder.encodeAsymmetric(
                    secureChannel,
                    requestId,
                    messageBuffer,
                    MessageType.OpenSecureChannel
                );

                if (!symmetricHandlerAdded) {
                    UascServerSymmetricHandler symmetricHandler = new UascServerSymmetricHandler(
                        stackServer,
                        serializationQueue,
                        secureChannel
                    );

                    ctx.pipeline().addBefore(ctx.name(), null, symmetricHandler);

                    symmetricHandlerAdded = true;
                }

                CompositeByteBuf chunkComposite = BufferUtil.compositeBuffer();

                for (ByteBuf chunk : encodedMessage.getMessageChunks()) {
                    chunkComposite.addComponent(chunk);
                    chunkComposite.writerIndex(chunkComposite.writerIndex() + chunk.readableBytes());
                }

                ctx.writeAndFlush(chunkComposite, ctx.voidPromise());

                logger.debug("Sent OpenSecureChannelResponse.");
            } catch (MessageEncodeException e) {
                logger.error("Error encoding OpenSecureChannelResponse: {}", e.getMessage(), e);
                ctx.fireExceptionCaught(e);
            } catch (UaSerializationException e) {
                logger.error("Error serializing OpenSecureChannelResponse: {}", e.getMessage(), e);
                ctx.fireExceptionCaught(e);
            } catch (UaException e) {
                logger.error("Error installing security token: {}", e.getStatusCode(), e);
                ctx.close();
            } finally {
                messageBuffer.release();
            }
        });
    }

    private OpenSecureChannelResponse openSecureChannel(
        ChannelHandlerContext ctx,
        OpenSecureChannelRequest request
    ) throws UaException {

        SecurityTokenRequestType requestType = request.getRequestType();

        if (requestType == SecurityTokenRequestType.Issue) {
            secureChannel.setMessageSecurityMode(request.getSecurityMode());

            String endpointUrl = ctx.channel().attr(UascServerHelloHandler.ENDPOINT_URL_KEY).get();

            EndpointDescription endpoint = stackServer.getEndpointDescriptions()
                .stream()
                .filter(e -> {
                    boolean transportMatch = Objects.equals(
                        e.getTransportProfileUri(),
                        transportProfile.getUri()
                    );

                    boolean pathMatch = Objects.equals(
                        EndpointUtil.getPath(e.getEndpointUrl()),
                        EndpointUtil.getPath(endpointUrl)
                    );

                    boolean securityPolicyMatch = Objects.equals(
                        e.getSecurityPolicyUri(),
                        secureChannel.getSecurityPolicy().getUri()
                    );

                    boolean securityModeMatch = Objects.equals(
                        e.getSecurityMode(),
                        request.getSecurityMode()
                    );

                    return transportMatch && pathMatch && securityPolicyMatch && securityModeMatch;
                })
                .findFirst()
                .orElseThrow(() -> {
                    String message = String.format(
                        "no matching endpoint found: transportProfile=%s, " +
                            "endpointUrl=%s, securityPolicy=%s, securityMode=%s",
                        transportProfile,
                        endpointUrl,
                        secureChannel.getSecurityPolicy(),
                        request.getSecurityMode()
                    );

                    return new UaException(StatusCodes.Bad_SecurityChecksFailed, message);
                });

            ctx.channel().attr(ENDPOINT_KEY).set(endpoint);
        }

        if (requestType == SecurityTokenRequestType.Renew &&
            secureChannel.getMessageSecurityMode() != request.getSecurityMode()) {

            throw new UaException(StatusCodes.Bad_SecurityChecksFailed,
                "secure channel renewal requested a different MessageSecurityMode.");
        }

        long channelLifetime = request.getRequestedLifetime().longValue();

        channelLifetime = Math.min(
            channelLifetime,
            stackServer.getConfig().getMaximumSecureChannelLifetime().longValue()
        );
        channelLifetime = Math.max(
            channelLifetime,
            stackServer.getConfig().getMinimumSecureChannelLifetime().longValue()
        );

        ChannelSecurityToken newToken = new ChannelSecurityToken(
            uint(secureChannel.getChannelId()),
            uint(stackServer.getNextTokenId()),
            DateTime.now(),
            uint(channelLifetime)
        );

        SecurityKeys newKeys = null;

        if (secureChannel.isSymmetricSigningEnabled()) {
            // Validate the remote nonce; it must be non-null and the correct length for the security algorithm.
            ByteString remoteNonce = request.getClientNonce();

            NonceUtil.validateNonce(remoteNonce, secureChannel.getSecurityPolicy());

            ByteString localNonce = generateNonce(secureChannel.getSecurityPolicy());

            secureChannel.setLocalNonce(localNonce);
            secureChannel.setRemoteNonce(remoteNonce);

            newKeys = ChannelSecurity.generateKeyPair(
                secureChannel,
                secureChannel.getRemoteNonce(),
                secureChannel.getLocalNonce()
            );
        }

        ChannelSecurity oldSecrets = secureChannel.getChannelSecurity();
        SecurityKeys oldKeys = oldSecrets != null ? oldSecrets.getCurrentKeys() : null;
        ChannelSecurityToken oldToken = oldSecrets != null ? oldSecrets.getCurrentToken() : null;

        ChannelSecurity newSecrets = new ChannelSecurity(
            newKeys,
            newToken,
            oldKeys,
            oldToken
        );

        secureChannel.setChannelSecurity(newSecrets);

        /*
         * Cancel the previous timeout, if it exists, and start a new one.
         */
        if (secureChannelTimeout == null || secureChannelTimeout.cancel()) {
            final long lifetime = channelLifetime;
            secureChannelTimeout = Stack.sharedWheelTimer().newTimeout(
                timeout -> {
                    logger.debug(
                        "SecureChannel renewal timed out after {}ms. id={}, channel={}",
                        lifetime, secureChannel.getChannelId(), ctx.channel());

                    ctx.close();
                },
                channelLifetime, TimeUnit.MILLISECONDS
            );
        }

        ResponseHeader responseHeader = new ResponseHeader(
            DateTime.now(),
            request.getRequestHeader().getRequestHandle(),
            StatusCode.GOOD,
            null,
            null,
            null
        );

        return new OpenSecureChannelResponse(
            responseHeader,
            uint(PROTOCOL_VERSION),
            newToken,
            secureChannel.getLocalNonce()
        );
    }

    private void checkMessageSize(ByteBuf messageBuffer) throws UaSerializationException {
        int messageSize = messageBuffer.readableBytes();
        int remoteMaxMessageSize = serializationQueue.getParameters().getRemoteMaxMessageSize();

        if (remoteMaxMessageSize > 0 && messageSize > remoteMaxMessageSize) {
            throw new UaSerializationException(
                StatusCodes.Bad_ResponseTooLarge,
                "response exceeds remote max message size: " +
                    messageSize + " > " + remoteMaxMessageSize);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        chunkBuffers.forEach(ReferenceCountUtil::safeRelease);
        chunkBuffers.clear();

        if (cause instanceof IOException) {
            ctx.close();
            logger.debug("[remote={}] IOException caught; channel closed",
                ctx.channel().remoteAddress(), cause);
        } else {
            ErrorMessage errorMessage = ExceptionHandler.sendErrorMessage(ctx, cause);

            if (cause instanceof UaException) {
                logger.debug("[remote={}] UaException caught; sent {}",
                    ctx.channel().remoteAddress(), errorMessage, cause);
            } else {
                logger.error("[remote={}] Exception caught; sent {}",
                    ctx.channel().remoteAddress(), errorMessage, cause);
            }
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        if (secureChannelTimeout != null) {
            secureChannelTimeout.cancel();
            secureChannelTimeout = null;
        }

        super.channelInactive(ctx);
    }

}
