package main

import (
	"context"
	"flag"
	"fmt"
	"message-service/commons/chatgpttemplate"
	"message-service/commons/config"
	"message-service/commons/utils"
	"message-service/internal/bean/dto"
	service "message-service/internal/service/central"
	"strconv"
	"strings"

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

const (
	defaultConfigPath   = "/etc/dbaas/config.yaml"
	defaultOutputPath   = "./question_answer.list"
	defaultQuestionPath = "./question.list"
)

func testGenQuestion(database string, metas []*dataapi.TableMeta, questionNum int) (string, error) {
	ctx := context.Background()
	prompt := `#Question Generator for Data Exploring Purpose

` + chatgpttemplate.FormatTableInstructionString(database, metas) + `

###Task:
Generate ` + strconv.Itoa(questionNum) + ` interesting questions about companies' ranking change based on the following MySQL database schema:`
	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 main() {
	var configPath string
	var outputPath string
	var questionPath string
	flag.StringVar(&configPath, "config-path", defaultConfigPath, "DBaaS configuration file path")
	flag.StringVar(&outputPath, "output", defaultOutputPath, "")
	flag.StringVar(&questionPath, "question", defaultQuestionPath, "")
	flag.Parse()
	err := config.InitConfig(configPath)
	if err != nil {
		panic(fmt.Sprintf("Initialize configuration failed, err: %+v", err))
	}
	ctx := context.Background()
	database := "sample_data"
	var meta = dataapi.TableMeta{
		Database: database,
		Table:    "global_fortune_500_2018_2022",
		Columns:  strings.Split("rank,company_name,country,employees_num,previous_rank,revenues_million,revene_change_percentage,profit_million,profit_change_percentage,asset_million,year", ","),
	}
	var metas = []*dataapi.TableMeta{&meta}
	questionList := []string{}
	if utils.PathExists(questionPath) {
		questionListPtr, _ := utils.ReadFileToList(questionPath)
		questionList = *questionListPtr
	} else {
		questListStr, err := testGenQuestion(database, metas, 30)
		if err != nil {
			panic(fmt.Sprintf("Gen Question failed, err: %+v", err))
		}
		questionList = strings.Split(questListStr, "\n")
		utils.WriteListToFile(questionPath, questionList)
	}

	dataLines := []string{}
	for _, question := range questionList {
		if question == "" {
			continue
		}
		formatPrompt := chatgpttemplate.FormatPromptForSqlEditorDatabase(question, metas[0].Database, metas)
		completion, err := service.CreateCompletionTextDavinci(ctx, formatPrompt)
		line := "=========== FormatPrompt: ===============\n"
		line += formatPrompt + "\n"
		line += "=========== Answer: ===============\n"
		if err != nil {
			line += "Gen Answer failed: " + err.Error() + "\n"
		} else {
			line += completion + "\n"
		}
		print(line)
		dataLines = append(dataLines, line)
	}
	utils.WriteListToFile(outputPath, dataLines)
}