package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"message-service/commons/chatgpttemplate"
	"message-service/commons/config"
	"message-service/commons/grpc"
	"message-service/commons/utils"
	"message-service/internal/bean/dto"
	"os"
	"strings"

	service "message-service/internal/service/central"

	"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"
)

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 ReadTableInfos(tableInfoFile string) ([]TableInfo, error) {
	content, _ := utils.ReadFileToString(tableInfoFile)
	contentBytes := []byte(content)
	fmt.Println(content)
	tableInfos := make([]TableInfo, 0)
	err := json.Unmarshal(contentBytes, &tableInfos)
	if err != nil {
		return nil, err
	}
	return tableInfos, nil
}

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, tableInfo := range tableInfos {
		questionStr, err := GenQuestionsByCompletion(tableInfo)
		if err != nil {
			continue
		}
		questionList := strings.Split(questionStr, "\n")
		tableInfos[index].Table = strings.Replace(tableInfos[index].Table, "\n", " ", -1)
		tableInfos[index].Table = strings.Replace(tableInfos[index].Table, "\t", " ", -1)
		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))
}
func GenQuestionsByCompletion(tableInfo TableInfo) (string, error) {
	ctx := context.Background()
	meta := dataapi.TableMeta{
		Database: tableInfo.Database,
		Table:    tableInfo.Table,
		Columns:  tableInfo.ColName,
	}
	prompt := chatgpttemplate.FormatPromptForGenQuestions(&meta, 20)
	param := dto.CreateCompletionParam{
		Model:       "text-davinci-003",
		Prompt:      []string{prompt},
		MaxTokens:   512,
		Temperature: 0,
		TopP:        0.4,
		N:           1,
		Stream:      false,
		Logprobs:    2,
		//	Stop:        []string{"#", "\n", "---", "###"},
		Stop: []string{"#", "---"},
	}
	completion, err := service.CreateCompletionToOpenAI(ctx, &param)
	if err != nil {
		return "", err
	}
	return completion, nil
}

func getClusterMeta(ctx context.Context, projectID uint64, clusterID uint64) ([]*dataapi.TableMeta, error) {
	req := &dataapi.GetClusterMetaReq{
		ProjectId: projectID,
		ClusterId: clusterID,
	}
	resp, err := grpc.GetClusterMeta(ctx, req)
	if err != nil {
		return nil, err
	}
	return resp.GetMetas(), nil
}