import java.io.IOException;
import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

import javax.annotation.PostConstruct;
import javax.sql.DataSource;

import org.apache.commons.dbutils.ResultSetHandler;
import org.springframework.beans.factory.BeanInitializationException;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.util.Assert;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Table;
import com.ibm.as400.access.AS400;
import com.ibm.as400.access.AS400JDBCDataSource;
import com.ibm.as400.access.AS400SecurityException;
import com.ibm.as400.access.ErrorCompletingRequestException;
import com.ibm.as400.access.ObjectDoesNotExistException;
import com.ibm.as400.access.User;

/**
 *
 * @author bsanders
 *
 */
public class AS400UserEnvironmentDataSource extends AbstractDataSource {

    private AS400 system;
    private User anonymousUser;

    private String
        userAttribute = "user",
        environmentAttribute = "environment";

    private boolean dataSourcePerUser;

    private Table<String, User, DataSource> targetDataSources;

    public AS400UserEnvironmentDataSource() {}

    /**
     * Called after dependencies have been injected to perform initialization.
     * If an object of this class is instantiated outside of an IOC container, then this method must be called manually.
     */
    @PostConstruct
    public void afterPropertiesSet() {
        try {
            logger.info("Initializing datasource...");

            Assert.notNull(system, "system property must not be null");
            anonymousUser = new User(system, system.getUserId());

            Connection connection = new AS400JDBCDataSource(system).getConnection();
            PreparedStatement stmt = connection.prepareStatement("SELECT LLLLID FROM KAUTL400DT.KALLM");
            targetDataSources = new UserEnvironmentDataSourceTableHandler("LLLLID").handle(stmt.executeQuery());
        }
        catch (SQLException | AS400SecurityException | ErrorCompletingRequestException | InterruptedException | IOException | ObjectDoesNotExistException e) {
            throw new BeanInitializationException("Unable to initialize datasource", e);
        }
    }

    /**
     * Sets the {@link AS400} session to use when initializing the datasource
     * @param system An {@link AS400} session with the proper authorization to access the list of available environments
     */
    public void setSystem(AS400 system) {
        this.system = system;
    }

    /**
     * Used by {@link #getTargetDataSource()} to determine the proper {@link DataSource} to return.
     * Default value is "user".
     *
     * @return The session attribute in which this datasource can find the {@link User}
     */
    public String getUserAttribute() {
        return userAttribute;
    }

    /**
     *
     * @param userAttribute The session attribute in which {@link #getTargetDataSource()} can expect to find the associated {@link #getUser() user}
     */
    public void setUserAttribute(String userAttribute) {
        this.userAttribute = userAttribute;
    }

    /**
     * Used by {@link #getTargetDataSource()} to determine the proper {@link DataSource} to return
     *
     * @return The session attribute in which this datasource can find the {@link #getEnvironment() environment}
     */
    public String getEnvironmentAttribute() {
        return environmentAttribute;
    }

    /**
     *
     * @param environmentAttribute
     */
    public void setEnvironmentAttribute(String environmentAttribute) {
        this.environmentAttribute = environmentAttribute;
    }

    /**
     *
     * @return
     */
    public boolean isDataSourcePerUser() {
        return dataSourcePerUser;
    }

    /**
     *
     * @param dataSourcePerUser
     */
    public void setDataSourcePerUser(boolean dataSourcePerUser) {
        this.dataSourcePerUser = dataSourcePerUser;
    }

    /**
     *
     * @return A {@link DataSource}
     * @throws IllegalArgumentException
     */
    private DataSource getTargetDataSource() {
        String environment = getEnvironment();
        User user = getUser();

        if (targetDataSources.get(environment, user) == null) {
            AS400JDBCDataSource datasource = new AS400JDBCDataSource(user.getSystem());
            datasource.setNaming("system");

            targetDataSources.put(environment, user, datasource);
        }

        return targetDataSources.get(environment, user);
    }

    /**
     * Convenient method to retrieve values from the session map
     *
     * @param name The key of the value you want to retrieve
     * @return The value, if any
     */
    private Object getSessionAttribute(String name) {
        return RequestContextHolder.currentRequestAttributes().getAttribute(name, RequestAttributes.SCOPE_SESSION);
    }

    /**
     * Called by {@link #getTargetDataSource()} to determine which {@link User} will act as a key in the datasource table.
     *
     * @return The value of the session attribute indicated by {@link #getUserAttribute()} if {@link #isDataSourcePerUser()} is true.
     * Otherwise, returns the user associated with {@link #system}.
     */
    private User getUser() {
        User user = isDataSourcePerUser() ? (User) getSessionAttribute(getUserAttribute()) : anonymousUser;
        Assert.notNull(user, "No user found in session attribute \"" + getUserAttribute() + "\"");

        return user;
    }

    /**
     * returns the current {@link String environment} from the session using the specified {@link #getEnvironmentAttribute() environmentAttribute}
     * @return The name of the environment
     */
    private String getEnvironment() {
        String environment = (String) getSessionAttribute(getEnvironmentAttribute());
        Assert.notNull(environment, "Unable to determine environment because session attribute \"" + getEnvironmentAttribute() + "\" is null");
        Assert.isTrue(targetDataSources.containsRow(environment), "Environment doesn't exist: " + environment);

        return environment;
    }

    @Override
    public Connection getConnection() throws SQLException {
        Connection connection = getTargetDataSource().getConnection();

        String cmd = String.format("KALLM LLID(%s)", getEnvironment().trim());

        //Tried using placeholders, but kept getting "Descriptor index not valid"
        CallableStatement stmt = connection.prepareCall(String.format("CALL QCMDEXC ('%s', %d)", cmd, cmd.length()));
        stmt.execute();
        stmt.close();

        return connection;
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return getTargetDataSource().getConnection(username, password);
    }

    private static class UserEnvironmentDataSourceTableHandler implements ResultSetHandler<Table<String, User, DataSource>> {

        private String rowKeyColumnName;

        public UserEnvironmentDataSourceTableHandler(String rowKeyColumnName) {
            this.rowKeyColumnName = rowKeyColumnName;
        }

        @Override
        public Table<String, User, DataSource> handle(ResultSet rs) throws SQLException {
            HashBasedTable<String, User, DataSource> result = HashBasedTable.create();

            User defaultUser = new User();
            AS400JDBCDataSource defaultDataSource = new AS400JDBCDataSource();

            while (rs.next()) {
                result.put(rs.getString(rowKeyColumnName), defaultUser, defaultDataSource);
            }

            return result;
        }

    }
}