package main

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

	"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 GenAnswers(index int, tableInfosPtr *[]TableInfo) {
	tableInfos := *tableInfosPtr
	tableInfos[index].Answers = make([]string, len(tableInfos[index].Questions))
	for qIndex, question := range tableInfos[index].Questions {
		answer, err := GenSqlByCompletion(question, tableInfos[index].Table, tableInfos[index].ColNames)
		if err != nil {
			continue
		}
		tableInfos[index].Answers[qIndex] = answer
	}
	wg.Done()
}

func main() {
	var configPath string
	var tableInfosPath string
	var outputPath 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.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))
	}
	for index, _ := range tableInfos {
		wg.Add(1)
		go GenAnswers(index, &tableInfos)
		/*
			tableInfos[index].Answers = make([]string, len(tableInfo.Questions))
			for qIndex, question := range tableInfo.Questions {
				answer, err := GenSqlByCompletion(question, tableInfo.Table, tableInfo.ColNames)
				if err != nil {
					continue
				}
				tableInfos[index].Answers[qIndex] = answer
			}
		*/
	}
	wg.Wait()
	outputStr, err := json.MarshalIndent(tableInfos, "", "\t")
	utils.WriteStringToFile(outputPath, string(outputStr))
}

func GenSqlByCompletion(question string, table string, colNames []string) (string, error) {
	ctx := context.Background()
	meta := dataapi.TableMeta{Table: table, Columns: colNames}
	metas := []*dataapi.TableMeta{&meta}
	formatPrompt := chatgpttemplate.FormatPromptForSqlEditorDatabase(question, metas)
	completion, err := service.CreateCompletionTextDavinci(ctx, formatPrompt)
	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
}