/*
    * 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
    *
    *      https://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.commons.dbcp2.datasources;
  
  import java.sql.Connection;
  import java.sql.SQLException;
  import java.time.Duration;
  
  import javax.sql.ConnectionEvent;
  import javax.sql.ConnectionEventListener;
  import javax.sql.ConnectionPoolDataSource;
  import javax.sql.PooledConnection;
  
  import org.apache.commons.dbcp2.PoolableConnection;
  import org.apache.commons.pool2.KeyedObjectPool;
  import org.apache.commons.pool2.KeyedPooledObjectFactory;
  import org.apache.commons.pool2.PooledObject;
  import org.apache.commons.pool2.impl.DefaultPooledObject;
  
  /**
   * A {@link KeyedPooledObjectFactory} that creates {@link PoolableConnection}s.
   *
   * @since 2.0
   */
  final class KeyedCPDSConnectionFactory extends AbstractConnectionFactory
          implements KeyedPooledObjectFactory<UserPassKey, PooledConnectionAndInfo>, ConnectionEventListener, PooledConnectionManager {
  
      private static final String NO_KEY_MESSAGE = "close() was called on a Connection, but I have no record of the underlying PooledConnection.";
      private KeyedObjectPool<UserPassKey, PooledConnectionAndInfo> pool;
  
      /**
       * Creates a new {@code KeyedCPDSConnectionFactory}.
       *
       * @param cpds
       *            the ConnectionPoolDataSource from which to obtain PooledConnections
       * @param validationQuery
       *            a query to use to {@link #validateObject validate} {@link Connection}s. Should return at least one
       *            row. May be {@code null} in which case3 {@link Connection#isValid(int)} will be used to validate
       *            connections.
       * @param validationQueryTimeoutDuration
       *            The Duration to allow for the validation query to complete
       * @param rollbackAfterValidation
       *            whether a rollback should be issued after {@link #validateObject validating} {@link Connection}s.
       * @since 2.10.0
       */
      public KeyedCPDSConnectionFactory(final ConnectionPoolDataSource cpds, final String validationQuery, final Duration validationQueryTimeoutDuration,
              final boolean rollbackAfterValidation) {
          super(cpds, validationQuery, validationQueryTimeoutDuration, rollbackAfterValidation);
      }
  
      @Override
      public void activateObject(final UserPassKey ignored, final PooledObject<PooledConnectionAndInfo> pooledObject) throws SQLException {
          validateLifetime(pooledObject);
      }
  
      /**
       * This implementation does not fully close the KeyedObjectPool, as this would affect all users. Instead, it clears
       * the pool associated with the given user. This method is not currently used.
       */
      @Override
      public void closePool(final String userName) throws SQLException {
          try {
              pool.clear(new UserPassKey(userName));
          } catch (final Exception ex) {
              throw new SQLException("Error closing connection pool", ex);
          }
      }
  
      /**
       * This will be called if the Connection returned by the getConnection method came from a PooledConnection, and the
       * user calls the close() method of this connection object. What we need to do here is to release this
       * PooledConnection from our pool...
       */
      @Override
      public void connectionClosed(final ConnectionEvent event) {
          final PooledConnection pc = (PooledConnection) event.getSource();
          // if this event occurred because we were validating, or if this
          // connection has been marked for removal, ignore it
          // otherwise return the connection to the pool.
          if (!validatingSet.contains(pc)) {
              final PooledConnectionAndInfo pci = pcMap.get(pc);
              if (pci == null) {
                  throw new IllegalStateException(NO_KEY_MESSAGE);
              }
              try {
                 pool.returnObject(pci.getUserPassKey(), pci);
             } catch (final Exception e) {
                 System.err.println("CLOSING DOWN CONNECTION AS IT COULD NOT BE RETURNED TO THE POOL");
                 pc.removeConnectionEventListener(this);
                 try {
                     pool.invalidateObject(pci.getUserPassKey(), pci);
                 } catch (final Exception e3) {
                     System.err.println("EXCEPTION WHILE DESTROYING OBJECT " + pci);
                     e3.printStackTrace();
                 }
             }
         }
     }
 
     /**
      * If a fatal error occurs, close the underlying physical connection so as not to be returned in the future
      */
     @Override
     public void connectionErrorOccurred(final ConnectionEvent event) {
         final PooledConnection pc = (PooledConnection) event.getSource();
         if (null != event.getSQLException()) {
             System.err.println("CLOSING DOWN CONNECTION DUE TO INTERNAL ERROR (" + event.getSQLException() + ")");
         }
         pc.removeConnectionEventListener(this);
 
         final PooledConnectionAndInfo info = pcMap.get(pc);
         if (info == null) {
             throw new IllegalStateException(NO_KEY_MESSAGE);
         }
         try {
             pool.invalidateObject(info.getUserPassKey(), info);
         } catch (final Exception e) {
             System.err.println("EXCEPTION WHILE DESTROYING OBJECT " + info);
             e.printStackTrace();
         }
     }
 
     /**
      * Closes the PooledConnection and stops listening for events from it.
      */
     @Override
     public void destroyObject(final UserPassKey ignored, final PooledObject<PooledConnectionAndInfo> pooledObject) throws SQLException {
         final PooledConnection pooledConnection = pooledObject.getObject().getPooledConnection();
         pooledConnection.removeConnectionEventListener(this);
         pcMap.remove(pooledConnection);
         pooledConnection.close();
     }
 
     /**
      * Returns the keyed object pool used to pool connections created by this factory.
      *
      * @return KeyedObjectPool managing pooled connections
      */
     public KeyedObjectPool<UserPassKey, PooledConnectionAndInfo> getPool() {
         return pool;
     }
 
     /**
      * Invalidates the PooledConnection in the pool. The KeyedCPDSConnectionFactory closes the connection and pool
      * counters are updated appropriately. Also clears any idle instances associated with the user name that was used to
      * create the PooledConnection. Connections associated with this user are not affected, and they will not be
      * automatically closed on return to the pool.
      */
     @Override
     public void invalidate(final PooledConnection pc) throws SQLException {
         final PooledConnectionAndInfo info = pcMap.get(pc);
         if (info == null) {
             throw new IllegalStateException(NO_KEY_MESSAGE);
         }
         final UserPassKey key = info.getUserPassKey();
         try {
             pool.invalidateObject(key, info); // Destroy and update pool counters
             pool.clear(key); // Remove any idle instances with this key
         } catch (final Exception ex) {
             throw new SQLException("Error invalidating connection", ex);
         }
     }
 
     /**
      * Creates a new {@code PooledConnectionAndInfo} from the given {@code UserPassKey}.
      *
      * @param userPassKey
      *            {@code UserPassKey} containing user credentials
      * @throws SQLException
      *             if the connection could not be created.
      * @see org.apache.commons.pool2.KeyedPooledObjectFactory#makeObject(Object)
      */
     @Override
     public synchronized PooledObject<PooledConnectionAndInfo> makeObject(final UserPassKey userPassKey) throws SQLException {
         PooledConnection pooledConnection = null;
         final String userName = userPassKey.getUserName();
         final String password = userPassKey.getPassword();
         if (userName == null) {
             pooledConnection = cpds.getPooledConnection();
         } else {
             pooledConnection = cpds.getPooledConnection(userName, password);
         }
         if (pooledConnection == null) {
             throw new IllegalStateException("Connection pool data source returned null from getPooledConnection");
         }
         // should we add this object as a listener or the pool.
         // consider the validateObject method in decision
         pooledConnection.addConnectionEventListener(this);
         final PooledConnectionAndInfo pci = new PooledConnectionAndInfo(pooledConnection, userPassKey);
         pcMap.put(pooledConnection, pci);
         return new DefaultPooledObject<>(pci);
     }
 
     @Override
     public void passivateObject(final UserPassKey ignored, final PooledObject<PooledConnectionAndInfo> pooledObject) throws SQLException {
         validateLifetime(pooledObject);
     }
 
     /**
      * Does nothing. This factory does not cache user credentials.
      */
     @Override
     public void setPassword(final char[] password) {
         // Does nothing. This factory does not cache user credentials.
     }
 
     /**
      * Does nothing. This factory does not cache user credentials.
      */
     @Override
     public void setPassword(final String password) {
         // Does nothing. This factory does not cache user credentials.
     }
 
     public void setPool(final KeyedObjectPool<UserPassKey, PooledConnectionAndInfo> pool) {
         this.pool = pool;
     }
 
     /**
      * Validates a pooled connection.
      * <p>
      * A query validation timeout greater than 0 and less than 1 second is converted to 1 second.
      * </p>
      *
      * @param ignored
      *            ignored
      * @param pooledObject
      *            wrapped {@code PooledConnectionAndInfo} containing the connection to validate
      * @return true if validation succeeds
      * @throws ArithmeticException if the query validation timeout does not fit as seconds in an int.
      */
     @Override
     public boolean validateObject(final UserPassKey ignored, final PooledObject<PooledConnectionAndInfo> pooledObject) {
         return super.validateObject(pooledObject);
     }
 }