package main

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

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

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

const (
	defaultConfigPath     = "/etc/dbaas/config.yaml"
	defaultTableInfosPath = "./table_info_with_question.json"
	defaultOutputPath     = "./table_info_with_sql.json"
)

var wg sync.WaitGroup
var exePath, _ = os.Executable()

type TableInfo struct {
	Database  string   `json:"database"`
	Table     string   `json:"table"`
	ColNames  []string `json:"col_names"`
	Questions []string `json:"questions"`
	Answers   []string `json:"answers"`
}

func GenAnswersByTableInfoFile(tableInfosPath string, outputPath string, topP float64) map[string][]TableInfo {
	var databaseMap = map[string][]TableInfo{}
	finalOutputPath := outputPath + "__top_p_" + fmt.Sprintf("%.6f", topP)
	var tableInfos []TableInfo
	_, err := os.Stat(finalOutputPath)
	if err == nil {
		tableInfos, err = ReadTableInfos(finalOutputPath)
	} else {
		tableInfos, err = ReadTableInfos(tableInfosPath)
	}
	if err != nil {
		panic(fmt.Sprintf("ReadTableInfos failed, err: %+v , file is %s", err, tableInfosPath))
	}
	for index, tableInfo := range tableInfos {
		if _, ok := databaseMap[tableInfo.Database]; !ok {
			databaseMap[tableInfo.Database] = []TableInfo{}
		}
		databaseMap[tableInfo.Database] = append(databaseMap[tableInfo.Database], tableInfo)
		if tableInfos[index].Answers == nil || len(tableInfo.Answers) != len(tableInfo.Questions) {
			tableInfos[index].Answers = make([]string, len(tableInfo.Questions))
		}
		for qIndex, question := range tableInfo.Questions {
			if tableInfos[index].Answers[qIndex] != "" {
				log.Info("skip answer - question: " + question + " answer: " + tableInfos[index].Answers[qIndex])
				continue
			}
			wg.Add(1)
			time.Sleep(200 * time.Millisecond)
			go ParallelGenSql(&tableInfos, index, qIndex, topP)
		}
		wg.Wait()
		outputStr, _ := json.MarshalIndent(tableInfos, "", "\t")
		utils.WriteStringToFile(finalOutputPath, string(outputStr))
	}
	outputStr, err := json.MarshalIndent(tableInfos, "", "\t")
	utils.WriteStringToFile(finalOutputPath, string(outputStr))
	return databaseMap
}

func GenDBAnswersByTableInfoFile(tableInfosPath string, outputPath string, topP float64, databaseMap map[string][]TableInfo) {
	finalOutputPath := outputPath + "__top_p_" + fmt.Sprintf("%.6f", topP)
	var tableInfos []TableInfo
	_, err := os.Stat(finalOutputPath)
	if err == nil {
		tableInfos, err = ReadTableInfos(finalOutputPath)
	} else {
		tableInfos, err = ReadTableInfos(tableInfosPath)
	}
	if err != nil {
		panic(fmt.Sprintf("ReadTableInfos failed, err: %+v , file is %s", err, tableInfosPath))
	}
	for index, tableInfo := range tableInfos {
		if tableInfos[index].Answers == nil || len(tableInfo.Answers) != len(tableInfo.Questions) {
			tableInfos[index].Answers = make([]string, len(tableInfo.Questions))
		}
		for qIndex, question := range tableInfo.Questions {
			if tableInfos[index].Answers[qIndex] != "" {
				log.Info("skip answer - question: " + question + " answer: " + tableInfos[index].Answers[qIndex])
				continue
			}
			wg.Add(1)
			time.Sleep(200 * time.Millisecond)
			go ParallelGenDBSql(&tableInfos, index, qIndex, topP, databaseMap)
		}
		wg.Wait()
		outputStr, _ := json.MarshalIndent(tableInfos, "", "\t")
		utils.WriteStringToFile(finalOutputPath, string(outputStr))
	}
	outputStr, err := json.MarshalIndent(tableInfos, "", "\t")
	utils.WriteStringToFile(finalOutputPath, string(outputStr))
}
func main() {
	var configPath string
	var tableInfosPath string
	var outputPath string
	var topP float64
	flag.StringVar(&configPath, "config-path", defaultConfigPath, "DBaaS configuration file path")
	flag.StringVar(&tableInfosPath, "table-infos", defaultTableInfosPath, "DBaaS configuration file path")
	flag.Float64Var(&topP, "top-p", 0.4, "DBaaS configuration file path")
	flag.StringVar(&outputPath, "output", defaultOutputPath, "")
	flag.Parse()
	err := config.InitConfig(configPath)
	if err != nil {
		panic(fmt.Sprintf("Initialize configuration failed, err: %+v", err))
	}
	databaseMap := GenAnswersByTableInfoFile(tableInfosPath, outputPath, topP)
	GenDBAnswersByTableInfoFile(tableInfosPath+"_database", outputPath+"_database", topP, databaseMap)
}

func ParallelGenSql(tableInfos *[]TableInfo, index int, qIndex int, topP float64) {
	if (*tableInfos)[index].Answers[qIndex] != "" {
		log.Info("result has been generated, return " + (*tableInfos)[index].Answers[qIndex])
		wg.Done()
		return
	}
	tableInfo := (*tableInfos)[index]
	question := tableInfo.Questions[qIndex]
	meta := dataapi.TableMeta{Table: tableInfo.Table, Columns: tableInfo.ColNames, Database: tableInfo.Database}
	metas := []*dataapi.TableMeta{&meta}
	answer, err := GenSqlByCompletion(question, metas, topP)
	if err == nil {
		(*tableInfos)[index].Answers[qIndex] = answer
	}
	wg.Done()
}

func ParallelGenDBSql(tableInfos *[]TableInfo, index int, qIndex int, topP float64, databaseMap map[string][]TableInfo) {
	if (*tableInfos)[index].Answers[qIndex] != "" {
		log.Info("result has been generated, return " + (*tableInfos)[index].Answers[qIndex])
		wg.Done()
		return
	}
	tableInfo := (*tableInfos)[index]
	question := tableInfo.Questions[qIndex]
	if _, ok := databaseMap[tableInfo.Database]; !ok {
		log.Info("No table found in databaseMap, return " + tableInfo.Database)
		wg.Done()
		return
	}
	metas := []*dataapi.TableMeta{}
	for _, table := range databaseMap[tableInfo.Database] {
		meta := dataapi.TableMeta{Table: table.Table, Columns: table.ColNames, Database: table.Database}
		metas = append(metas, &meta)
	}
	answer, err := GenSqlByCompletion(question, metas, topP)
	if err == nil {
		(*tableInfos)[index].Answers[qIndex] = answer
	}
	wg.Done()
}

func GenSqlByCompletion(question string, metas []*dataapi.TableMeta, topP float64) (string, error) {
	ctx := context.Background()
	formatPrompt := chatgpttemplate.FormatPromptForSqlEditorDatabase(question, metas[0].Database, metas)
	fmt.Print("prompt " + formatPrompt)
	param := dto.CreateCompletionParam{
		Model:       "text-davinci-003",
		Prompt:      []string{formatPrompt},
		MaxTokens:   512,
		Temperature: 0,
		TopP:        float32(topP),
		N:           1,
		Stream:      false,
		Logprobs:    2,
		Stop:        []string{"#", "---"},
	}
	completion, err := service.CreateCompletionToOpenAI(ctx, &param)
	if err != nil {
		return "", err
	}
	fmt.Print("completion " + completion)
	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
}