package com.njzscloud.common.gen.support;

import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

@Slf4j
@RequiredArgsConstructor
public class DbMetaData {
    private final DataSource dataSource;

    public String getDatabaseName() {
        try (Connection conn = dataSource.getConnection()) {
            return conn.getCatalog();
        } catch (Exception e) {
            log.error("获取数据库名称失败", e);
            return null;
        }
    }

    public List<Map<String, Object>> getTableMetaData(String tableName) {
        try (Connection connection = dataSource.getConnection()) {
            return getTables(connection, tableName);
        } catch (Exception e) {
            log.error("获取数据库元数据失败", e);
            return Collections.emptyList();
        }
    }

    private List<Map<String, Object>> getTableColumns(Connection conn, String tableName) throws Exception {
        List<Map<String, Object>> columns = new ArrayList<>();

        DatabaseMetaData metaData = conn.getMetaData();
        try (ResultSet columnsData = metaData.getColumns(null, null, tableName, "%")) {
            while (columnsData.next()) {
                columns.add(MapUtil.<String, Object>builder()
                    .put("name", columnsData.getString("COLUMN_NAME"))
                    .put("dataType", columnsData.getString("TYPE_NAME"))
                    .put("size", columnsData.getInt("COLUMN_SIZE"))
                    .put("nullable", columnsData.getInt("NULLABLE") == DatabaseMetaData.columnNullable)
                    .put("defaultValue", columnsData.getString("COLUMN_DEF"))
                    .put("comment", columnsData.getString("REMARKS"))
                    .build()
                );
            }
        }
        try (ResultSet primaryKeyData = metaData.getPrimaryKeys(null, null, tableName)) {
            while (primaryKeyData.next()) {
                String primaryKeyColumn = primaryKeyData.getString("COLUMN_NAME");
                columns.forEach(column -> column.put("primaryKey", column.get("name").equals(primaryKeyColumn)));
            }
        }
        return columns;
    }

    private List<Map<String, Object>> getTables(Connection conn, String tableName) throws Exception {
        List<Map<String, Object>> tableInfos = new ArrayList<>();
        DatabaseMetaData metaData = conn.getMetaData();
        try (ResultSet tablesData = metaData.getTables(null, null, StrUtil.isNotBlank(tableName) ? tableName : "%", new String[]{"TABLE"})) {
            while (tablesData.next()) {
                String tableName_ = tablesData.getString("TABLE_NAME");
                List<Map<String, Object>> columns = getTableColumns(conn, tableName_);
                tableInfos.add(MapUtil.<String, Object>builder()
                    .put("name", tableName_)
                    .put("comment", tablesData.getString("REMARKS"))
                    .put("columns", columns)
                    .build());
            }
        }
        return tableInfos;
    }

}
