/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.activemq.transport.amqp.client;

import static org.apache.activemq.transport.amqp.AmqpSupport.CONNECTION_OPEN_FAILED;

import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.activemq.transport.InactivityIOException;
import org.apache.activemq.transport.amqp.client.sasl.SaslAuthenticator;
import org.apache.activemq.transport.amqp.client.transport.NettyTransportListener;
import org.apache.activemq.transport.amqp.client.util.AsyncResult;
import org.apache.activemq.transport.amqp.client.util.ClientFuture;
import org.apache.activemq.transport.amqp.client.util.IdGenerator;
import org.apache.activemq.transport.amqp.client.util.NoOpAsyncResult;
import org.apache.activemq.transport.amqp.client.util.UnmodifiableProxy;
import org.apache.qpid.proton.amqp.Symbol;
import org.apache.qpid.proton.engine.Collector;
import org.apache.qpid.proton.engine.Connection;
import org.apache.qpid.proton.engine.EndpointState;
import org.apache.qpid.proton.engine.Event;
import org.apache.qpid.proton.engine.Event.Type;
import org.apache.qpid.proton.engine.Sasl;
import org.apache.qpid.proton.engine.Transport;
import org.apache.qpid.proton.engine.impl.CollectorImpl;
import org.apache.qpid.proton.engine.impl.TransportImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.util.ReferenceCountUtil;

public class AmqpConnection extends AmqpAbstractResource<Connection> implements NettyTransportListener {

    private static final Logger LOG = LoggerFactory.getLogger(AmqpConnection.class);

    private static final NoOpAsyncResult NOOP_REQUEST = new NoOpAsyncResult();

    private static final int DEFAULT_MAX_FRAME_SIZE = 1024 * 1024 * 1;
    // NOTE: Limit default channel max to signed short range to deal with
    //       brokers that don't currently handle the unsigned range well.
    private static final int DEFAULT_CHANNEL_MAX = 32767;
    private static final IdGenerator CONNECTION_ID_GENERATOR = new IdGenerator();

    public static final long DEFAULT_CONNECT_TIMEOUT = 515000;
    public static final long DEFAULT_CLOSE_TIMEOUT = 30000;
    public static final long DEFAULT_DRAIN_TIMEOUT = 60000;

    private ScheduledThreadPoolExecutor serializer;
    private final AtomicBoolean closed = new AtomicBoolean();
    private final AtomicBoolean connected = new AtomicBoolean();
    private final AtomicLong sessionIdGenerator = new AtomicLong();
    private final AtomicLong txIdGenerator = new AtomicLong();
    private final Collector protonCollector = new CollectorImpl();
    private final org.apache.activemq.transport.amqp.client.transport.NettyTransport transport;
    private final Transport protonTransport = Transport.Factory.create();

    private final String username;
    private final String password;
    private final URI remoteURI;
    private final String connectionId;
    private List<Symbol> offeredCapabilities = Collections.emptyList();
    private Map<Symbol, Object> offeredProperties = Collections.emptyMap();

    private volatile AmqpFrameValidator sentFrameInspector;
    private volatile AmqpFrameValidator receivedFrameInspector;
    private AmqpConnectionListener listener;
    private SaslAuthenticator authenticator;
    private String mechanismRestriction;
    private String authzid;

    private int idleTimeout = 0;
    private boolean idleProcessingDisabled;
    private String containerId;
    private boolean authenticated;
    private int channelMax = DEFAULT_CHANNEL_MAX;
    private long connectTimeout = DEFAULT_CONNECT_TIMEOUT;
    private long closeTimeout = DEFAULT_CLOSE_TIMEOUT;
    private long drainTimeout = DEFAULT_DRAIN_TIMEOUT;
    private boolean trace;

    public AmqpConnection(org.apache.activemq.transport.amqp.client.transport.NettyTransport transport, String username, String password) {
        setEndpoint(Connection.Factory.create());
        getEndpoint().collect(protonCollector);

        this.transport = transport;
        this.username = username;
        this.password = password;
        this.connectionId = CONNECTION_ID_GENERATOR.generateId();
        this.remoteURI = transport.getRemoteLocation();

        this.serializer = new ScheduledThreadPoolExecutor(1, new ThreadFactory() {

            @Override
            public Thread newThread(Runnable runner) {
                Thread serial = new Thread(runner);
                serial.setDaemon(true);
                serial.setName(toString());
                return serial;
            }
        });

        // Ensure timely shutdown
        this.serializer.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
        this.serializer.setContinueExistingPeriodicTasksAfterShutdownPolicy(false);

        this.transport.setTransportListener(this);
        this.transport.setMaxFrameSize(getMaxFrameSize());
    }

    public void connect() throws Exception {
        if (connected.compareAndSet(false, true)) {
            transport.connect();

            final ClientFuture future = new ClientFuture();
            serializer.execute(new Runnable() {
                @Override
                public void run() {
                    getEndpoint().setContainer(safeGetContainerId());
                    getEndpoint().setHostname(remoteURI.getHost());
                    if (!getOfferedCapabilities().isEmpty()) {
                        getEndpoint().setOfferedCapabilities(getOfferedCapabilities().toArray(new Symbol[0]));
                    }
                    if (!getOfferedProperties().isEmpty()) {
                        getEndpoint().setProperties(getOfferedProperties());
                    }

                    if (getIdleTimeout() > 0) {
                        protonTransport.setIdleTimeout(getIdleTimeout());
                    }
                    protonTransport.setMaxFrameSize(getMaxFrameSize());
                    protonTransport.setChannelMax(getChannelMax());
                    protonTransport.bind(getEndpoint());
                    Sasl sasl = protonTransport.sasl();
                    if (sasl != null) {
                        sasl.client();
                    }
                    authenticator = new SaslAuthenticator(sasl, username, password, authzid, mechanismRestriction);
                    ((TransportImpl) protonTransport).setProtocolTracer(new AmqpProtocolTracer(AmqpConnection.this));
                    open(future);

                    pumpToProtonTransport(future);
                }
            });

            try {
                if (connectTimeout <= 0) {
                    future.sync();
                } else {
                    future.sync(connectTimeout, TimeUnit.MILLISECONDS);
                    if (getEndpoint().getRemoteState() != EndpointState.ACTIVE) {
                        throw new IOException("Failed to connect after configured timeout.");
                    }
                }
            } catch (Throwable error) {
                try {
                    close();
                } catch (Throwable ignore) {}

                throw error;
            }
        }
    }

    public boolean isConnected() {
        return transport.isConnected() && connected.get();
    }

    public void close() {
        if (closed.compareAndSet(false, true)) {
            final ClientFuture request = new ClientFuture();
            serializer.execute(new Runnable() {

                @Override
                public void run() {
                    try {

                        // If we are not connected then there is nothing we can do now
                        // just signal success.
                        if (!transport.isConnected()) {
                            request.onSuccess();
                        }

                        if (getEndpoint() != null) {
                            close(request);
                        } else {
                            request.onSuccess();
                        }

                        pumpToProtonTransport(request);
                    } catch (Exception e) {
                        LOG.debug("Caught exception while closing proton connection");
                    }
                }
            });

            try {
                if (closeTimeout <= 0) {
                    request.sync();
                } else {
                    request.sync(closeTimeout, TimeUnit.MILLISECONDS);
                }
            } catch (IOException e) {
                LOG.warn("Error caught while closing Provider: {}", e.getMessage());
            } finally {
                if (transport != null) {
                    try {
                        transport.close();
                    } catch (Exception e) {
                        LOG.debug("Cuaght exception while closing down Transport: {}", e.getMessage());
                    }
                }

                serializer.shutdownNow();
                try {
                    if (!serializer.awaitTermination(10, TimeUnit.SECONDS)) {
                        LOG.warn("Serializer didn't shutdown cleanly");
                    }
                } catch (InterruptedException e) {
                }
            }
        }
    }

    /**
     * Creates a new Session instance used to create AMQP resources like
     * senders and receivers.
     *
     * @return a new AmqpSession that can be used to create links.
     *
     * @throws Exception if an error occurs during creation.
     */
    public AmqpSession createSession() throws Exception {
        checkClosed();

        final AmqpSession session = new AmqpSession(AmqpConnection.this, getNextSessionId());
        final ClientFuture request = new ClientFuture();

        serializer.execute(new Runnable() {

            @Override
            public void run() {
                checkClosed();
                session.setEndpoint(getEndpoint().session());
                session.setStateInspector(getStateInspector());
                session.open(request);
                pumpToProtonTransport(request);
            }
        });

        request.sync();

        return session;
    }

    //----- Access to low level IO for specific test cases -------------------//

    public void sendRawBytes(final byte[] rawData) throws Exception {
        checkClosed();

        final ClientFuture request = new ClientFuture();

        serializer.execute(new Runnable() {

            @Override
            public void run() {
                checkClosed();
                try {
                    transport.send(Unpooled.wrappedBuffer(rawData));
                } catch (IOException e) {
                    fireClientException(e);
                } finally {
                    request.onSuccess();
                }
            }
        });

        request.sync();
    }

    //----- Configuration accessors ------------------------------------------//

    /**
     * @return the user name that was used to authenticate this connection.
     */
    public String getUsername() {
        return username;
    }

    /**
     * @return the password that was used to authenticate this connection.
     */
    public String getPassword() {
        return password;
    }

    public void setAuthzid(String authzid) {
        this.authzid = authzid;
    }

    public String getAuthzid() {
        return authzid;
    }

    /**
     * @return the URI of the remote peer this connection attached to.
     */
    public URI getRemoteURI() {
        return remoteURI;
    }

    /**
     * @return the container ID that will be set as the container Id.
     */
    public String getContainerId() {
        return this.containerId;
    }

    /**
     * Sets the container Id that will be configured on the connection prior to
     * connecting to the remote peer.  Calling this after connect has no effect.
     *
     * @param containerId
     * 		  the container Id to use on the connection.
     */
    public void setContainerId(String containerId) {
        this.containerId = containerId;
    }

    /**
     * @return the currently set Max Frame Size value.
     */
    public int getMaxFrameSize() {
        return DEFAULT_MAX_FRAME_SIZE;
    }

    public int getChannelMax() {
        return channelMax;
    }

    public void setChannelMax(int channelMax) {
        this.channelMax = channelMax;
    }

    public long getConnectTimeout() {
        return connectTimeout;
    }

    public void setConnectTimeout(long connectTimeout) {
        this.connectTimeout = connectTimeout;
    }

    public long getCloseTimeout() {
        return closeTimeout;
    }

    public void setCloseTimeout(long closeTimeout) {
        this.closeTimeout = closeTimeout;
    }

    public long getDrainTimeout() {
        return drainTimeout;
    }

    public void setDrainTimeout(long drainTimeout) {
        this.drainTimeout = drainTimeout;
    }

    public List<Symbol> getOfferedCapabilities() {
        return offeredCapabilities;
    }

    public void setOfferedCapabilities(List<Symbol> offeredCapabilities) {
        if (offeredCapabilities != null) {
            offeredCapabilities = Collections.emptyList();
        }

        this.offeredCapabilities = offeredCapabilities;
    }

    public Map<Symbol, Object> getOfferedProperties() {
        return offeredProperties;
    }

    public void setOfferedProperties(Map<Symbol, Object> offeredProperties) {
        if (offeredProperties != null) {
            offeredProperties = Collections.emptyMap();
        }

        this.offeredProperties = offeredProperties;
    }

    public Connection getConnection() {
        return UnmodifiableProxy.connectionProxy(getEndpoint());
    }

    public AmqpConnectionListener getListener() {
        return listener;
    }

    public void setListener(AmqpConnectionListener listener) {
        this.listener = listener;
    }

    public int getIdleTimeout() {
        return idleTimeout;
    }

    public void setIdleTimeout(int idleTimeout) {
        this.idleTimeout = idleTimeout;
    }

    public void setIdleProcessingDisabled(boolean value) {
        this.idleProcessingDisabled = value;
    }

    public boolean isIdleProcessingDisabled() {
        return idleProcessingDisabled;
    }

    /**
     * Sets a restriction on the SASL mechanism to use (if offered by the server).
     *
     * @param mechanismRestriction the mechanism to use
     */
    public void setMechanismRestriction(String mechanismRestriction) {
        this.mechanismRestriction = mechanismRestriction;
    }

    public String getMechanismRestriction() {
        return mechanismRestriction;
    }

    public boolean isTraceFrames() {
        return trace;
    }

    public void setTraceFrames(boolean trace) {
        this.trace = trace;
    }

    public AmqpFrameValidator getSentFrameInspector() {
        return sentFrameInspector;
    }

    public void setSentFrameInspector(AmqpFrameValidator amqpFrameInspector) {
        this.sentFrameInspector = amqpFrameInspector;
    }

    public AmqpFrameValidator getReceivedFrameInspector() {
        return receivedFrameInspector;
    }

    public void setReceivedFrameInspector(AmqpFrameValidator amqpFrameInspector) {
        this.receivedFrameInspector = amqpFrameInspector;
    }

    //----- Internal getters used from the child AmqpResource classes --------//

    ScheduledExecutorService getScheduler() {
        return this.serializer;
    }

    Connection getProtonConnection() {
        return getEndpoint();
    }

    String getConnectionId() {
        return this.connectionId;
    }

    AmqpTransactionId getNextTransactionId() {
        return new AmqpTransactionId(connectionId + ":" + txIdGenerator.incrementAndGet());
    }

    void pumpToProtonTransport() {
        pumpToProtonTransport(NOOP_REQUEST);
    }

    void pumpToProtonTransport(AsyncResult request) {
        try {
            boolean done = false;
            while (!done) {
                ByteBuffer toWrite = protonTransport.getOutputBuffer();
                if (toWrite != null && toWrite.hasRemaining()) {
                    ByteBuf outbound = transport.allocateSendBuffer(toWrite.remaining());
                    outbound.writeBytes(toWrite);
                    transport.send(outbound);
                    protonTransport.outputConsumed();
                } else {
                    done = true;
                }
            }
        } catch (IOException e) {
            fireClientException(e);
            request.onFailure(e);
        }
    }

    //----- Transport listener event hooks -----------------------------------//

    @Override
    public void onData(final ByteBuf incoming) {

        // We need to retain until the serializer gets around to processing it.
        ReferenceCountUtil.retain(incoming);

        serializer.execute(new Runnable() {

            @Override
            public void run() {
                ByteBuffer source = incoming.nioBuffer();
                LOG.trace("Client Received from Broker {} bytes:", source.remaining());

                if (protonTransport.isClosed()) {
                    LOG.debug("Ignoring incoming data because transport is closed");
                    return;
                }

                do {
                    ByteBuffer buffer = protonTransport.getInputBuffer();
                    int limit = Math.min(buffer.remaining(), source.remaining());
                    ByteBuffer duplicate = source.duplicate();
                    duplicate.limit(source.position() + limit);
                    buffer.put(duplicate);
                    protonTransport.processInput();
                    source.position(source.position() + limit);
                } while (source.hasRemaining());

                ReferenceCountUtil.release(incoming);

                // Process the state changes from the latest data and then answer back
                // any pending updates to the Broker.
                processUpdates();
                pumpToProtonTransport();
            }
        });
    }

    @Override
    public void onTransportClosed() {
        LOG.debug("The transport has unexpectedly closed");
        failed(getOpenAbortException());
    }

    @Override
    public void onTransportError(Throwable cause) {
        fireClientException(cause);
    }

    //----- Internal implementation ------------------------------------------//

    @Override
    protected void doOpenCompletion() {
        // If the remote indicates that a close is pending, don't open.
        if (getEndpoint().getRemoteProperties() == null ||
            !getEndpoint().getRemoteProperties().containsKey(CONNECTION_OPEN_FAILED)) {

            if (!isIdleProcessingDisabled()) {
                // Using nano time since it is not related to the wall clock, which may change
                long initialNow = TimeUnit.NANOSECONDS.toMillis(System.nanoTime());
                long initialKeepAliveDeadline = protonTransport.tick(initialNow);
                if (initialKeepAliveDeadline != 0) {

                    getScheduler().schedule(new Runnable() {

                        @Override
                        public void run() {
                            try {
                                if (getEndpoint().getLocalState() != EndpointState.CLOSED) {
                                    LOG.debug("Client performing next idle check");
                                    // Using nano time since it is not related to the wall clock, which may change
                                    long now = TimeUnit.NANOSECONDS.toMillis(System.nanoTime());
                                    long deadline = protonTransport.tick(now);

                                    pumpToProtonTransport();
                                    if (protonTransport.isClosed()) {
                                        LOG.debug("Transport closed after inactivity check.");
                                        throw new InactivityIOException("Channel was inactive for too long");
                                    } else {
                                        if(deadline != 0) {
                                            getScheduler().schedule(this, deadline - now, TimeUnit.MILLISECONDS);
                                        }
                                    }
                                }
                            } catch (Exception e) {
                                try {
                                    transport.close();
                                } catch (IOException e1) {
                                }
                                fireClientException(e);
                            }
                        }
                    }, initialKeepAliveDeadline - initialNow, TimeUnit.MILLISECONDS);
                }
            }
            super.doOpenCompletion();
        }
    }

    @Override
    protected void doOpenInspection() {
        try {
            getStateInspector().inspectOpenedResource(getConnection());
        } catch (Throwable error) {
            getStateInspector().markAsInvalid(error.getMessage());
        }
    }

    @Override
    protected void doClosedInspection() {
        try {
            getStateInspector().inspectClosedResource(getConnection());
        } catch (Throwable error) {
            getStateInspector().markAsInvalid(error.getMessage());
        }
    }

    protected void fireClientException(Throwable ex) {
        AmqpConnectionListener listener = this.listener;
        if (listener != null) {
            listener.onException(ex);
        }
    }

    protected void checkClosed() throws IllegalStateException {
        if (closed.get()) {
            throw new IllegalStateException("The Connection is already closed");
        }
    }

    private void processUpdates() {
        try {
            Event protonEvent = null;
            while ((protonEvent = protonCollector.peek()) != null) {
                if (!protonEvent.getType().equals(Type.TRANSPORT)) {
                    LOG.trace("Client: New Proton Event: {}", protonEvent.getType());
                }

                AmqpEventSink amqpEventSink = null;
                switch (protonEvent.getType()) {
                    case CONNECTION_REMOTE_CLOSE:
                        amqpEventSink = (AmqpEventSink) protonEvent.getConnection().getContext();
                        amqpEventSink.processRemoteClose(this);
                        break;
                    case CONNECTION_REMOTE_OPEN:
                        amqpEventSink = (AmqpEventSink) protonEvent.getConnection().getContext();
                        amqpEventSink.processRemoteOpen(this);
                        break;
                    case SESSION_REMOTE_CLOSE:
                        amqpEventSink = (AmqpEventSink) protonEvent.getSession().getContext();
                        amqpEventSink.processRemoteClose(this);
                        break;
                    case SESSION_REMOTE_OPEN:
                        amqpEventSink = (AmqpEventSink) protonEvent.getSession().getContext();
                        amqpEventSink.processRemoteOpen(this);
                        break;
                    case LINK_REMOTE_CLOSE:
                        amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                        amqpEventSink.processRemoteClose(this);
                        break;
                    case LINK_REMOTE_DETACH:
                        amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                        amqpEventSink.processRemoteDetach(this);
                        break;
                    case LINK_REMOTE_OPEN:
                        amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                        amqpEventSink.processRemoteOpen(this);
                        break;
                    case LINK_FLOW:
                        amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                        amqpEventSink.processFlowUpdates(this);
                        break;
                    case DELIVERY:
                        amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                        amqpEventSink.processDeliveryUpdates(this);
                        break;
                    default:
                        break;
                }

                protonCollector.pop();
            }

            // We have to do this to pump SASL bytes in as SASL is not event driven yet.
            if (!authenticated) {
                processSaslAuthentication();
            }
        } catch (Exception ex) {
            LOG.warn("Caught Exception during update processing: {}", ex.getMessage(), ex);
            fireClientException(ex);
        }
    }

    private void processSaslAuthentication() {
        if (authenticated || authenticator == null) {
            return;
        }

        try {
            if (authenticator.authenticate()) {
                authenticator = null;
                authenticated = true;
            }
        } catch (SecurityException ex) {
            failed(ex);
        }
    }

    private String getNextSessionId() {
        return connectionId + ":" + sessionIdGenerator.incrementAndGet();
    }

    private String safeGetContainerId() {
        String containerId = getContainerId();
        if (containerId == null || containerId.isEmpty()) {
            containerId = UUID.randomUUID().toString();
        }

        return containerId;
    }

    @Override
    public String toString() {
        return "AmqpConnection { " + connectionId + " }";
    }
}
