
import com.jivesoftware.base.*;
import com.jivesoftware.base.database.ConnectionManager;
import com.jivesoftware.base.database.UserIterator;
import com.jivesoftware.util.LongList;

import javax.naming.Context;
import javax.naming.InitialContext;
import javax.sql.DataSource;
import java.sql.*;
import java.util.Iterator;

/**
 * Sample UserManager implementation, which creates and manages CustomUser
 * objects.
 */
public class CustomUserManager extends UserManagerAdapter {

    protected DataSource dataSource = null;

    // Database queries for the sample database schema this class uses. See
    // the CustomUser class for more information.
    private static final String USER_COUNT = "SELECT count(*) FROM person";
    private static final String ALL_USERS = "SELECT id FROM person";

    public CustomUserManager() {
        try {
            // Lookup the database source at a hard-coded name. You'd probably want
            // to make this configurable for a real implementation.
            Context context = new InitialContext();
            dataSource = (DataSource) context.lookup("java:comp/env/jdbc/database");
        } catch (Exception e) {
            Log.error(e);
        }
    }

    public int getUserCount() {
        // If your external user store doesn't support this operation, return 0.

        int count = 0;
        Connection con = null;
        PreparedStatement pstmt = null;
        try {
            con = ConnectionManager.getConnection();
            pstmt = con.prepareStatement(USER_COUNT);
            ResultSet rs = pstmt.executeQuery();
            if (rs.next()) {
                count = rs.getInt(1);
            }
        } catch (SQLException sqle) {
            sqle.printStackTrace();
        } finally {
            try {
                if (pstmt != null) {
                    pstmt.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
            try {
                if (con != null) {
                    con.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
        }
        return count;
    }

    public Iterator users() {
        // If your external user store doesn't support this operation, return
        // Collections.EMPTY_LIST.iterator();

        LongList users = new LongList(500);
        Connection con = null;
        PreparedStatement pstmt = null;

        try {
            con = ConnectionManager.getConnection();
            pstmt = con.prepareStatement(ALL_USERS);
            ResultSet rs = pstmt.executeQuery();
            // Set the fetch size. This will prevent some JDBC drivers from trying
            // to load the entire result set into memory.
            ConnectionManager.setFetchSize(rs, 500);
            while (rs.next()) {
                users.add(rs.getLong(1));
            }
        } catch (SQLException e) {
            Log.error(e);
        } finally {
            try {
                if (pstmt != null) {
                    pstmt.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
            try {
                if (con != null) {
                    con.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
        }
        return new UserIterator(users.toArray());
    }

    public Iterator users(int startIndex, int numResults) {
        // If your external user store doesn't support this operation, return
        // Collections.EMPTY_LIST.iterator();

        LongList users = new LongList();
        Connection con = null;
        PreparedStatement pstmt = null;

        try {
            con = ConnectionManager.getConnection();
            pstmt = con.prepareStatement(ALL_USERS);
            ResultSet rs = pstmt.executeQuery();
            ConnectionManager.setFetchSize(rs, startIndex + numResults);
            // Move to start of index
            for (int i = 0; i < startIndex; i++) {
                rs.next();
            }
            // Now read in desired number of results (or stop if we run out of results).
            for (int i = 0; i < numResults; i++) {
                if (rs.next()) {
                    users.add(rs.getLong(1));
                } else {
                    break;
                }
            }
        } catch (SQLException e) {
            Log.error(e);
        } finally {
            try {
                if (pstmt != null) {
                    pstmt.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
            try {
                if (con != null) {
                    con.close();
                }
            } catch (Exception e) {
                Log.error(e);
            }
        }
        return new UserIterator(users.toArray());
    }

    protected User lookupUser(long userID) throws UserNotFoundException {
        return new CustomUser(userID, dataSource);
    }

    protected User lookupUser(String username) throws UserNotFoundException {
        return new CustomUser(username, dataSource);
    }

    public boolean isGetUserCountSupported() {   // override to indicate that this is implemented.
        return true;
    }

    public boolean isUsersSupported() {   // override to indicate that this is implemented.
        return true;
    }
}
