/**
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.atlas.trino.client;

import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.StringUtils;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TrinoClientHelper {
    private final String jdbcUrl;
    private final String username;
    private final String password;

    public TrinoClientHelper(Configuration atlasConf) {
        this.jdbcUrl  = atlasConf.getString("atlas.trino.jdbc.address");
        this.username = atlasConf.getString("atlas.trino.jdbc.user");
        this.password = atlasConf.getString("atlas.trino.jdbc.password", "");
    }

    public Map<String, String> getAllTrinoCatalogs() {
        Map<String, String> catalogs = new HashMap<>();

        try (Connection connection = getTrinoConnection();
                Statement stmt = connection.createStatement()) {
            String    query = "SELECT catalog_name, connector_name FROM system.metadata.catalogs";
            ResultSet rs    = stmt.executeQuery(query);

            while (rs.next()) {
                catalogs.put(rs.getString("catalog_name"), rs.getString("connector_name"));
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }

        return catalogs;
    }

    public List<String> getTrinoSchemas(String catalog, String schemaToImport) throws SQLException {
        List<String> schemas = new ArrayList<>();

        try (Connection connection = getTrinoConnection();
                Statement stmt = connection.createStatement()) {
            StringBuilder query = new StringBuilder();

            query.append("SELECT schema_name FROM ").append(catalog).append(".information_schema.schemata");

            if (StringUtils.isNotEmpty(schemaToImport)) {
                query.append(" where schema_name = '").append(schemaToImport).append("'");
            }

            ResultSet rs = stmt.executeQuery(query.toString());

            while (rs.next()) {
                schemas.add(rs.getString("schema_name"));
            }
        }

        return schemas;
    }

    public Map<String, Map<String, Object>> getTrinoTables(String catalog, String schema, String tableToImport) throws SQLException {
        Map<String, Map<String, Object>> tables = new HashMap<>();

        try (Connection connection = getTrinoConnection();
                Statement stmt = connection.createStatement()) {
            StringBuilder query = new StringBuilder();

            query.append("SELECT table_name, table_type FROM ").append(catalog).append(".information_schema.tables WHERE table_schema = '").append(schema).append("'");

            if (StringUtils.isNotEmpty(tableToImport)) {
                query.append(" and table_name = '").append(tableToImport).append("'");
            }

            ResultSet rs = stmt.executeQuery(query.toString());

            while (rs.next()) {
                Map<String, Object> tableMetadata = new HashMap<>();

                tableMetadata.put("table_name", rs.getString("table_name"));
                tableMetadata.put("table_type", rs.getString("table_type"));

                tables.put(rs.getString("table_name"), tableMetadata);
            }
        }

        return tables;
    }

    public Map<String, Map<String, Object>> getTrinoColumns(String catalog, String schema, String table) throws SQLException {
        Map<String, Map<String, Object>> columns = new HashMap<>();

        try (Connection connection = getTrinoConnection();
                Statement stmt = connection.createStatement()) {
            StringBuilder query = new StringBuilder();

            query.append("SELECT column_name, ordinal_position, column_default, is_nullable, data_type FROM ").append(catalog).append(".information_schema.columns WHERE table_schema = '").append(schema).append("' AND table_name = '").append(table).append("'");

            ResultSet rs = stmt.executeQuery(query.toString());

            while (rs.next()) {
                Map<String, Object> columnMetadata = new HashMap<>();

                columnMetadata.put("ordinal_position", rs.getInt("ordinal_position"));
                columnMetadata.put("column_default", rs.getString("column_default"));
                columnMetadata.put("column_name", rs.getString("column_name"));

                if (StringUtils.isNotEmpty(rs.getString("is_nullable"))) {
                    if (StringUtils.equalsIgnoreCase(rs.getString("is_nullable"), "YES")) {
                        columnMetadata.put("is_nullable", true);
                    } else {
                        columnMetadata.put("is_nullable", false);
                    }
                }

                columnMetadata.put("data_type", rs.getString("data_type"));

                columns.put(rs.getString("column_name"), columnMetadata);
            }
        }

        return columns;
    }

    private Connection getTrinoConnection() throws SQLException {
        return DriverManager.getConnection(jdbcUrl, username, password);
    }
}
