package org.jeecg.modules.stm.polardb;

import lombok.extern.slf4j.Slf4j;
import org.jeecg.common.api.vo.Result;
import org.springframework.web.bind.annotation.*;

import java.sql.Connection;
import java.sql.*;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@RestController
@RequestMapping("/polar")
@Slf4j
public class MySQLAccess {

    @GetMapping(value = "queryTest")
    public Result<Map<String,Object>> polarDbTest(@RequestParam(name = "queryCond") String queryCond) {

        StringBuffer resSb=  new StringBuffer();
        Map<String,Object> retMap = new HashMap<>();

        try {

            List<String> resSqlList = new ArrayList<String>();
            Connection connection = null;
            PreparedStatement statement = null;
            ResultSet resultSet = null;
            List tableList = new ArrayList();
            LinkedHashMap<String,String> tableHeads = new LinkedHashMap<>();
            List<String> dataColumns = new ArrayList<>();

            try {
                connection = DBUtil.getConnection();
                String sql = "/*polar4ai*/SELECT * FROM PREDICT (MODEL _polar4ai_nl2sql, select '{" + queryCond + "}') WITH (basic_index_name='schema_index')";
                statement = connection.prepareStatement(sql);
                resultSet = statement.executeQuery();

                while (resultSet.next()) {
                    String result = resultSet.getString("result");
                    log.info(result);
                    resSqlList.add(result);
                }
            } catch (Exception e) {
                log.error(e.toString());
                e.printStackTrace();
            } finally {
                if(statement!=null) statement.close();
                if(resultSet!=null) resultSet.close();
                DBUtil.closeConnection(connection);
            }

            String url = "jdbc:mysql://localhost:3306/stm_db?useUnicode=true&characterEncoding=utf-8&useSSL=false&serverTimezone=UTC";
            String username = "root";
            String password = "ZhongRunChangHong/123";
            Connection connection2 = DBUtil.getConnection(url, username, password);

            try {
                for (String tmpSql : resSqlList) {
                    tableHeads = getColumnComment(connection2.prepareStatement(tmpSql),getTableName(tmpSql));
                    statement = connection2.prepareStatement(tmpSql);
                    resultSet = statement.executeQuery();
                    ResultSetMetaData metaData = resultSet.getMetaData();
                    int columnCount = metaData.getColumnCount();
                    for (int i = 1; i <= columnCount; i++) {
                        String columnName = metaData.getColumnName(i);
                        dataColumns.add(columnName);
                    }
                    while (resultSet.next()) {
                        Map<String, Object> tableRowMap = new HashMap<String, Object>();
                        for (int i = 1; i <= columnCount; i++) {
                            String columnName = metaData.getColumnName(i);
                            String strLabelName = metaData.getColumnLabel(i);
                            Object value = resultSet.getObject(i);
                            tableRowMap.put(columnName, value);

                        }
                        tableList.add(tableRowMap);
                    }

                }

            } catch (Exception e) {
                log.error(e.toString());
                e.printStackTrace();
            } finally {
                if(statement!=null) statement.close();
                if(resultSet!=null) resultSet.close();
                DBUtil.closeConnection(connection2);
            }
            retMap.put("sqls", resSqlList);
            retMap.put("datas", tableList);
            retMap.put("heads", tableHeads);
            retMap.put("showColumns", dataColumns);

        } catch (Exception e) {
            log.error(e.toString());
            e.printStackTrace();
            Result.error("查询失败:");
        }
        return  Result.OK(retMap);
    }

    private String getTableName(String strSql) {
        String regex = "(?i)FROM\\s+(\\w+)";
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(strSql);
        String tableName = null;
        if (matcher.find()) {
            tableName = matcher.group(1);
        }
        return tableName;
    }

    private LinkedHashMap<String,String> getColumnComment( PreparedStatement statement,String tableName) {
        // 5. 查询列的注释
        String commentsQuery = "SELECT COLUMN_NAME, COLUMN_COMMENT FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '"+tableName+"' ORDER BY ordinal_position";
        LinkedHashMap<String, String> retMap = new LinkedHashMap<>();
        ResultSet commentsResultSet = null;
        try {
            commentsResultSet = statement.executeQuery(commentsQuery);
            while (commentsResultSet.next()) {
                String columnName = commentsResultSet.getString("COLUMN_NAME");
                String columnComment = commentsResultSet.getString("COLUMN_COMMENT");
//                if(columnComment==null||columnComment.trim().length()==0) {
//                    columnComment = columnName;
//                }
                retMap.put(columnName,columnComment);
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if(commentsResultSet!=null){
                try {
                    commentsResultSet.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
            return retMap;
        }
    }

    @GetMapping(value = "polarAITest")
    public Result<Map<String,Object>> polarAITest(@RequestParam(name = "queryCond") String queryCond) {
        StringBuffer resSb = new StringBuffer();
        Map<String, Object> retMap = new HashMap<>();
        try {
            List<String> resSqlList = new ArrayList<String>();
            Connection connection = null;
            PreparedStatement statement = null;
            ResultSet resultSet = null;
            List<String> dataColumns = new ArrayList<>();
            try {

                String sql = "/*polar4ai*/SELECT * FROM PREDICT(MODEL _polar4ai_db_search, SELECT '"+queryCond+"') WITH (basic_index_name='schema_index', index_name='vector_table', content_column_name='content', vec_column_name='content_vector', top_k=3, score_threshold=0.01);";
                connection = DBUtil.getConnection();
                statement = connection.prepareStatement(sql);
                resultSet = statement.executeQuery();

                while (resultSet.next()) {

                    String result = resultSet.getString("final_result");
                    if(result!=null&&!result.equals("")){
                        retMap.put("result", result);
                    }

                    String n2sqlResult = resultSet.getString("n2sql_result");
                    if(n2sqlResult!=null&&!n2sqlResult.equals("")) {
                        retMap.put("n2sqlResult", n2sqlResult);
                    }

                    String docResult = resultSet.getString("doc_result");
                    if(docResult!=null&&!docResult.equals("")) {
                        retMap.put("docResult", docResult);
                    }

                    String sqlResult = resultSet.getString("sql_result");
                    if(sqlResult!=null&&!sqlResult.equals("")) {
                        retMap.put("sqlResult", sqlResult);
                    }
                }


            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (statement != null) statement.close();
                if (resultSet != null) resultSet.close();
                DBUtil.closeConnection(connection);
            }

            return Result.OK(retMap);
        }  catch (Exception e) {
            e.printStackTrace();
        }

        return null;

    }

    @PostMapping(value = "polarAISummary")
    public Result<Map<String,Object>> polarAISummary(@RequestBody Map<String,String> formData) {
        StringBuffer resSb = new StringBuffer();
        Map<String, Object> retMap = new HashMap<>();
        try {
            Connection connection = null;
            PreparedStatement statement = null;
            ResultSet resultSet = null;
            //queryCond = "评级年度内，董事会召开涉及信息科技议题的会议次数为0次。这表明在该评定年度内，董事会层面没有专门就信息科技相关议题召开会议。希望根据上述现状生成审计发现、影响分析及改进建议。";
            try {
                String queryCond = formData.get("queryCond");
                String sql = "/*polar4ai*/SELECT * FROM PREDICT (MODEL _polar4ai_tongyi, SELECT '"+queryCond+"') with ();";
                connection = DBUtil.getConnection();
                statement = connection.prepareStatement(sql);
                resultSet = statement.executeQuery();
                while (resultSet.next()) {
                    String result = resultSet.getString("result");
                    retMap.put("summary",result);
                    break;
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (statement != null) statement.close();
                if (resultSet != null) resultSet.close();
                DBUtil.closeConnection(connection);
            }

            return Result.OK(retMap);
        }  catch (Exception e) {
            e.printStackTrace();
        }

        return null;

    }

    @PostMapping(value = "polarAIHight")
    public Result<Map<String,Object>> polarAIHight(@RequestBody Map<String,String> formData) {
        StringBuffer resSb = new StringBuffer();
        Map<String, Object> retMap = new HashMap<>();
        try {
            Connection connection = null;
            PreparedStatement statement = null;
            ResultSet resultSet = null;
            //queryCond = "评级年度内，董事会召开涉及信息科技议题的会议次数为0次。这表明在该评定年度内，董事会层面没有专门就信息科技相关议题召开会议。希望根据上述现状生成审计发现、影响分析及改进建议。";
            try {
                String queryCond = formData.get("queryCond");
                String sql = "/*polar4ai*/SELECT * FROM PREDICT(MODEL _polar4ai_rag, SELECT '" + queryCond + "') WITH (" +
                        "    index_name='file_index_table'," +
                        "    primary_key='chunk_id'," +
                        "    vec_top_k=5," +
                        "    vec_score_threshold=0.0," +
                        "    pos_top_k=3,\n" +
                        "    pos_score_threshold=0.0," +
                        "    use_vec=1," +
                        "    use_pos=1," +
                        "    hybrid_type=0," +
                        "    hybrid_score=0.8," +
                        "    ranker_model=''," +
                        "    category_ids=''," +
                        "    file_id=''" +
                        ")";
                connection = DBUtil.getConnection();
                statement = connection.prepareStatement(sql);
                resultSet = statement.executeQuery();
                while (resultSet.next()) {
                    String recallContent = resultSet.getString("recall_content");
                    String generateResult = resultSet.getString("generate_result");
                    retMap.put("generateResult", generateResult);
                    retMap.put("recallContent", recallContent);
                    break;
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (statement != null) statement.close();
                if (resultSet != null) resultSet.close();
                DBUtil.closeConnection(connection);
            }

            return Result.OK(retMap);
        } catch (Exception e) {
            e.printStackTrace();
        }

        return null;
    }


}

