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

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import jakarta.jms.JMSException;
import jakarta.jms.ServerSession;
import jakarta.jms.ServerSessionPool;
import jakarta.jms.Session;
import jakarta.resource.spi.UnavailableException;
import jakarta.resource.spi.endpoint.MessageEndpoint;

import org.apache.activemq.ActiveMQConnection;
import org.apache.activemq.ActiveMQSession;
import org.apache.activemq.command.MessageDispatch;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *  $Date$
 */
public class ServerSessionPoolImpl implements ServerSessionPool {

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

    private final ActiveMQEndpointWorker activeMQAsfEndpointWorker;
    private final int maxSessions;

    private final List<ServerSessionImpl> idleSessions = new ArrayList<ServerSessionImpl>();
    private final List<ServerSessionImpl> activeSessions = new ArrayList<ServerSessionImpl>();
    private final Lock sessionLock = new ReentrantLock();
    private final AtomicBoolean closing = new AtomicBoolean(false);

    public ServerSessionPoolImpl(ActiveMQEndpointWorker activeMQAsfEndpointWorker, int maxSessions) {
        this.activeMQAsfEndpointWorker = activeMQAsfEndpointWorker;
        this.maxSessions = maxSessions;
    }

    private ServerSessionImpl createServerSessionImpl() throws JMSException {
        MessageActivationSpec activationSpec = activeMQAsfEndpointWorker.endpointActivationKey.getActivationSpec();
        int acknowledge = (activeMQAsfEndpointWorker.transacted) ? Session.SESSION_TRANSACTED : activationSpec.getAcknowledgeModeForSession();
        final ActiveMQConnection connection = activeMQAsfEndpointWorker.getConnection();
        if (connection == null) {
            // redispatch of pending prefetched messages after disconnect can have a null connection
            return null;
        }
        final ActiveMQSession session = (ActiveMQSession)connection.createSession(activeMQAsfEndpointWorker.transacted, acknowledge);
        MessageEndpoint endpoint;
        try {
            int batchSize = 0;
            if (activationSpec.getEnableBatchBooleanValue()) {
                batchSize = activationSpec.getMaxMessagesPerBatchIntValue();
            }
            if (activationSpec.isUseRAManagedTransactionEnabled()) {
                // The RA will manage the transaction commit.
                endpoint = createEndpoint(null);
                return new ServerSessionImpl(this, session, activeMQAsfEndpointWorker.workManager, endpoint, true, batchSize);
            } else {
                // Give the container an object to manage to transaction with.
                endpoint = createEndpoint(new LocalAndXATransaction(session.getTransactionContext()));
                return new ServerSessionImpl(this, session, activeMQAsfEndpointWorker.workManager, endpoint, false, batchSize);
            }
        } catch (UnavailableException e) {
            // The container could be limiting us on the number of endpoints
            // that are being created.
            if (LOG.isDebugEnabled()) {
                LOG.debug("Could not create an endpoint.", e);
            }
            session.close();
            return null;
        }
    }

    private MessageEndpoint createEndpoint(LocalAndXATransaction txResourceProxy) throws UnavailableException {
        MessageEndpoint endpoint;
        endpoint = activeMQAsfEndpointWorker.endpointFactory.createEndpoint(txResourceProxy);
        MessageEndpointProxy endpointProxy = new MessageEndpointProxy(endpoint);
        return endpointProxy;
    }

    /**
     */
    @Override
    public ServerSession getServerSession() throws JMSException {
        if (LOG.isDebugEnabled()) {
            LOG.debug("ServerSession requested.");
        }
        if (closing.get()) {
            throw new JMSException("Session Pool Shutting Down.");
        }
        ServerSessionImpl ss = null;
        sessionLock.lock();
        try {
            ss = getExistingServerSession(false);
        } finally {
            sessionLock.unlock();
        }
        if (ss != null) {
            return ss;
        }
        ss = createServerSessionImpl();
        sessionLock.lock();
        try {
            // We may not be able to create a session due to the container
            // restricting us.
            if (ss == null) {
                if (activeSessions.isEmpty() && idleSessions.isEmpty()) {
                    throw new JMSException("Endpoint factory did not allow creation of any endpoints.");
                }

                ss = getExistingServerSession(true);
            } else {
                activeSessions.add(ss);
            }
        } finally {
            sessionLock.unlock();
        }
        if (LOG.isDebugEnabled()) {
            LOG.debug("Created a new session: " + ss);
        }
        return ss;

    }

    /**
     * Must be called with sessionLock held.
     * Returns an idle session if one exists or an active session if no more
     * sessions can be created.  Sessions can not be created if force is true
     * or activeSessions >= maxSessions.
     * @param force do not check activeSessions >= maxSessions, return an active connection anyway.
     * @return an already existing session.
     */
    private ServerSessionImpl getExistingServerSession(boolean force) {
        ServerSessionImpl ss = null;
        if (idleSessions.size() > 0) {
            ss = idleSessions.remove(idleSessions.size() - 1);
        }
        if (ss != null) {
            activeSessions.add(ss);
            if (LOG.isDebugEnabled()) {
                LOG.debug("Using idle session: " + ss);
            }
        } else if (force || activeSessions.size() >= maxSessions) {
            // If we are at the upper limit
            // then reuse the already created sessions..
            // This is going to queue up messages into a session for
            // processing.
            ss = getExistingActiveServerSession();
        }
        return ss;
    }

    /**
     * Must be called with sessionLock held.
     * Returns the first session from activeSessions, shifting it to last.
     * @return session
     */
    private ServerSessionImpl getExistingActiveServerSession() {
        ServerSessionImpl ss = null;
        // try to find a session that is not in use already (all message consumed and back to the pool or stale to be removed)
        for (ServerSessionImpl activeSession : activeSessions) {
            if (!activeSession.isRunning()) {
                activeSessions.remove(activeSession);
                activeSessions.add(activeSession);
                if (LOG.isDebugEnabled()) {
                    LOG.debug("Selected non-running active session: {}", activeSession);
                }
                return activeSession;
            }
        }
        // if none, we'll have to pull a running active session
        if (!activeSessions.isEmpty()) {
            if (activeSessions.size() > 1) {
                // round-robin
                ss = activeSessions.remove(0);
                activeSessions.add(ss);
            } else {
                ss = activeSessions.get(0);
            }
        }
        if (LOG.isDebugEnabled()) {
            LOG.debug("Reusing a running active session: {}", ss);
        }
        return ss;
    }

    public void returnToPool(ServerSessionImpl ss) {
        sessionLock.lock();
        activeSessions.remove(ss);
        try {
            // make sure we only return non-stale sessions to the pool
            if ( ss.isStale() ) {
                if ( LOG.isDebugEnabled() ) {
                    LOG.debug("Discarding stale ServerSession to be returned to pool: " + ss);
                }
                ss.close();
            } else {
                if (LOG.isDebugEnabled()) {
                    LOG.debug("ServerSession returned to pool: " + ss);
                }
            idleSessions.add(ss);
            }
        } finally {
            sessionLock.unlock();
        }
        synchronized (closing) {
            closing.notify();
        }
    }

    public void removeFromPool(ServerSessionImpl ss) {
        sessionLock.lock();
        try {
            activeSessions.remove(ss);
        } finally {
            sessionLock.unlock();
        }
        try {
            ActiveMQSession session = (ActiveMQSession)ss.getSession();
            List<MessageDispatch> l = session.getUnconsumedMessages();
            if (!isClosing() && !l.isEmpty()) {
                ActiveMQConnection connection = activeMQAsfEndpointWorker.getConnection();
                if (connection != null) {
                    for (Iterator<MessageDispatch> i = l.iterator(); i.hasNext();) {
                        MessageDispatch md = i.next();
                        if (connection.hasDispatcher(md.getConsumerId())) {
                            dispatchToSession(md);
                            LOG.trace("on remove of {} redispatch of {}", session, md);
                        } else {
                            LOG.trace("on remove not redispatching {}, dispatcher no longer present on {}", md, session.getConnection());
                        }
                    }
                } else {
                    LOG.trace("on remove of {} not redispatching while disconnected", session);
                }
            }
        } catch (Throwable t) {
            LOG.error("Error redispatching unconsumed messages from stale server session {}", ss, t);
        }
        ss.close();
        synchronized (closing) {
            closing.notify();
        }
    }

    /**
     * @param messageDispatch
     *            the message to dispatch
     * @throws JMSException
     */
    private void dispatchToSession(MessageDispatch messageDispatch)
            throws JMSException {

        ServerSession serverSession = getServerSession();
        Session s = serverSession.getSession();
        ActiveMQSession session = null;
        if (s instanceof ActiveMQSession) {
            session = (ActiveMQSession) s;
        } else {
            activeMQAsfEndpointWorker.getConnection()
                    .onAsyncException(new JMSException(
                            "Session pool provided an invalid session type: "
                                    + s.getClass()));
        }
        session.dispatch(messageDispatch);
        serverSession.start();
    }

    public void close() {
        closing.set(true);
        LOG.debug("{} close", this);
        int activeCount = closeSessions();
        // we may have to wait erroneously 250ms if an
        // active session is removed during our wait and we
        // are not notified
        while (activeCount > 0) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("Active Sessions = " + activeCount);
            }
            try {
                synchronized (closing) {
                    closing.wait(250);
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }
            activeCount = closeSessions();
        }
    }


    protected int closeSessions() {
        sessionLock.lock();
        try {
            List<ServerSessionImpl> alreadyClosedServerSessions = new ArrayList<>(activeSessions.size());
            for (ServerSessionImpl ss : activeSessions) {
                try {
                    ActiveMQSession session = (ActiveMQSession) ss.getSession();
                    if (!session.isClosed()) {
                        session.close();
                    } else {
                        LOG.debug("Session {} already closed", session);
                        alreadyClosedServerSessions.add(ss);

                    }
                } catch (JMSException ignored) {
                    if (LOG.isDebugEnabled()) {
                        LOG.debug("Failed to close active running server session {}, reason:{}", ss, ignored.toString(), ignored);
                    }
                }
            }
            for (ServerSessionImpl ss : alreadyClosedServerSessions) {
                removeFromPool(ss);
            }
            alreadyClosedServerSessions.clear();

            for (ServerSessionImpl ss : idleSessions) {
                ss.close();
            }
            idleSessions.clear();
            return activeSessions.size();
        } finally {
            sessionLock.unlock();
        }
    }

    /**
     * @return Returns the closing.
     */
    public boolean isClosing() {
        return closing.get();
    }

    /**
     * @param closing The closing to set.
     */
    public void setClosing(boolean closing) {
        this.closing.set(closing);
    }

}
