package ru.yandex.webmaster3.storage.util.sql;

import org.springframework.dao.DataAccessException;
import org.springframework.dao.IncorrectResultSizeDataAccessException;
import org.springframework.jdbc.core.ArgumentPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.ResultSetExtractor;
import org.springframework.jdbc.core.RowCallbackHandler;
import org.springframework.jdbc.core.RowMapper;
import ru.yandex.common.util.collections.Pair;
import ru.yandex.webmaster3.storage.WebmasterSQLException;
import ru.yandex.wmtools.common.util.ParameterizedMapRowMapper;
import ru.yandex.wmtools.common.util.SqlUtil;

import javax.sql.DataSource;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Date;
import java.util.List;
import java.util.NavigableMap;
import java.util.TreeMap;

/**
 * @author avhaliullin
 */
public class ReadJdbcTemplate implements IReadJdbcTemplate {
    protected final JdbcTemplate jdbcTemplate;

    protected final String dbIdentity;

    public ReadJdbcTemplate(DataSource dataSource, String dbIdentity) {
        jdbcTemplate = new JdbcTemplate(dataSource);
        this.dbIdentity = dbIdentity;
    }

    @Override
    public String getDbIdentity() {
        return dbIdentity;
    }

    @Override
    public int queryForInt(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForObject(sqlString, Integer.class, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public long queryForLong(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForObject(sqlString, Long.class, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public <T> T queryForObject(final String sqlString, final Class<T> aClass, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForObject(sqlString, aClass, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public <T> T queryForObject(final String sqlString, RowMapper<T> rowMapper, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForObject(sqlString, rowMapper, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public <T> List<T> query(final String sqlString, RowMapper<T> rowMapper, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.query(sqlString, rowMapper, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public java.util.Map<String, Object> queryForMap(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForMap(sqlString, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public List<java.util.Map<String, Object>> queryForList(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return jdbcTemplate.queryForList(sqlString, requestParams);
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public Date safeQueryForTimestamp(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        RowMapper<Date> dateRowMapper = new RowMapper<Date>() {
            @Override
            public Date mapRow(ResultSet resultSet, int i) throws SQLException {
                return SqlUtil.safeGetTimestamp(resultSet, 1);
            }
        };

        try {
            List<Date> results = query(sqlString, dateRowMapper, requestParams);
            if ((results == null) || (results.size() != 1)) {
                return null;
            }

            return results.get(0);
        } catch (WebmasterSQLException e) {
            if (e.getCause() instanceof IncorrectResultSizeDataAccessException) {
                return null;
            }
            throw e;
        }
    }

    @Override
    public Integer safeQueryForInt(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return queryForObject(sqlString, Integer.class, requestParams);
        } catch (WebmasterSQLException e) {
            if (e.getCause() instanceof IncorrectResultSizeDataAccessException) {
                return null;
            }
            throw e;
        }
    }

    @Override
    public Long safeQueryForLong(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return queryForObject(sqlString, Long.class, requestParams);
        } catch (WebmasterSQLException e) {
            if (e.getCause() instanceof IncorrectResultSizeDataAccessException) {
                return null;
            }
            throw e;
        }
    }

    @Override
    public String safeQueryForString(final String sqlString, final Object... requestParams) throws WebmasterSQLException {
        try {
            return queryForObject(sqlString, String.class, requestParams);
        } catch (WebmasterSQLException e) {
            if (e.getCause() instanceof IncorrectResultSizeDataAccessException) {
                return null;
            }
            throw e;
        }
    }

    @Override
    public <E> E safeQueryForObject(final String sqlString, RowMapper<E> rowMapper, final Object... requestParams) throws WebmasterSQLException {
        try {
            return queryForObject(sqlString, rowMapper, requestParams);
        } catch (WebmasterSQLException e) {
            if (e.getCause() instanceof IncorrectResultSizeDataAccessException) {
                return null;
            }
            throw e;
        }
    }

    @Override
    public <K, V> NavigableMap<K, V> queryForNavigableMap(final String sqlString, ParameterizedMapRowMapper<K, V> rowMapper, Object... requestParams) throws WebmasterSQLException {
        try {
            // Considers an Object array passed into a varargs parameter as collection of arguments rather than as single argument.
            if (requestParams.length == 1 && requestParams[0] instanceof Object[]) {
                requestParams = (Object[]) requestParams[0];
            }

            // noinspection unchecked
            return (NavigableMap<K, V>) jdbcTemplate.query(sqlString, requestParams, new MapResultSetExtractor<>(rowMapper));
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    @Override
    public void query(final String sqlString, final RowCallbackHandler rch, final Object... requestParams) throws WebmasterSQLException {
        try {
            jdbcTemplate.query(con -> {
                        final PreparedStatement statement = con.prepareStatement(sqlString, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
                        statement.setFetchSize(Integer.MIN_VALUE);
                        return statement;
                    }, new ArgumentPreparedStatementSetter(requestParams),
                    rs -> {
                        while (rs.next()) {
                            rch.processRow(rs);
                        }
                        return null;
                    });
        } catch (RuntimeException e) {
            throw new WebmasterSQLException(dbIdentity, e.getMessage(), e);
        }
    }

    private static class MapResultSetExtractor<K, V> implements ResultSetExtractor {
        private int rowNum = 0;
        private ParameterizedMapRowMapper<K, V> rowMapper;

        public MapResultSetExtractor(ParameterizedMapRowMapper<K, V> rowMapper) {
            this.rowMapper = rowMapper;
        }

        @Override
        public NavigableMap<K, V> extractData(ResultSet resultSet) throws SQLException, DataAccessException {
            NavigableMap<K, V> results = new TreeMap<K, V>();
            while (resultSet.next()) {
                Pair<K, V> pair = rowMapper.mapRow(resultSet, rowNum++);
                results.put(pair.first, pair.second);
            }
            return results;
        }
    }
}
