package com.alibaba.datax.plugin.writer.hbase11xsqlwriter;

import com.alibaba.datax.common.element.Column;
import com.alibaba.datax.common.element.Record;
import com.alibaba.datax.common.exception.DataXException;
import com.alibaba.datax.common.plugin.RecordReceiver;
import com.alibaba.datax.common.plugin.TaskPluginCollector;
import com.alibaba.datax.common.util.Configuration;
import com.google.common.collect.Lists;
import org.apache.phoenix.schema.PTable;
import org.apache.phoenix.schema.types.PDataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.math.BigDecimal;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Types;
import java.util.Arrays;
import java.util.List;

/**
 * @author yanghan.y
 */
public class HbaseSQLWriterTask {
    private final static Logger LOG = LoggerFactory.getLogger(HbaseSQLWriterTask.class);

    private TaskPluginCollector taskPluginCollector;
    private HbaseSQLWriterConfig cfg;
    private Connection connection = null;
    private PreparedStatement ps = null;
    // 需要向hbsae写入的列的数量,即用户配置的column参数中列的个数。时间戳不包含在内
    private int numberOfColumnsToWrite;
    // 期待从源头表的Record中拿到多少列
    private int numberOfColumnsToRead;
    private boolean needExplicitVersion = false;
    private int[] columnTypes;

    public HbaseSQLWriterTask(Configuration configuration) {
        // 这里仅解析配置，不访问远端集群，配置的合法性检查在writer的init过程中进行
        cfg = HbaseSQLHelper.parseConfig(configuration);
    }

    public void startWriter(RecordReceiver lineReceiver, TaskPluginCollector taskPluginCollector) {
        this.taskPluginCollector = taskPluginCollector;
        Record record;
        try {
            // 准备阶段
            prepare();

            List<Record> buffer = Lists.newArrayListWithExpectedSize(cfg.getBatchSize());
            while ((record = lineReceiver.getFromReader()) != null) {
                // 校验列数量是否符合预期
                if (record.getColumnNumber() != numberOfColumnsToRead) {
                    throw DataXException.asDataXException(HbaseSQLWriterErrorCode.ILLEGAL_VALUE,
                            "数据源给出的列数量[" + record.getColumnNumber() + "]与您配置中的列数量[" + numberOfColumnsToRead +
                            "]不同, 请检查您的配置 或者 联系 Hbase 管理员.");
                }

                buffer.add(record);
                if (buffer.size() > cfg.getBatchSize()) {
                    doBatchUpsert(buffer);
                    buffer.clear();
                }
            } // end while loop

            // 处理剩余的record
            if (!buffer.isEmpty()) {
                doBatchUpsert(buffer);
                buffer.clear();
            }
        } catch (Throwable t) {
            // 确保所有异常都转化为DataXException
            throw DataXException.asDataXException(HbaseSQLWriterErrorCode.PUT_HBASE_ERROR, t);
        } finally {
            close();
        }
    }

    private void prepare() throws SQLException {
        if (connection == null) {
            connection = HbaseSQLHelper.getJdbcConnection(cfg);
            connection.setAutoCommit(false);    // 批量提交
        }

        if (ps == null) {
            // 一个Task的生命周期中只使用一个PreparedStatement对象，所以，在
            ps = createPreparedStatement();
            columnTypes = getColumnSqlType(cfg.getColumns());
        }
    }

    private void close()  {
        if (ps != null) {
            try {
                ps.close();
            } catch (SQLException e) {
                // 不会出错
                LOG.error("Failed closing PreparedStatement", e);
            }
        }
        if (connection != null) {
            try {
                connection.close();
            } catch (SQLException e) {
                // 不会出错
                LOG.error("Failed closing Connection", e);
            }
        }
    }

    /**
     * 批量提交一组数据，如果失败，则尝试一行行提交，如果仍然失败，抛错给用户
     */
    private void doBatchUpsert(List<Record> records) throws SQLException {
        try {
            // 将所有record提交到connection缓存
            for (Record r : records) {
                setupStatement(r);
                ps.executeUpdate();
            }

            // 将缓存的数据提交到hbase
            connection.commit();
        } catch (SQLException e) {
            LOG.error("Failed batch committing " + records.size() + " records", e);

            // 批量提交失败，则一行行重试，以确定那一行出错
            connection.rollback();
            doSingleUpsert(records);
        } catch (Exception e) {
            throw DataXException.asDataXException(HbaseSQLWriterErrorCode.PUT_HBASE_ERROR, e);
        }
    }

    /**
     * 单行提交，将出错的行记录到脏数据中。由脏数据收集模块判断任务是否继续
     */
    private void doSingleUpsert(List<Record> records) throws SQLException {
        for (Record r : records) {
            try {
                setupStatement(r);
                ps.executeUpdate();
                connection.commit();
            } catch (SQLException e) {
                //出错了，记录脏数据
                LOG.error("Failed writing hbase", e);
                this.taskPluginCollector.collectDirtyRecord(r, e);
            }
        }
    }

    /**
     * 生成sql模板，并根据模板创建PreparedStatement
     */
    private PreparedStatement createPreparedStatement() throws SQLException {
        // 生成列名集合，列之间用逗号分隔： col1,col2,col3,...
        StringBuilder columnNamesBuilder = new StringBuilder();
        for (String col : cfg.getColumns()) {
            // 列名使用双引号，则不自动转换为全大写，而是保留用户配置的大小写
            columnNamesBuilder.append("\"");
            columnNamesBuilder.append(col);
            columnNamesBuilder.append("\"");
            columnNamesBuilder.append(",");
        }
        columnNamesBuilder.setLength(columnNamesBuilder.length() - 1);   // 移除末尾多余的逗号
        String columnNames = columnNamesBuilder.toString();
        numberOfColumnsToWrite = cfg.getColumns().size();
        numberOfColumnsToRead = numberOfColumnsToWrite;   // 开始的时候，要读的列数娱要写的列数相等

        // 生成UPSERT模板
        String tableName = cfg.getTableName();
        // 表名使用双引号，则不自动转换为全大写，而是保留用户配置的大小写
        StringBuilder upsertBuilder =
                new StringBuilder("upsert into \"" + tableName + "\" (" + columnNames + " ) values (");
        for (int i = 0; i < cfg.getColumns().size(); i++) {
            upsertBuilder.append("?,");
        }
        upsertBuilder.setLength(upsertBuilder.length() - 1);  // 移除末尾多余的逗号
        upsertBuilder.append(")");

        String sql = upsertBuilder.toString();
        PreparedStatement ps = connection.prepareStatement(sql);
        LOG.debug("SQL template generated: " + sql);
        return ps;
    }

    /**
     * 根据列名来从数据库元数据中获取这一列对应的SQL类型
     */
    private int[] getColumnSqlType(List<String> columnNames) throws SQLException {
        int[] types = new int[numberOfColumnsToWrite];
        PTable ptable = HbaseSQLHelper.getTableSchema(connection, cfg.getTableName());

        for (int i = 0; i < columnNames.size(); i++) {
            String name = columnNames.get(i);
            PDataType type = ptable.getColumnForColumnName(name).getDataType();
            types[i] = type.getSqlType();
            LOG.debug("Column name : " + name + ", sql type = " + type.getSqlType() + " " + type.getSqlTypeName());
        }
        return types;
    }

    private void setupStatement(Record record) throws SQLException {
        // 一开始的时候就已经校验过record中的列数量与ps中需要的值数量相等
        for (int i = 0; i < numberOfColumnsToWrite; i++) {
            Column col = record.getColumn(i);
            int sqlType = columnTypes[i];
            // PreparedStatement中的索引从1开始，所以用i+1
            setupColumn(i + 1, sqlType, col);
        }
    }

    private void setupColumn(int pos, int sqlType, Column col) throws SQLException {
        if (col.getRawData() != null) {
            switch (sqlType) {
                case Types.CHAR:
                case Types.VARCHAR:
                    ps.setString(pos, col.asString());
                    break;

                case Types.BINARY:
                case Types.VARBINARY:
                    ps.setBytes(pos, col.asBytes());
                    break;

                case Types.BOOLEAN:
                    ps.setBoolean(pos, col.asBoolean());
                    break;

                case Types.TINYINT:
                case Constant.TYPE_UNSIGNED_TINYINT:
                    ps.setByte(pos, col.asLong().byteValue());
                    break;

                case Types.SMALLINT:
                case Constant.TYPE_UNSIGNED_SMALLINT:
                    ps.setShort(pos, col.asLong().shortValue());
                    break;

                case Types.INTEGER:
                case Constant.TYPE_UNSIGNED_INTEGER:
                    ps.setInt(pos, col.asLong().intValue());
                    break;

                case Types.BIGINT:
                case Constant.TYPE_UNSIGNED_LONG:
                    ps.setLong(pos, col.asLong());
                    break;

                case Types.FLOAT:
                    ps.setFloat(pos, col.asDouble().floatValue());
                    break;

                case Types.DOUBLE:
                    ps.setDouble(pos, col.asDouble());
                    break;

                case Types.DECIMAL:
                    ps.setBigDecimal(pos, col.asBigDecimal());
                    break;

                case Types.DATE:
                case Constant.TYPE_UNSIGNED_DATE:
                    ps.setDate(pos, new java.sql.Date(col.asDate().getTime()));
                    break;

                case Types.TIME:
                case Constant.TYPE_UNSIGNED_TIME:
                    ps.setTime(pos, new java.sql.Time(col.asDate().getTime()));
                    break;

                case Types.TIMESTAMP:
                case Constant.TYPE_UNSIGNED_TIMESTAMP:
                    ps.setTimestamp(pos, new java.sql.Timestamp(col.asDate().getTime()));
                    break;

                default:
                    throw DataXException.asDataXException(HbaseSQLWriterErrorCode.ILLEGAL_VALUE,
                            "不支持您配置的列类型:" + sqlType + ", 请检查您的配置 或者 联系 Hbase 管理员.");

            } // end switch
        } else {
            // 没有值，按空值的配置情况处理
            switch (cfg.getNullMode()){
                case Skip:
                    // 跳过空值，则不插入该列,
                    ps.setNull(pos, sqlType);
                    break;

                case Empty:
                    // 插入"空值"，请注意不同类型的空值不同
                    // 另外，对SQL来说，空值本身是有值的，这与直接操作HBASE Native API时的空值完全不同
                    ps.setObject(pos, getEmptyValue(sqlType));
                    break;

                default:
                    // nullMode的合法性在初始化配置的时候已经校验过，这里一定不会出错
                    throw DataXException.asDataXException(HbaseSQLWriterErrorCode.ILLEGAL_VALUE,
                        "Hbasewriter 不支持该 nullMode 类型: " + cfg.getNullMode() +
                        ", 目前支持的 nullMode 类型是:" + Arrays.asList(NullModeType.values()));
            }
        }
    }

    /**
     * 根据类型获取"空值"
     * 值类型的空值都是0，bool是false，String是空字符串
     * @param sqlType sql数据类型，定义于{@link Types}
     */
    private Object getEmptyValue(int sqlType) {
        switch (sqlType) {
            case Types.VARCHAR:
                return "";

            case Types.BOOLEAN:
                return false;

            case Types.TINYINT:
            case Constant.TYPE_UNSIGNED_TINYINT:
                return (byte) 0;

            case Types.SMALLINT:
            case Constant.TYPE_UNSIGNED_SMALLINT:
                return (short) 0;

            case Types.INTEGER:
            case Constant.TYPE_UNSIGNED_INTEGER:
                return (int) 0;

            case Types.BIGINT:
            case Constant.TYPE_UNSIGNED_LONG:
                return (long) 0;

            case Types.FLOAT:
                return (float) 0.0;

            case Types.DOUBLE:
                return (double) 0.0;

            case Types.DECIMAL:
                return new BigDecimal(0);

            case Types.DATE:
            case Constant.TYPE_UNSIGNED_DATE:
                return new java.sql.Date(0);

            case Types.TIME:
            case Constant.TYPE_UNSIGNED_TIME:
                return new java.sql.Time(0);

            case Types.TIMESTAMP:
            case Constant.TYPE_UNSIGNED_TIMESTAMP:
                return new java.sql.Timestamp(0);

            case Types.BINARY:
            case Types.VARBINARY:
                return new byte[0];

            default:
                throw DataXException.asDataXException(HbaseSQLWriterErrorCode.ILLEGAL_VALUE,
                        "不支持您配置的列类型:" + sqlType + ", 请检查您的配置 或者 联系 Hbase 管理员.");
        }
    }
}
