package main

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

	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 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 ConvertTableInfoToMeta(tableInfos []TableInfo) []*dataapi.TableMeta {
	metas := []*dataapi.TableMeta{}
	for _, tableInfo := range tableInfos {
		meta := dataapi.TableMeta{
			Database: tableInfo.Database,
			Table:    tableInfo.Table,
			Columns:  tableInfo.ColNames,
		}
		metas = append(metas, &meta)
	}
	return metas
}

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))
	}
	metas := ConvertTableInfoToMeta(tableInfos)
	for index, tableInfo := range tableInfos {
		tableInfos[index].Answers = make([]string, len(tableInfo.Questions))
		for qIndex, question := range tableInfo.Questions {
			answer, err := GuessDBByCompletion(question, metas)
			if err != nil {
				continue
			}
			tableInfos[index].Answers[qIndex] = answer
		}
	}
	outputStr, err := json.MarshalIndent(tableInfos, "", "\t")
	utils.WriteStringToFile(outputPath, string(outputStr))
}

func GuessDBByCompletion(question string, metas []*dataapi.TableMeta) (string, error) {
	ctx := context.Background()
	formatPrompt := chatgpttemplate.FormatPromptForGuessDatabase(question, metas)
	completion, err := service.CreateCompletionTextDavinci(ctx, formatPrompt)
	if err != nil {
		log.Info("err is " + err.Error())
		return "", err
	}
	log.Info("answer is " + 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
}