/*
 * 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
 *
 *    http://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.seatunnel.datasource.plugin.sqlserver.jdbc;

import org.apache.seatunnel.api.configuration.util.OptionRule;
import org.apache.seatunnel.datasource.plugin.api.DataSourceChannelExt;
import org.apache.seatunnel.datasource.plugin.api.DataSourcePluginException;
import org.apache.seatunnel.datasource.plugin.api.model.TableField;
import org.apache.seatunnel.datasource.plugin.api.utils.JdbcUtils;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

import java.sql.*;
import java.util.*;

import static com.google.common.base.Preconditions.checkNotNull;

@Slf4j
public class SqlServerDataSourceChannel implements DataSourceChannelExt {
    @Override
    public OptionRule getDataSourceOptions(@NonNull String pluginName) {
        return SqlServerDataSourceConfig.OPTION_RULE;
    }

    @Override
    public OptionRule getDatasourceMetadataFieldsByDataSourceName(@NonNull String pluginName) {
        return SqlServerDataSourceConfig.METADATA_RULE;
    }

    @Override
    public List<String> getTables(
            @NonNull String pluginName,
            Map<String, String> requestParams,
            String database,
            Map<String, String> option) {
        List<String> tableNames = new ArrayList<>();
        try (Connection connection = getConnection(requestParams);
                ResultSet resultSet =
                        connection
                                .getMetaData()
                                .getTables(database, null, null, new String[] {"TABLE"})) {
            while (resultSet.next()) {
                String tableName = resultSet.getString("TABLE_NAME");
                if (StringUtils.isNotBlank(tableName)) {
                    tableNames.add(tableName);
                }
            }
            return tableNames;
        } catch (ClassNotFoundException | SQLException e) {
            throw new DataSourcePluginException("get table names failed", e);
        }
    }

    @Override
    public List<String> getDatabases(
            @NonNull String pluginName, @NonNull Map<String, String> requestParams) {
        List<String> dbNames = new ArrayList<>();
        try (Connection connection = getConnection(requestParams);
                PreparedStatement statement =
                        connection.prepareStatement(
                                "SELECT name FROM sys.databases WHERE name NOT IN ('master', 'tempdb', 'model', 'msdb');");
                ResultSet re = statement.executeQuery()) {
            // filter system databases
            while (re.next()) {
                String dbName = re.getString("name");
                if (StringUtils.isNotBlank(dbName)
                        && !SqlServerDataSourceConfig.SQLSERVER_SYSTEM_DATABASES.contains(dbName)) {
                    dbNames.add(dbName);
                }
            }
            return dbNames;
        } catch (Exception ex) {
            throw new RuntimeException("get databases failed", ex);
        }
    }

    @Override
    public boolean checkDataSourceConnectivity(
            @NonNull String pluginName, @NonNull Map<String, String> requestParams) {
        try (Connection ignored = getConnection(requestParams)) {
            return true;
        } catch (Exception e) {
            throw new DataSourcePluginException("check jdbc connectivity failed", e);
        }
    }

    @Override
    public List<TableField> getTableFields(
            @NonNull String pluginName,
            @NonNull Map<String, String> requestParams,
            @NonNull String database,
            @NonNull String table) {
        List<TableField> tableFields = new ArrayList<>();
        try (Connection connection = getConnection(requestParams, null)) {
            DatabaseMetaData metaData = connection.getMetaData();
            String primaryKey = getPrimaryKey(metaData, database, table);

            String sql =
                    "SELECT\n"
                            + "\tA.name AS TABLE_NAME,\n"
                            + "\tf.value AS TABLE_COMMENT,\n"
                            + "\tB.name AS COLUMN_NAME,\n"
                            + "\tB.max_length AS COLUMN_LENGTH,\n"
                            + "\tB.scale AS COLUMN_SCALE,\n"
                            + "\tC.value AS REMARKS,\n"
                            + "\tD.name AS TYPE_NAME,\n"
                            + "\tB.is_identity AS IS_IDENTITY,\n"
                            + " CASE\n"
                            + "\t\tA.type \n"
                            + "\t\tWHEN 'u' THEN\n"
                            + "\t\t'用户定义表' \n"
                            + "\tEND AS Type,\n"
                            + "\tB.is_nullable AS IS_NULLABLE, \n"
                            + "\tde.definition COLUMN_DEFAULT "
                            + "FROM\n"
                            + "\tsys.tables AS A\n"
                            + "\tINNER JOIN sys.columns AS B ON B.object_id = A.object_id\n"
                            + "\tLEFT   JOIN sys.extended_properties AS C ON C.major_id = B.object_id \n"
                            + "\tAND C.minor_id = B.column_id\n"
                            + "\tLEFT  JOIN sys.types AS D ON D.system_type_id = B.system_type_id  and d.name!='sysname'    \n"
                            + "\tleft  JOIN sys.extended_properties AS f ON A.object_id = f.major_id \n"
                            + "\tAND f.minor_id = 0 \n"
                            + "\t  LEFT JOIN   sys.default_constraints de\ton  de.object_id = B.default_object_id \n"
                            + "WHERE\n"
                            + "\tA.name = '"
                            + table
                            + "'";
            Statement statement = connection.createStatement();
            statement.execute("use " + database);
            ResultSet resultSet = statement.executeQuery(sql);
            while (resultSet.next()) {
                TableField tableField = new TableField();
                String columnName = resultSet.getString("COLUMN_NAME");
                tableField.setPrimaryKey(false);
                if (StringUtils.isNotBlank(primaryKey) && primaryKey.equals(columnName)) {
                    tableField.setPrimaryKey(true);
                }
                tableField.setLen(resultSet.getString("COLUMN_LENGTH"));
                tableField.setScale(resultSet.getString("COLUMN_SCALE"));
                tableField.setName(columnName);

                String extra = resultSet.getString("IS_IDENTITY");
                if (StringUtils.isNotBlank(extra)) {
                    tableField.setHasAutoIncrement(true);
                }

                tableField.setType(resultSet.getString("TYPE_NAME"));
                tableField.setComment(resultSet.getString("REMARKS"));
                String columnDefault = resultSet.getString("COLUMN_DEFAULT");
                if (StringUtils.isNotBlank(columnDefault)) {
                    tableField.setDefaultValue(columnDefault.replace("(", "").replace(")", ""));
                }
                Object nullable = resultSet.getObject("IS_NULLABLE");
                tableField.setNullable(Boolean.TRUE.toString().equals(nullable.toString()));
                tableFields.add(tableField);
            }

        } catch (ClassNotFoundException | SQLException e) {
            throw new DataSourcePluginException("get table fields failed", e);
        }
        return tableFields;
    }

    @Override
    public Map<String, List<TableField>> getTableFields(
            @NonNull String pluginName,
            @NonNull Map<String, String> requestParams,
            @NonNull String database,
            @NonNull List<String> tables) {
        return null;
    }

    private String getPrimaryKey(DatabaseMetaData metaData, String dbName, String tableName)
            throws SQLException {
        ResultSet primaryKeysInfo = metaData.getPrimaryKeys(dbName, "%", tableName);
        while (primaryKeysInfo.next()) {
            return primaryKeysInfo.getString("COLUMN_NAME");
        }
        return null;
    }

    private Connection getConnection(Map<String, String> requestParams)
            throws SQLException, ClassNotFoundException {
        return getConnection(requestParams, null);
    }

    private Connection getConnection(Map<String, String> requestParams, String databaseName)
            throws SQLException, ClassNotFoundException {
        checkNotNull(requestParams.get(SqlServerOptionRule.DRIVER.key()));
        checkNotNull(requestParams.get(SqlServerOptionRule.URL.key()), "Jdbc url cannot be null");
        String url =
                JdbcUtils.replaceDatabase(
                        requestParams.get(SqlServerOptionRule.URL.key()), databaseName);
        if (requestParams.containsKey(SqlServerOptionRule.USER.key())) {
            String username = requestParams.get(SqlServerOptionRule.USER.key());
            String password = requestParams.get(SqlServerOptionRule.PASSWORD.key());
            return DriverManager.getConnection(url, username, password);
        }
        return DriverManager.getConnection(url);
    }

    @Override
    public List<Map> getTableAlls(
            @NonNull String pluginName,
            Map<String, String> requestParams,
            String database,
            String tableName,
            Map<String, String> options) {
        List<Map> tableNames = new ArrayList<>();
        try (Connection connection = getConnection(pluginName, requestParams)) {

            ResultSet resultSet = null;

            if (StringUtils.isNotBlank(database)) {

                StringBuilder sql =
                        new StringBuilder(
                                "select ROW_NUMBER() OVER (ORDER BY a.name) AS No, a.name TABLE_NAME    ,g.class_desc,g.major_id,g.class,g.value  TABLE_COMMENT  from "
                                        + " sys.tables a left join sys.extended_properties g on (a.object_id = g.major_id AND g.minor_id = 0) where 1=1 ");

                if (StringUtils.isNotBlank(tableName)) {
                    sql.append(" and  a.name='" + tableName + "'");
                }

                Statement statement = connection.createStatement();
                statement.execute("use " + database);
                resultSet = statement.executeQuery(sql.toString());
                fillTableNames(tableNames, resultSet);
            } else {

                String sqlDatabases = "SELECT * FROM sys.sysdatabases";
                Statement statement = connection.createStatement();
                resultSet = statement.executeQuery(sqlDatabases.toString());
                while (resultSet.next()) {

                    //                    ResultSetMetaData metaData = resultSet.getMetaData();
                    //                    //  获取表列数
                    //                    int columnCount = metaData.getColumnCount();
                    //                    for (int i = 1; i <= columnCount; i++) {
                    //                        tables.put(metaData.getColumnName(i),
                    // resultSet.getObject(i));
                    //                    }

                    String databaseName = resultSet.getString("name");
                    if (StringUtils.isNotBlank(databaseName)
                            && !Arrays.asList("master", "tempdb", "model", "msdb")
                                    .contains(databaseName)) {
                        StringBuilder sql =
                                new StringBuilder(
                                        "select ROW_NUMBER() OVER (ORDER BY a.name) AS No, a.name TABLE_NAME    ,g.class_desc,g.major_id,g.class,g.value  TABLE_COMMENT  from "
                                                + " sys.tables a left join sys.extended_properties g on (a.object_id = g.major_id AND g.minor_id = 0) where 1=1 ");

                        if (StringUtils.isNotBlank(tableName)) {
                            sql.append(" and  a.name='" + tableName + "'");
                        }

                        statement.execute("use " + database);
                        resultSet = statement.executeQuery(sql.toString());
                        fillTableNames(tableNames, resultSet);
                    }
                }
            }

            return tableNames;
        } catch (SQLException e) {
            throw new DataSourcePluginException("get table names failed", e);
        }
    }

    private static void fillTableNames(List<Map> tableNames, ResultSet resultSet)
            throws SQLException {
        if (Objects.nonNull(resultSet)) {
            while (resultSet.next()) {
                Map tables = new HashMap();

                ResultSetMetaData metaData = resultSet.getMetaData();
                //  获取表列数
                int columnCount = metaData.getColumnCount();
                for (int i = 1; i <= columnCount; i++) {
                    tables.put(metaData.getColumnName(i), resultSet.getObject(i));
                }

                String tableNameOld = resultSet.getString("TABLE_NAME");
                if (StringUtils.isNotBlank(tableNameOld)) {
                    tables.put("TABLE_NAME", tableNameOld);
                }

                String tableComment = resultSet.getString("TABLE_COMMENT");

                if (StringUtils.isNotBlank(tableComment)) {
                    tables.put("TABLE_COMMENT", tableComment);
                }
                tableNames.add(tables);
            }
        }
    }

    @Override
    public Map getTableName(
            @NonNull String pluginName,
            Map<String, String> requestParams,
            String database,
            Map<String, String> options,
            String tableName) {

        try (Connection connection = getConnection(requestParams, database)) {
            String sql =
                    "SELECT * FROM sys.all_objects WHERE object_id = OBJECT_ID(N'["
                            + database
                            + "].[dbo].["
                            + tableName
                            + "]') AND type IN ('U')";
            Statement statement = connection.createStatement();
            ResultSet resultSet = statement.executeQuery(sql);
            ResultSetMetaData metaData = resultSet.getMetaData();
            // 字段的个数
            int count = metaData.getColumnCount();
            // 创建List集合
            List<Map> list = new ArrayList<>();
            // 处理结果集
            while (resultSet.next()) {
                // 创建Map集合
                Map map = new HashMap<>();
                // 根据字段的个数, 循环
                for (int i = 0; i < count; i++) {
                    // 获取字段的名字
                    String columnName = metaData.getColumnName(i + 1);
                    // 获取字段对应的值
                    Object object = resultSet.getObject(columnName);
                    // 将字段名和字段值, 存入map集合中
                    map.put(columnName, object);
                }
                // 将map集合添加到List集合中
                list.add(map);
            }
            if (CollectionUtils.isNotEmpty(list)) {
                return list.get(0);
            }
        } catch (SQLException e) {
            throw new DataSourcePluginException("get table fields failed", e);
        } catch (ClassNotFoundException e) {
            throw new DataSourcePluginException("get table fields failed", e);
        }

        return null;
    }
}
