package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"message-service/commons/chatgpttemplate"
	"message-service/commons/config"
	"message-service/commons/utils"
	"os"
	"strconv"
	"strings"

	"message-service/internal/bean/dto"
	service "message-service/internal/service/central"

	"github.com/pingcap/log"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
)

const (
	host                    = "http://127.0.0.1:10010"
	defaultConfigPath       = "/etc/dbaas/config.yaml"
	defaultTableInfosPath   = "./table_info.json"
	defaultOutputPath       = "./table_info_with_question.json"
	defaultQuestionFilePath = "./questionFile.list"
)

var exePath, _ = os.Executable()

type TableInfo struct {
	Database  string   `json:"database"`
	Table     string   `json:"table"`
	Sql       string   `json:"sql"`
	ColName   []string `json:"col_names"`
	Questions []string `json:"questions"`
}

func main() {
	var configPath string
	var tableInfosPath string
	var outputPath string
	var questionFilePath string
	flag.StringVar(&configPath, "config-path", defaultConfigPath, "DBaaS configuration file path")
	flag.StringVar(&tableInfosPath, "table-infos", defaultTableInfosPath, "DBaaS configuration file path")
	flag.StringVar(&outputPath, "output", defaultOutputPath, "")
	flag.StringVar(&questionFilePath, "question-file", defaultQuestionFilePath, "DBaaS configuration file path")
	flag.Parse()
	err := config.InitConfig(configPath)
	if err != nil {
		panic(fmt.Sprintf("Initialize configuration failed, err: %+v", err))
	}
	tableInfos, err := ReadTableInfos(tableInfosPath)
	if err != nil {
		panic(fmt.Sprintf("ReadTableInfos failed, err: %+v , file is %s", err, tableInfosPath))
	}
	questionNum := 20
	var databaseMap = make(map[string][]*dataapi.TableMeta)
	for index, tableInfo := range tableInfos {
		meta := dataapi.TableMeta{
			Database: tableInfo.Database,
			Table:    tableInfo.Table,
			Columns:  tableInfo.ColName,
		}
		if _, ok := databaseMap[tableInfo.Database]; !ok {
			databaseMap[tableInfo.Database] = []*dataapi.TableMeta{}
		}
		databaseMap[tableInfo.Database] = append(databaseMap[tableInfo.Database], &meta)
		if len(tableInfo.Questions) == questionNum {
			log.Info("skip gen question for table " + tableInfo.Table + " database " + tableInfo.Database)
			continue
		}
		questionStr, err := GenQuestionsByCompletion([]*dataapi.TableMeta{&meta}, questionNum)
		if err != nil {
			continue
		}
		questionList := strings.Split(questionStr, "\n")
		tableInfos[index].Questions = []string{}
		for _, question := range questionList {
			if question == "" {
				continue
			}
			tableInfos[index].Questions = append(tableInfos[index].Questions, question)
		}
	}
	outputStr, err := json.MarshalIndent(tableInfos, "", "\t")
	utils.WriteStringToFile(outputPath, string(outputStr))

	// Gen Questions for database
	var summaryTableInfos []TableInfo
	for database, metas := range databaseMap {
		log.Info("Gen questions for summaryTable. " + database + " meta num " + strconv.Itoa(len(metas)))
		if len(metas) <= 1 {
			continue
		}
		questionStr, err := GenQuestionsByCompletion(metas, questionNum+10)
		if err != nil {
			continue
		}
		questionList := strings.Split(questionStr, "\n")
		questions, _ := utils.ReadFileToList(questionFilePath)
		tableInfo := TableInfo{
			Database:  database,
			Questions: *questions,
		}

		for _, question := range questionList {
			if question == "" {
				continue
			}
			tableInfo.Questions = append(tableInfo.Questions, question)
		}
		summaryTableInfos = append(summaryTableInfos, tableInfo)
	}

	outputStr, err = json.MarshalIndent(summaryTableInfos, "", "\t")
	utils.WriteStringToFile(outputPath+"_database", string(outputStr))
}

func GenQuestionsByCompletion(metas []*dataapi.TableMeta, num int) (string, error) {
	ctx := context.Background()
	prompt := chatgpttemplate.FormatPromptForGenDatabaseQuestions(metas[0].Database, metas, num)
	log.Info("prompt is " + prompt)
	param := dto.CreateCompletionParam{
		Model:       "text-davinci-003",
		Prompt:      []string{prompt},
		MaxTokens:   1024,
		Temperature: 1,
		TopP:        1,
		N:           1,
		Stream:      false,
		Logprobs:    2,
		Stop:        []string{"#", "---"},
	}
	completion, err := service.CreateCompletionToOpenAI(ctx, &param)
	if err != nil {
		return "", err
	}
	return completion, nil
}

func ReadTableInfos(tableInfoFile string) ([]TableInfo, error) {
	content, _ := utils.ReadFileToString(tableInfoFile)
	contentBytes := []byte(content)
	tableInfos := make([]TableInfo, 0)
	err := json.Unmarshal(contentBytes, &tableInfos)
	if err != nil {
		return nil, err
	}
	return tableInfos, nil
}