package central

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"message-service/commons/chatgpttemplate"
	"message-service/commons/config"
	"message-service/commons/httpclient"
	"message-service/internal/bean/dto"

	"github.com/avast/retry-go"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
	"github.com/tidbcloud/tiinf/errs"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

const (
	createCompletionPath   = "/v1/completions"
	authorizationFormatter = "Bearer %s"
)

const (
	COMPLETION_CODE_OK                     = 0
	COMPLETION_CODE_GEN_SQL_FAILED         = 1
	COMPLETION_CODE_GUESS_DATABASE_FAILED  = 2
	COMPLETION_CODE_TOO_MANY_TABLES_FAILED = 3
	COMPLETION_CODE_UNKNOWN_FAILED         = 255
)
const defaultUnknownAIErrorMsg = "We couldn't generate SQL for this instruction, please try another one."

func GetRespByCompletion(templateType string, response string) (string, int) {
	if strings.Contains(response, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
		return "We couldn't generate SQL for this instruction, please try another one.", COMPLETION_CODE_GEN_SQL_FAILED
	} else if strings.Contains(response, chatgpttemplate.DefaultNoDatabaseGuessedMsg) {
		return "Please type \"use {database};\" to specify a database first.", COMPLETION_CODE_GUESS_DATABASE_FAILED
	} else if strings.Contains(response, chatgpttemplate.DefaultPromptTooLongMsg) {
		return "We couldn't generate SQL for this database so far, because it has too many tables.", COMPLETION_CODE_TOO_MANY_TABLES_FAILED
	}
	return response, COMPLETION_CODE_OK
}

func IsPromptTooLong(formatPrompt string) bool {
	return len(formatPrompt) > 4096*3
}

func modifySqlCompletion(ctx context.Context, completion string, metas []*dataapi.TableMeta) string {
	for _, meta := range metas {
		databaseTableName := meta.Database + "." + meta.GetTable()
		comparedStr := "`" + databaseTableName + "`"
		if strings.Contains(completion, comparedStr) {
			tilog.Info(ctx, "modify completion: "+completion+" replace databasename: "+comparedStr)
			completion = strings.Replace(completion, comparedStr, databaseTableName, -1)
		}
	}
	return completion
}
func createCompletionForSqlEditorDatabase(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var formatPrompt chatgpttemplate.Prompt
	// if req.BotType != "sqleditor_database" || req.ProjectID == 0 || req.Database == "" {
	if req.ProjectID == 0 || req.Database == "" {
		tilog.Warn(ctx, "createCompletionForSqlEditorDatabase projectID or database is empty", zap.String("database", req.Database), zap.Any("projectID", req.ProjectID))
		return formatPrompt, "", nil, errs.NewBadRequest("invalid params for sqleditor_database.", nil)
	}
	clusterId, err := strconv.ParseUint(req.ClusterID, 10, 64)
	if err != nil {
		return formatPrompt, "", nil, errs.NewBadRequest("invalid params for sqleditor_database, parse clusterId failed.", nil)
	}

	metas, err := getSqlEditorDatabaseColumns(ctx, req.ProjectID, clusterId, req.Database)
	tilog.Info(ctx, "getSqlEditorDatabaseColumns result", zap.Any("metas", metas), zap.Any("err", err))
	if err != nil {
		return formatPrompt, "", nil, err
	}

	if req.PromptTemplate != "" {
		formatPrompt = chatgpttemplate.FormatCustomizedPromptForSQLEditorDatabase(req.PromptTemplate, req.PromptTemplateVersion, req.Question, req.Database, metas)
	} else {
		formatPrompt = chatgpttemplate.FormatPromptForSqlEditorDatabase(req.Question, req.Database, metas)
	}

	if IsPromptTooLong(formatPrompt.Content) {
		return formatPrompt, chatgpttemplate.DefaultPromptTooLongMsg, nil, nil
	}
	completion, result, err := CreateCompletionTextDavinci(ctx, formatPrompt.Content)
	return formatPrompt, modifySqlCompletion(ctx, completion, metas), result, err
}

func findDatabaseInMetas(database string, metas []*dataapi.TableMeta) bool {
	for _, meta := range metas {
		if meta.Database == database {
			return true
		}
	}
	return false
}
func createCompletionForGuessDatabase(ctx context.Context, req *dto.InteractRequest) (string, *dto.CreateCompletionResult, error) {
	// if req.BotType != "sqleditor_no_database" || req.ProjectID == 0 {
	if req.ProjectID == 0 {
		return "", nil, errs.NewBadRequest("invalid params for sqleditor_database.", nil)
	}
	clusterId, err := strconv.ParseUint(req.ClusterID, 10, 64)
	if err != nil {
		return "", nil, errs.NewBadRequest("invalid params for sqleditor_database, parse clusterId failed.", nil)
	}

	metas, err := getSqlEditorClusterColumns(ctx, req.ProjectID, clusterId)
	if err != nil {
		return "", nil, err
	}
	formatPrompt := chatgpttemplate.FormatPromptForGuessDatabase(req.Question, metas)
	resp, result, err := CreateCompletionTextDavinci(ctx, formatPrompt)
	if err != nil {
		return "", result, err
	}
	database := strings.ReplaceAll(resp, "\n\nAnswer: ", "")
	if !findDatabaseInMetas(database, metas) {
		return chatgpttemplate.DefaultNoDatabaseGuessedMsg, result, nil
	}
	return database, result, err
}
func CreateCompletionWithReq(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var (
		prompt           chatgpttemplate.Prompt
		completion       string
		completionResult *dto.CreateCompletionResult
		err              error
	)
	switch req.BotType {
	case "chat2sql":
		prompt, completion, completionResult, err = CreateCompletionForChat2Sql(ctx, req)
	case "chat2chart":
		prompt, completion, completionResult, err = CreateCompletionForChat2Chart(ctx, req)
	default:
		prompt, completion, completionResult, err = CreateCompletionForChat2Sql(ctx, req)
	}
	if prompt.Version != "" {
		if req.BotType == "" {
			prompt.Version = "chat2sql-" + prompt.Version
		} else {
			prompt.Version = req.BotType + "-" + prompt.Version
		}
	}
	return prompt, completion, completionResult, err
}

func CreateCompletionForChat2Chart(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var (
		prompt           chatgpttemplate.Prompt
		completion       string
		completionResult *dto.CreateCompletionResult
		err              error
	)
	// if req.BotType != "sqleditor_database" || req.ProjectID == 0 || req.Database == "" {
	if req.ProjectID == 0 || req.Database == "" {
		return prompt, "", nil, errs.NewBadRequest("invalid params for sqleditor_database.", nil)
	}
	clusterId, err := strconv.ParseUint(req.ClusterID, 10, 64)
	if err != nil {
		return prompt, "", nil, errs.NewBadRequest("invalid params for sqleditor_database, parse clusterId failed.", nil)
	}

	metas, err := getSqlEditorDatabaseColumns(ctx, req.ProjectID, clusterId, req.Database)
	if err != nil {
		return prompt, "", nil, err
	}
	prompt = chatgpttemplate.FormatPromptForChat2Chart(req.Question, req.Database, metas)
	if IsPromptTooLong(prompt.Content) {
		return prompt, chatgpttemplate.DefaultPromptTooLongMsg, nil, nil
	}
	completion, completionResult, err = CreateCompletionTextDavinci(ctx, prompt.Content)
	return prompt, completion, completionResult, err
}
func CreateCompletionForChat2Sql(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	if req.Database == "" {
		var prompt chatgpttemplate.Prompt
		//		return "", chatgpttemplate.DefaultNoDatabaseGuessedMsg, nil
		database, result, err := createCompletionForGuessDatabase(ctx, req)
		if err != nil {
			return prompt, database, nil, err
		}
		if database == chatgpttemplate.DefaultNoDatabaseGuessedMsg {
			return prompt, database, nil, nil
		}
		tilog.Info(ctx, "guessdatabase : "+database)
		sqlEditorReq := &dto.InteractRequest{
			Question:  req.Question,
			BotType:   "sqleditor_database",
			ProjectID: req.ProjectID,
			ClusterID: req.ClusterID,
			Database:  database,
		}
		prompt, completion, result, err := createCompletionForSqlEditorDatabase(ctx, sqlEditorReq)
		if err != nil || completion == chatgpttemplate.DefaultNoSqlGeneratedMsg {
			return prompt, completion, result, err
		}
		return prompt, "use " + database + ";\n" + strings.Trim(completion, "\n"), result, err
	} else {
		return createCompletionForSqlEditorDatabase(ctx, req)
	}
}

func CreateCompletionTextDavinci(ctx context.Context, prompt string) (string, *dto.CreateCompletionResult, error) {
	tilog.Info(ctx, "openai create completion request", zap.String("prompt", prompt))
	// botType: text-davinci-003
	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{"#", "---"},
	}
	tilog.Info(ctx, "openai create completion prompt", zap.Any("prompt", prompt))
	return CreateCompletionToOpenAI(ctx, &param)
}
func CreateCompletionToOpenAI(ctx context.Context, param *dto.CreateCompletionParam) (string, *dto.CreateCompletionResult, error) {
	conf := config.Global().MessageSvcConf

	headers := map[string]string{
		"Content-Type":  "application/json",
		"Authorization": fmt.Sprintf(authorizationFormatter, conf.OpenAIApiKey),
	}

	url := conf.OpenAIEndpoint + createCompletionPath

	var result dto.CreateCompletionResult

	start := time.Now()

	err := retry.Do(func() error {
		body, code, err := httpclient.PostWithHeader(url, param, &headers)

		tilog.Info(ctx, "openai create completion response", zap.Int("'code", code), zap.String("body", body), zap.Error(err))

		if err != nil || code != http.StatusOK {
			tilog.Error(ctx, "openai create completion error", zap.Int("'code", code), zap.Error(err))
			return errs.NewInternalHTTP("request openai failed. ", err)
		}

		err = json.Unmarshal([]byte(body), &result)
		if err != nil {
			tilog.Error(ctx, "openai create completion unmarshal body error", zap.String("body", body), zap.Error(err))
			return err
		}
		if len(result.Choices) == 0 {
			tilog.Info(ctx, "openai create completion get result choices empty", zap.Any("result", result))
			return errs.NewInternalServerError("openai create completion get result choices empty", nil)
		}

		return nil
	}, retry.MaxDelay(time.Duration(conf.MaxDelaySeconds)), retry.Attempts(uint(conf.MaxRetryCount)), retry.LastErrorOnly(true), retry.RetryIf(
		func(err error) bool {
			if err == nil {
				return false
			}
			return time.Since(start) <= time.Second*time.Duration(conf.MaxDelaySeconds)
		}),
	)

	if err != nil {
		tilog.Error(ctx, "openai fail create completion at last", zap.Error(err))
		return "", &result, err
	}

	completion := result.Choices[0].Text
	tilog.Info(ctx, "openai create completion get completion", zap.Any("completion", completion))
	return completion, &result, nil
}