package openaisvc

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/avast/retry-go"
	"github.com/google/uuid"
	jsoniter "github.com/json-iterator/go"
	"github.com/pkoukk/tiktoken-go"
	"github.com/sashabaranov/go-openai"
	"github.com/spf13/viper"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"github.com/tidbcloud/tiinf/sdk/measure"
	"go.uber.org/zap"

	"github.com/tidbcloud/aipower/commons/chatgpttemplate"
	"github.com/tidbcloud/aipower/commons/config"
	"github.com/tidbcloud/aipower/commons/measure_report"
	"github.com/tidbcloud/aipower/commons/metrics"
	"github.com/tidbcloud/aipower/commons/utils"
	"github.com/tidbcloud/aipower/consts"
	"github.com/tidbcloud/aipower/idl/pbgen/proto/api"
	"github.com/tidbcloud/aipower/internal/bean/dto"
	"github.com/tidbcloud/aipower/internal/dao"
	customError "github.com/tidbcloud/aipower/internal/errors"
	"github.com/tidbcloud/aipower/internal/service/cachesvc"
)

var (
	openAIClient *openai.Client
)

func InitOpenAIClient() error {
	conf := config.Global().MessageSvcConf

	openAIClient = openai.NewClient(conf.OpenAIApiKey)

	return nil
}

const (
	prefixStr     = "Answer:"
	extraStr      = "Explanation:"
	groupByClause = "GROUP BY"
	setSQLMode    = "SET SESSION sql_mode=(SELECT REPLACE(@@sql_mode, 'ONLY_FULL_GROUP_BY', ''));"
)

const NoSQLGeneratedErrorMsg = "Unable to generate SQL for this instruction. Please provide a different instruction and try again."
const GetDatabaseColumnsErrorMsg = "Unable to generate SQL due to failure to obtain database schemas. Please try again later."
const NoDatabaseGuessedErrorMsg = "Unable to generate SQL as no database is specified. Please identify database in statement \"use database\"."
const DatabaseEmptyErrorMsg = "Unable to generate SQL as no table in the database. Make sure there is at least one table in the database and try again."
const PromptExceedErrorMsg = "Unable to generate SQL for this database due to its extensive tables."
const InvalidParamsErrorMsg = "Unable to generate SQL due to invalid parameters."
const RequestOpenAIErrorMsg = "Unable to generate SQL due to error from OpenAI. Please try again later."
const defaultUnknownAIErrorMsg = "Unable to generate SQL due to system error."

func createMeasure(req *dto.InteractRequest, cost int64, sku string) error {
	return measure_report.Report.PutMeasure(measure.BillingMeasure{
		TenantID:      req.OrgID,
		CostID:        uuid.NewString(),
		Cost:          cost,
		CostTimestamp: time.Now().Unix(),
		Labels: map[string]string{
			consts.MeasureLabelEndpoint: req.Endpoint,
			consts.MeasureLabelAppID:    fmt.Sprintf("%d", req.AppID),
			consts.MeasureLabelModel:    req.Model,
		},
		Unit: sku,
	})
}

func measureToken(ctx context.Context, fullCompletion string, req *dto.InteractRequest) {
	tkm, err := tiktoken.EncodingForModel(req.Model)
	if err != nil {
		tilog.Error(ctx, "failed to get encoding", zap.String("model", req.Model), zap.Error(err))
		return
	}

	// encode
	token := tkm.Encode(fullCompletion, nil, nil)

	tokenUsed := len(token)

	err = createMeasure(req, int64(tokenUsed), consts.SKUMeasureTokens)
	if err != nil {
		tilog.Error(ctx, "measureRequest  put err ", zap.Error(err), zap.Any("req", req), zap.Int("tokenUsed", tokenUsed))
	}

	metrics.OpenAIEDATokenUsageCounterVec.WithLabelValues(
		fmt.Sprintf("%d", req.OrgID), fmt.Sprintf("%d", req.AppID), req.Endpoint, req.AppKey,
	).Add(float64(tokenUsed))
}

func measureRequest(ctx context.Context, req *dto.InteractRequest) {
	err := createMeasure(req, 1, consts.SKUMeasureRequests)
	if err != nil {
		tilog.Error(ctx, "measureRequest  put err ", zap.Error(err))
	}
	endpoint := "POST - /chat2data"
	metrics.OpenAIEDARequestCounterVec.WithLabelValues(
		fmt.Sprintf("%d", req.OrgID), fmt.Sprintf("%d", req.AppID), endpoint, req.AppKey,
	).Inc()
}

func ConvertToTiErr(err error) error {
	response := err.Error()

	if strings.Contains(response, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
		return errs.NewError(customError.ErrCompletionGenSQLFailed, NoSQLGeneratedErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultNoDatabaseGuessedMsg) {
		return errs.NewError(customError.ErrCompletionGuessDatabaseFailed, NoDatabaseGuessedErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultPromptTooLongMsg) {
		return errs.NewError(customError.ErrCompletionTooManyTables, PromptExceedErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultRequestOpenAIError) {
		return errs.NewError(customError.ErrCompletionRequestOpenAIFailed, RequestOpenAIErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultInvalidParamError) {
		return errs.NewError(customError.ErrCompletionInvalidParams, InvalidParamsErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultGetDatabaseColumnsError) {
		return errs.NewError(customError.ErrCompletionGetDBMetasFailed, GetDatabaseColumnsErrorMsg, err)
	} else if strings.Contains(response, chatgpttemplate.DefaultDatabaseSchemaEmptyError) {
		return errs.NewError(customError.ErrCompletionGetDBMetasEmpty, DatabaseEmptyErrorMsg, err)
	} else {
		return errs.NewInternalGeneralBusiness(defaultUnknownAIErrorMsg, err)
	}
}

// refineSqlCompletion refine the completion sql, it's a generic function for chat2sql/chat2data...
func refineSqlCompletion(ctx context.Context, botType string, req *dto.InteractRequest, completion string, nt []string) string {
	tilog.Info(ctx, "gen sql before optimized", zap.String("completion", completion))
	// when only one table listed in NT and the sql generated used `table.column` format
	if len(nt) == 1 {
		comparedStr := "`" + nt[0] + "`."
		completion = strings.Replace(completion, comparedStr, "", -1)
	}
	// add semicolon
	completion = strings.TrimSpace(strings.Trim(strings.TrimSpace(completion), "\n"))
	if !strings.HasSuffix(completion, ";") {
		completion = completion + ";"
	}
	if strings.Contains(completion, groupByClause) {
		if utils.SetSQLModeEnable(ctx, botType) && req != nil {
			tilog.Info(ctx, "set SQL mode enable", zap.String("botType", botType))
			if !validateSQLWithDryRun(ctx, req, completion) {
				completion = setSQLMode + "\n" + completion
			}
		}
	}
	tilog.Info(ctx, "gen sql after optimized", zap.String("botType", botType), zap.String("completion", completion))
	return completion
}

func getPromptForSQLEditor(ctx context.Context, req *dto.InteractRequest, metas []*dataapi.TableMeta) chatgpttemplate.Prompt {
	if req.PromptTemplate != "" {
		tilog.Info(ctx, "using template from request", zap.Any("template", req.PromptTemplate))
		return chatgpttemplate.FormatCustomizedPromptForSQLEditorDatabase(req, metas)
	}

	template, err := dao.GetLatestPromptByAIModelAndDeploy(consts.AIModelChat2SQL, config.Global().ServerConf.Deployment)
	if err == nil {
		tilog.Info(ctx, "using template from database", zap.Any("template", template))
		return chatgpttemplate.Prompt{
			Content: template.Content,
			Version: template.Version,
		}
	}

	tilog.Info(ctx, "using template from code")
	return chatgpttemplate.FormatPromptForSqlEditorDatabase(req, metas)
}

func createCompletionForSqlEditorDatabase(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var formatPrompt chatgpttemplate.Prompt
	if req.ProjectID == 0 || req.Database == "" {
		tilog.Info(ctx, "invalid params for query2sql", zap.Any("req", req))
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}
	var (
		metas []*dataapi.TableMeta
		err   error
	)
	if utils.GetInformationSchemaEnable(ctx, req.BotType) {
		tilog.Info(ctx, "get information schema enable for query2sql", zap.String("botType", req.BotType))
		metas, err = getDatabaseDetailedMeta(ctx, req)
	} else {
		metas, err = getSqlEditorDatabaseColumns(ctx, req)
	}

	tilog.Info(ctx, "get database metas for query2sql", zap.Any("metas", metas), zap.Any("err", err))
	if err != nil {
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultGetDatabaseColumnsError, err)
	}
	// if it's not "chat2sql", we need to return error
	if len(metas) == 0 {
		tilog.Info(ctx, "database schema is empty", zap.String("botType", req.BotType), zap.String("database", req.Database))
		if req.BotType != consts.AIModelChat2SQL {
			return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultDatabaseSchemaEmptyError, nil)
		}
	}
	formatPrompt = getPromptForSQLEditor(ctx, req, metas)
	// record the status of db schema
	if len(metas) == 0 {
		formatPrompt.SchemaEmpty = true
	}
	// use same params
	completion, result, err := CreateCompletionGeneral(ctx, formatPrompt, getModelName(req), consts.AIModelChat2SQL, req)
	if err != nil {
		tilog.Info(ctx, "query2sql error", zap.String("botType", req.BotType), zap.Error(err))
		return formatPrompt, completion, result, err
	}

	if req.Async {
		return formatPrompt, completion, result, err
	}

	// when we received the response in sync mode, the completion had already been modified
	return formatPrompt, completion, result, err
}

// modifySQLEditorResponse modify the full completion for SQLEditor(chat2sql)
func modifySQLEditorResponse(ctx context.Context, completion string, botType string, req *dto.InteractRequest) (string, error) {
	if len(completion) == 0 {
		return "", errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, nil)
	}
	// pre-check completion
	tilog.Info(ctx, "openai get completion - original", zap.Any("completion", completion))
	// post process
	completion = strings.Replace(completion, "\n", " ", -1)
	// trim prefix string
	if strings.Contains(completion, prefixStr) {
		completion = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(completion), prefixStr))
	}
	// trim extra string
	if index := strings.Index(completion, extraStr); index != -1 {
		completion = strings.TrimSpace(completion[:index])
	}

	// post process
	if index := strings.Index(completion, "\"sql\":"); index != -1 {
		completionTopHalf := strings.TrimSpace(completion[:index])
		completionBottomHalf := strings.TrimSpace(completion[index:])
		if !strings.HasSuffix(completionTopHalf, ",") {
			completionTopHalf += ","
		}
		completion = completionTopHalf + " " + completionBottomHalf
	}

	tilog.Info(ctx, "openai get completion - optimized", zap.Any("completion", completion))

	completionInfo := dto.Chat2QueryCompletion{}
	if unmarshalErr := json.Unmarshal([]byte(completion), &completionInfo); unmarshalErr != nil {
		tilog.Error(ctx, "unmarshal result for chat2sql error", zap.String("completion", completion), zap.Error(unmarshalErr))
		return "", errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, unmarshalErr)
	}
	tilog.Info(ctx, "openai get completion info - final", zap.Any("completionInfo", completionInfo))

	// hot fix: if completionInfo.rewriteSql = "No SQL Generated" and completionInfo.sql != "No SQL Generated", output the result of completionInfo.sql
	if completionInfo.RewriteSql == chatgpttemplate.DefaultNoSqlGeneratedMsg && completionInfo.Sql != chatgpttemplate.DefaultNoSqlGeneratedMsg || completionInfo.RewriteSql == "" {
		completionInfo.RewriteSql = completionInfo.Sql
	}

	if strings.Contains(completionInfo.RewriteSql, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
		tilog.Info(ctx, "No SQL Generated")
		return "", errs.NewBadRequest(chatgpttemplate.DefaultNoSqlGeneratedMsg, nil)
	}

	return refineSqlCompletion(ctx, botType, req, completionInfo.RewriteSql, completionInfo.NT), nil
}

func findDatabaseInMetas(database string, metas []*dataapi.TableMeta) bool {
	for _, meta := range metas {
		if meta.Database == database {
			return true
		}
	}
	return false
}

func getModelName(req *dto.InteractRequest) string {
	if req.Model != "" && (req.Model == consts.OpenAIModelTextDavinci003 || req.Model == consts.OpenAIModelGPT35Turbo || req.Model == consts.OpenAIModelGPT35Turbo16K) {
		return req.Model
	}

	return consts.OpenAIModelTextDavinci003
}

func createCompletionForGuessDatabase(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var prompt chatgpttemplate.Prompt
	if req.ProjectID == 0 {
		tilog.Info(ctx, "invalid params for guessDatabase", zap.Any("req", req))
		return prompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}

	metas, err := getSqlEditorClusterColumns(ctx, req.ProjectID, req.ClusterID, req.Caller)
	tilog.Info(ctx, "get database metas for guessDatabase", zap.Any("metas", metas), zap.Any("err", err))
	if err != nil {
		return prompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultGetDatabaseColumnsError, err)
	}
	if len(metas) == 0 {
		return prompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultDatabaseSchemaEmptyError, nil)
	}
	prompt = chatgpttemplate.FormatPromptForGuessDatabase(req.Question, metas)
	completion, result, err := CreateCompletionGeneral(ctx, prompt, getModelName(req), consts.AIModelGuessDatabase, req)
	// this error need to be distinguished additionally
	if err != nil {
		tilog.Info(ctx, "guess database name error", zap.Error(err))
		return prompt, "", result, errs.NewBadRequest(chatgpttemplate.DefaultNoDatabaseGuessedMsg, err)
	}
	dbName := completion
	completionInfo := dto.GuessDatabaseCompletion{}
	if unmarshalErr := json.Unmarshal([]byte(completion), &completionInfo); unmarshalErr != nil {
		tilog.Error(ctx, "post processing for guessDatabase error", zap.Error(unmarshalErr))
	} else {
		dbName = completionInfo.DBName
	}
	tilog.Info(ctx, "get database name success", zap.String("dbName", dbName))
	if !findDatabaseInMetas(dbName, metas) {
		tilog.Info(ctx, "guess database name failed")
		return prompt, "", result, errs.NewNotFound(chatgpttemplate.DefaultNoDatabaseGuessedMsg, nil)
	}
	return prompt, dbName, result, nil
}

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
		start            = time.Now()
	)

	defer func() {
		duration := time.Since(start).Milliseconds()
		status := "success"
		msg := "nil"
		if err != nil {
			status = "failed"
			msg = fmt.Sprintf("%s", err)
		}
		metrics.OpenAIHistogramVec.WithLabelValues(req.BotType, status, msg).Observe(float64(duration))
		metrics.OpenAICounterVec.WithLabelValues(req.BotType, status, msg).Inc()
	}()

	switch req.BotType {
	case consts.AIModelChat2SQL:
		// force it using gpt-3.5-turbo-16k
		req.Model = consts.OpenAIModelGPT35Turbo16K
		prompt, completion, completionResult, err = CreateCompletionForChat2Sql(ctx, req)
	case consts.AIModelChat2Data:
		// force it using gpt-3.5-turbo-16k
		req.Model = consts.OpenAIModelGPT35Turbo16K
		// not support stream mode, force change
		req.Async = false
		prompt, completion, completionResult, err = CreateCompletionForChat2Sql(ctx, req)
	case consts.AIModelChat2Chart:
		// force it using gpt-3.5-turbo-16k
		req.Model = consts.OpenAIModelGPT35Turbo16K
		// not support stream mode, force change
		req.Async = false
		prompt, completion, completionResult, err = CreateCompletionForChat2Chart(ctx, req)
	case consts.AIModelChat2Question:
		prompt, completion, completionResult, err = CreateCompletionForChat2Question(ctx, req)
	case consts.AIModelColumnMatching:
		prompt, completion, completionResult, err = CreateCompletionForData2ColumnMatching(ctx, req)
	default:
		tilog.Error(ctx, "unknown request ai model", zap.String("model", req.BotType))
		// force it using gpt-3.5-turbo-16k
		req.Model = consts.OpenAIModelGPT35Turbo16K
		prompt, completion, completionResult, err = CreateCompletionForChat2Sql(ctx, req)
	}

	if req.Source == consts.RequestSourceDataService {
		measureRequest(ctx, req)
	}

	return prompt, completion, completionResult, err
}

func CreateCompletionForData2ColumnMatching(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var formatPrompt chatgpttemplate.Prompt
	if req.SampleData == "" {
		tilog.Error(ctx, "invalid arguments for data2ColumnMatching", zap.Any("req", req))
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}
	if req.PromptTemplate != "" {
		formatPrompt = chatgpttemplate.FormatCustomizedPromptForData2ColumnMatching(req)
	} else {
		formatPrompt = chatgpttemplate.FormatPromptForData2ColumnMatching(req.SampleData)
	}
	completion, result, err := CreateCompletionGeneral(ctx, formatPrompt, getModelName(req), consts.AIModelColumnMatching, req)
	if err != nil {
		tilog.Warn(ctx, "data2ColumnMatching error", zap.Error(err))
	}
	return formatPrompt, completion, result, err
}

func CreateCompletionForChat2Question(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var formatPrompt chatgpttemplate.Prompt

	if req.ClusterID == 0 || req.ProjectID == 0 || req.Database == "" || req.QuestionNums == 0 {
		tilog.Info(ctx, "invalid arguments for chat2question", zap.Any("req", req))
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}

	metas, err := getSqlEditorDatabaseColumns(ctx, req)
	tilog.Info(ctx, "get database metas for chat2question", zap.Any("metas", metas), zap.Any("err", err))
	if err != nil {
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultGetDatabaseColumnsError, err)
	}
	if len(metas) == 0 {
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultDatabaseSchemaEmptyError, nil)
	}
	if req.PromptTemplate != "" {
		formatPrompt = chatgpttemplate.FormatCustomizedPromptForChat2Question(req, metas)
	} else {
		formatPrompt = chatgpttemplate.FormatPromptForGenDatabaseQuestions(req, metas, req.QuestionNums)
	}
	completion, result, err := CreateCompletionGeneral(ctx, formatPrompt, getModelName(req), consts.AIModelChat2Question, req)
	if err != nil {
		tilog.Warn(ctx, "chat2Question error", zap.Error(err))
		return formatPrompt, completion, result, err
	}
	tilog.Info(ctx, "completion return by chat2Question", zap.String("completion", completion))
	// Add post-processing, format the completion
	completionInfo := dto.Chat2QuestionCompletion{}
	if unmarshalErr := json.Unmarshal([]byte(completion), &completionInfo); unmarshalErr == nil {
		jsonBytes, marshalErr := json.Marshal(completionInfo.Questions)
		if marshalErr == nil {
			completion = string(jsonBytes)
		} else {
			tilog.Error(ctx, "json marshal for chat2Question error", zap.Error(marshalErr))
		}
	} else {
		tilog.Error(ctx, "json unmarshal for chat2Question error", zap.Error(unmarshalErr))
	}
	return formatPrompt, completion, result, err
}

func getPromptForChat2Chart(ctx context.Context, req *dto.InteractRequest, metas []*dataapi.TableMeta) chatgpttemplate.Prompt {
	if req.PromptTemplate != "" {
		tilog.Info(ctx, "using template from request", zap.String("template", req.PromptTemplate))
		return chatgpttemplate.FormatCustomizedPromptForChat2Chart(req, metas)
	}

	template, err := dao.GetLatestPromptByAIModelAndDeploy(consts.AIModelChat2Chart, config.Global().ServerConf.Deployment)
	if err == nil {
		tilog.Info(ctx, "using template from database", zap.Any("template", template))
		return chatgpttemplate.Prompt{
			Content: template.Content,
			Version: template.Version,
		}
	}

	tilog.Info(ctx, "using template from code")
	return chatgpttemplate.FormatPromptForChat2Chart(req, metas)
}

func CreateCompletionForChat2Chart(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	var (
		formatPrompt     chatgpttemplate.Prompt
		completion       string
		completionResult *dto.CreateCompletionResult
		metas            []*dataapi.TableMeta
		err              error
	)
	if req.ClusterID == 0 || req.ProjectID == 0 || req.Database == "" || req.Question == "" {
		tilog.Info(ctx, "invalid arguments for chat2chart", zap.Any("req", req))
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}

	if utils.GetInformationSchemaEnable(ctx, consts.AIModelChat2Chart) {
		tilog.Info(ctx, "get information schema enable", zap.String("botType", consts.AIModelChat2Chart))
		metas, err = getDatabaseDetailedMeta(ctx, req)
	} else {
		metas, err = getSqlEditorDatabaseColumns(ctx, req)
	}
	tilog.Info(ctx, "get database metas for chat2chart", zap.Any("metas", metas), zap.Any("err", err))
	if err != nil {
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultGetDatabaseColumnsError, err)
	}
	if len(metas) == 0 {
		return formatPrompt, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultDatabaseSchemaEmptyError, nil)
	}
	formatPrompt = getPromptForChat2Chart(ctx, req, metas)
	completion, completionResult, err = CreateCompletionGeneral(ctx, formatPrompt, getModelName(req), consts.AIModelChat2Chart, req)
	if err != nil {
		tilog.Warn(ctx, "chat2chart error", zap.Error(err))
		return formatPrompt, completion, completionResult, err
	}

	completion = modifyChat2ChartResponse(ctx, req, completion)
	return formatPrompt, completion, completionResult, err
}

// modifyChat2ChartResponse modify the full completion for chat2chart
func modifyChat2ChartResponse(ctx context.Context, req *dto.InteractRequest, completion string) string {
	// add post-processing
	completionInfo := dto.Chat2ChartCompletion{}
	if unmarshalErr := json.Unmarshal([]byte(completion), &completionInfo); unmarshalErr == nil {
		completionInfo.Sql = refineSqlCompletion(ctx, req.BotType, req, completionInfo.Sql, nil)
		jsonBytes, marshalErr := json.Marshal(completionInfo)
		if marshalErr == nil {
			completion = string(jsonBytes)
		} else {
			tilog.Error(ctx, "post processing for chat2chart error", zap.Error(marshalErr))
		}
	} else {
		tilog.Error(ctx, "post processing for chat2chart error", zap.Error(unmarshalErr))
	}
	return completion
}

func CreateCompletionForChat2Sql(ctx context.Context, req *dto.InteractRequest) (chatgpttemplate.Prompt, string, *dto.CreateCompletionResult, error) {
	if req.ProjectID == 0 || req.ClusterID == 0 || req.Question == "" {
		tilog.Info(ctx, "invalid params for Chat2Sql", zap.Any("req", req))
		return chatgpttemplate.Prompt{}, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultInvalidParamError, nil)
	}
	if req.Database == "" {
		// Fix: if it is "chat2sql", without database, return error directly, because it is async.
		if req.BotType == consts.AIModelChat2SQL {
			tilog.Info(ctx, "database name is not specified", zap.Any("req", req))
			return chatgpttemplate.Prompt{}, "", nil, errs.NewBadRequest(chatgpttemplate.DefaultNoDatabaseGuessedMsg, nil)
		}
		// Guess Database Name, can not use async mode
		guessDBReq := *req
		guessDBReq.Async = false
		guessDBReq.BotType = consts.AIModelGuessDatabase

		prompt, database, result, err := createCompletionForGuessDatabase(ctx, &guessDBReq)
		if err != nil {
			return prompt, database, result, err
		}
		tilog.Info(ctx, "guess database success", zap.String("database", database))

		// this is very hacky now to change input pointer field value as response value, caller will use this filed after call it
		req.Database = database

		sqlEditorReq := *req
		sqlEditorReq.Database = database
		// SQL Generation
		prompt, completion, result, err := createCompletionForSqlEditorDatabase(ctx, &sqlEditorReq)
		if err != nil {
			return prompt, completion, result, err
		}
		// special case. if completion == "No SQL Generated", It is treated as an error here
		if strings.Contains(completion, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
			tilog.Info(ctx, "No SQL Generated")
			return prompt, completion, result, errs.NewBadRequest(chatgpttemplate.DefaultNoSqlGeneratedMsg, nil)
		}
		return prompt, "use " + database + ";\n" + strings.Trim(completion, "\n"), result, err
	} else { //nolint
		prompt, completion, result, err := createCompletionForSqlEditorDatabase(ctx, req)
		// special case. if completion == "No SQL Generated", It is treated as an error
		if strings.Contains(completion, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
			tilog.Info(ctx, "No SQL Generated")
			return prompt, completion, result, errs.NewBadRequest(chatgpttemplate.DefaultNoSqlGeneratedMsg, nil)
		}
		return prompt, completion, result, err
	}
}

func CreateCompletionGeneral(ctx context.Context, prompt chatgpttemplate.Prompt, model string, aiModel string, req *dto.InteractRequest) (string, *dto.CreateCompletionResult, error) {
	timeout := time.Second * 20
	openAIParams := dto.OpenAIParams{
		Temperature: 0,
		TopP:        0.4,
		N:           1,
		Stream:      false,
		MaxTokens:   512,
	}

	if aiModel == consts.AIModelChat2Question {
		timeout = time.Minute // chat2question timeout should be 60s
		openAIParams = dto.OpenAIParams{
			MaxTokens:   2048,
			Temperature: 1.0,
			TopP:        1.0,
			N:           1,
			Stream:      false,
		}
	} else if aiModel == consts.AIModelColumnMatching {
		timeout = time.Minute // column matching timeout should be 60s
		openAIParams = dto.OpenAIParams{
			MaxTokens:   2048,
			Temperature: 0,
			TopP:        0.4,
			N:           1,
			Stream:      false,
		}
	} else if aiModel == consts.AIModelChat2SQL {
		// Support set timeout with param
		if req.Timeout == 0 {
			timeout = time.Second * time.Duration(viper.GetInt(config.TimeOutForChat2Query))
			tilog.Info(ctx, "time out for chat2query, Set by system", zap.Any("timeout", timeout))
		} else {
			timeout = time.Second * time.Duration(req.Timeout/1000)
			tilog.Info(ctx, "time out for chat2query, Set by user", zap.Any("timeout", timeout))
		}
		openAIParams = dto.OpenAIParams{
			MaxTokens:   768,
			Temperature: 0,
			TopP:        0.4,
			N:           1,
			Stream:      false,
		}
	} else if aiModel == consts.AIModelGuessDatabase {
		timeout = time.Second * 30
		openAIParams = dto.OpenAIParams{
			MaxTokens:   256,
			Temperature: 0,
			TopP:        0.4,
			N:           1,
			Stream:      false,
		}
	} else if aiModel == consts.AIModelChat2Chart {
		timeout = time.Second * 40
	}

	if model == consts.OpenAIModelGPT35Turbo16K {
		openAIParams.Model = openai.GPT3Dot5Turbo16K
		openAIParams.MaxTokens = 4096
		tilog.Info(ctx, "openai create completion params", zap.String("channel", aiModel), zap.Any("params", openAIParams))
		return CreateCompletionToOpenAI(ctx, prompt, openAIParams, timeout, req)
	} else { //nolint
		openAIParams.Model = openai.GPT3TextDavinci003

		tilog.Info(ctx, "openai create completion params", zap.String("channel", aiModel), zap.Any("params", openAIParams))
		return CreateCompletionToOpenAI(ctx, prompt, openAIParams, timeout, req)
	}
}

type OpenAIStreamable[T any] interface {
	Recv() (response T, err error)
	Close()
}

func requestOpenAIInStreamWithGPT35(ctx context.Context, prompt chatgpttemplate.Prompt, params dto.OpenAIParams, jobID string, ch chan dto.CreateCompletionResult) {
	defer close(ch)

	openAIReq := openai.ChatCompletionRequest{
		Model: params.Model,
		Messages: []openai.ChatCompletionMessage{
			{
				Role:    openai.ChatMessageRoleUser,
				Content: prompt.Content,
			},
		},
		MaxTokens:        params.MaxTokens,
		Temperature:      params.Temperature,
		TopP:             params.TopP,
		N:                params.N,
		Stream:           true,
		PresencePenalty:  params.PresencePenalty,
		FrequencyPenalty: params.FrequencyPenalty,
	}

	stream, err := openAIClient.CreateChatCompletionStream(ctx, openAIReq)
	if err != nil {
		tilog.Error(ctx, "failed to create completion in stream", zap.Error(err), zap.Any("params", params))
		ch <- dto.CreateCompletionResult{
			JobID:   jobID,
			Err:     errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, err),
			Object:  "",
			Created: 0,
			Model:   params.Model,
			Content: "",
			Usage:   openai.Usage{},
		}
		return
	}
	defer stream.Close() //nolint

	for {
		response, err := stream.Recv()
		tilog.Info(ctx, "received response from stream", zap.Any("response", response), zap.Error(err))

		var content string
		if len(response.Choices) > 0 {
			content = response.Choices[0].Delta.Content
		}

		if err != nil && !errors.Is(err, io.EOF) {
			err = errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, err)
		}

		ch <- dto.CreateCompletionResult{
			JobID:   jobID,
			Err:     err,
			Object:  response.Object,
			Created: response.Created,
			Model:   response.Model,
			Content: content,
			Usage:   openai.Usage{},
		}
		if err != nil && !errs.IsError(err, io.EOF) {
			tilog.Error(ctx, "failed to receive response from stream", zap.Error(err))
			return
		}
	}
}

func requestOpenAIInStreamWithDavinci(ctx context.Context, prompt chatgpttemplate.Prompt, params dto.OpenAIParams, jobID string, ch chan dto.CreateCompletionResult) {
	defer close(ch)

	openAIReq := openai.CompletionRequest{
		Model:            openai.GPT3TextDavinci003,
		Temperature:      params.Temperature,
		TopP:             params.TopP,
		N:                params.N,
		Stream:           true,
		MaxTokens:        params.MaxTokens,
		Prompt:           prompt.Content,
		PresencePenalty:  params.PresencePenalty,
		FrequencyPenalty: params.FrequencyPenalty,
	}
	stream, err := openAIClient.CreateCompletionStream(ctx, openAIReq)
	if err != nil {
		tilog.Error(ctx, "failed to create completion in stream", zap.Error(err), zap.Any("params", params))
		ch <- dto.CreateCompletionResult{
			JobID: jobID,
			Err:   errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, err),
		}
		return
	}
	defer stream.Close() //nolint

	for {
		response, err := stream.Recv()
		tilog.Info(ctx, "received response from stream", zap.Any("response", response), zap.Error(err))

		var content string
		if len(response.Choices) > 0 {
			content = response.Choices[0].Text
		}

		if err != nil && !errors.Is(err, io.EOF) {
			err = errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, err)
		}

		ch <- dto.CreateCompletionResult{
			JobID:   jobID,
			Err:     err,
			Object:  response.Object,
			Created: response.Created,
			Model:   response.Model,
			Content: content,
			Usage:   response.Usage,
		}
		if err != nil && !errs.IsError(err, io.EOF) {
			tilog.Error(ctx, "failed to receive response from stream", zap.Error(err))
			return
		}
	}
}

func shouldHandleAsHackyQuestion(req *dto.InteractRequest) bool {
	return req.Question == consts.HackyQuestionAIPowerGotInvalidJSON ||
		req.Question == consts.HackyQuestionAIPowerGotBadJSON ||
		req.Question == consts.HackyQuestionTokenOverflow ||
		req.Question == consts.HackyQuestionTimeout
}

func handleHackyQuestion(ctx context.Context, req *dto.InteractRequest, jobID string, respCh chan dto.CreateCompletionResult) {
	defer close(respCh)

	var (
		completion string
		err        error
	)
	if req.Question == consts.HackyQuestionAIPowerGotInvalidJSON {
		completion = `{
		"RQ": "Find top 3 slow queries",
		"NT": ["MEMORY_USAGE_OPS_HISTORY"],
		"NC": {"MEMORY_USAGE_OPS_HISTORY": ["SQL_TEXT"]}
		"sql": "SELECT SQL_TEXT FROM MEMORY_USAGE_OPS_HISTORY ORDER BY TIME DESC LIMIT 3",
		"rewriteSql":` + "\"SELECT `SQL_TEXT` FROM `MEMORY_USAGE_OPS_HISTORY` ORDER BY `TIME` DESC LIMIT 3\"}"
	} else if req.Question == consts.HackyQuestionAIPowerGotBadJSON {
		completion = "hacky: you_can_never_parse_me_as_json"
	} else if req.Question == consts.HackyQuestionTokenOverflow {
		err = errors.New("This model's maximum context length is 16385 tokens. However, your messages resulted in 25507 tokens. Please reduce the length of the messages.")
	} else if req.Question == consts.HackyQuestionTimeout {
		tilog.Info(ctx, "start to sleep 5 minutes, which will cause the request timeout")
		time.Sleep(time.Minute * time.Duration(5))
	}

	respCh <- dto.CreateCompletionResult{
		JobID:   jobID,
		Err:     err,
		Model:   openai.GPT3Dot5Turbo16K,
		Content: completion,
		Usage:   openai.Usage{},
	}
}

func requestOpenAIInStream(
	req *dto.InteractRequest, prompt chatgpttemplate.Prompt, params dto.OpenAIParams, jobID string,
	ch chan dto.CreateCompletionResult, timeout time.Duration,
) {
	// TODO: ctx: value
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	var (
		startAt        = time.Now()
		respCh         = make(chan dto.CreateCompletionResult, 1)
		completionResp *api.QueryCompletionResp
		response       dto.CreateCompletionResult
		fullCompletion string
	)

	defer func() {
		var errMsg = ""

		if req.Source == consts.RequestSourceDataService {
			measureToken(ctx, fullCompletion, req)
		}

		// Get detail errMsg
		if response.Err != nil {
			if newErr, ok := errs.AsTiError(response.Err); ok && newErr.Cause() != nil {
				errMsg = newErr.ErrMsg() + ":" + newErr.Cause().Error()
			} else {
				errMsg = "Request Open API Error: " + response.Err.Error()
			}
		}
		completionResp = postHandleStreamCompletion(ctx, req.BotType, jobID, fullCompletion, prompt.SchemaEmpty, consts.AIPowerStreamStatusDone)

		endAt := time.Now()
		saveRequestRecord(ctx, req, prompt, completionResp.Completion, req.Caller, errMsg, endAt.Sub(startAt).Milliseconds())

		response.Content = completionResp.Completion
		ch <- response
		tilog.Info(ctx, "stream closed")
		close(ch)
	}()

	if shouldHandleAsHackyQuestion(req) {
		tilog.Info(ctx, "handled as hacky question", zap.Any("req", req), zap.Any("prompt", prompt), zap.Any("params", params), zap.String("jobID", jobID))
		go handleHackyQuestion(ctx, req, jobID, respCh)
	} else {
		tilog.Info(ctx, "start to request openai in stream mode", zap.Any("req", req), zap.Any("prompt", prompt), zap.Any("params", params), zap.String("jobID", jobID))
		if req.Model == openai.GPT3Dot5Turbo16K {
			go requestOpenAIInStreamWithGPT35(ctx, prompt, params, jobID, respCh)
		} else {
			go requestOpenAIInStreamWithDavinci(ctx, prompt, params, jobID, respCh)
		}
	}

	var (
		ok bool
	)
	for {
		select {
		case <-ctx.Done():
			tilog.Info(ctx, "stop to generate openai stream as time is out")
			msg := fmt.Sprintf("timeout is reatched, botType:%s, timeout: %d", req.BotType, timeout)
			response.Err = errs.NewInternalTimeout(msg, nil)
			return
		case response, ok = <-respCh:
			if !ok {
				tilog.Info(ctx, "stop to generate openai stream as channel closed")
				return
			}
		}

		if errors.Is(response.Err, io.EOF) {
			tilog.Info(ctx, "stream finished as received EOF")
			response.Err = nil // reset the error
			return
		}

		// TODO: 若AI stream中途error，这里直接返回，会不会导致前端只展示部分结果，是否需要在这里将fullCompletion清空
		if response.Err != nil {
			tilog.Error(ctx, "request openai in stream got error", zap.Error(response.Err), zap.Any("params", params))
			return
		}

		tilog.Info(ctx, "received stream response", zap.String("question", req.Question), zap.Any("response", response))
		fullCompletion += response.Content
		completionResp = postHandleStreamCompletion(ctx, req.BotType, jobID, fullCompletion, prompt.SchemaEmpty, consts.AIPowerStreamStatusRunning)
		tilog.Info(ctx, "received full completion", zap.String("question", req.Question), zap.Any("fullCompletion", fullCompletion))
	}
}

// postHandleStreamCompletion is a post handler for stream response, which take care of parsing streaming response
func postHandleStreamCompletion(ctx context.Context, botType string, jobID string, fullCompletion string, schemaEmpty bool, status int32) *api.QueryCompletionResp {
	var (
		completion string
		message    string
		err        error
	)

	switch botType {
	case consts.AIModelChat2Data:
		// chat2data use a sync way, we just do nothing here, pass the result as is.
		if status == consts.AIPowerStreamStatusDone {
			tilog.Info(ctx, "stream is done, start to parse & modify the full completion")
			completion, err = modifySQLEditorResponse(ctx, fullCompletion, botType, nil)
			if err != nil {
				tilog.Info(ctx, "failed to modify sql editor response", zap.String("fullCompletion", fullCompletion), zap.Error(err))
				message = defaultUnknownAIErrorMsg
			}
		}
	case consts.AIModelChat2SQL:
		var result dto.Chat2QueryCompletion
		iter := jsoniter.ParseString(jsoniter.ConfigDefault, fullCompletion+"\"")
		iter.ReadVal(&result)
		completion = result.RewriteSql
		if iter.Error != nil {
			message = defaultUnknownAIErrorMsg
			tilog.Info(ctx, "failed to parse Chat2SQL json", zap.String("fullCompletion", fullCompletion), zap.Error(iter.Error))
		}

		if status == consts.AIPowerStreamStatusDone {
			tilog.Info(ctx, "stream is done, start to parse & modify the full completion")
			completion, err = modifySQLEditorResponse(ctx, fullCompletion, botType, nil)
			if schemaEmpty && !strings.HasPrefix(strings.ToUpper(completion), "CREATE") {
				tilog.Info(ctx, "db schema is empty but gen SQL does not start with 'CREATE'", zap.String("botType", botType), zap.String("completion", completion))
				message = DatabaseEmptyErrorMsg
				completion = ""
			}
			if err != nil {
				tilog.Info(ctx, "failed to modify sql editor response", zap.String("fullCompletion", fullCompletion), zap.Error(err))
				message = defaultUnknownAIErrorMsg
			}
		}
	case consts.AIModelChat2Question:
		var result dto.Chat2QuestionCompletion
		iter := jsoniter.ParseString(jsoniter.ConfigDefault, fullCompletion+"\"")
		iter.ReadVal(&result)
		if iter.Error != nil {
			message = defaultUnknownAIErrorMsg
			tilog.Info(ctx, "failed to parse Chat2Question json", zap.String("fullCompletion", fullCompletion), zap.Error(iter.Error))
		}

		var data []byte
		data, err = json.Marshal(result.Questions)
		if err != nil {
			message = defaultUnknownAIErrorMsg
			tilog.Info(ctx, "failed to marshal Chat2Question json", zap.Error(err), zap.Any("result", result))
		}
		completion = string(data)
	case consts.AIModelColumnMatching:
		var result []dto.ColumnDescription
		iter := jsoniter.ParseString(jsoniter.ConfigDefault, fullCompletion+"\"")
		iter.ReadVal(&result)
		if iter.Error != nil {
			message = defaultUnknownAIErrorMsg
			tilog.Error(ctx, "failed to parse ColumnMatching json", zap.String("fullCompletion", fullCompletion), zap.Error(iter.Error))
		}
		var data []byte
		data, err = json.Marshal(result)
		if err != nil {
			message = defaultUnknownAIErrorMsg
			tilog.Error(ctx, "failed to marshal ColumnMatching json", zap.Error(err), zap.Any("result", result))
		}
		completion = string(data)
	case consts.AIModelGuessDatabase:
		completion = fullCompletion
	case consts.AIModelChat2Chart:
		completion = fullCompletion
	default:
		completion = fullCompletion
	}

	if status != consts.AIPowerStreamStatusDone {
		message = ""
	}

	if strings.Contains(completion, chatgpttemplate.DefaultNoSqlGeneratedMsg) {
		tilog.Info(ctx, "No SQL Generated", zap.String("jobID", jobID), zap.String("fullCompletion", fullCompletion))
		message = NoSQLGeneratedErrorMsg
	}

	completionResp := &api.QueryCompletionResp{
		Completion: completion,
		Status:     status,
		Message:    message,
	}

	cachesvc.SaveCompletionResult(ctx, jobID, completionResp)

	return completionResp
}

func requestOpenAI(ctx context.Context, req *dto.InteractRequest, prompt chatgpttemplate.Prompt, params dto.OpenAIParams) dto.CreateCompletionResult {
	if req.Model == openai.GPT3Dot5Turbo {
		openAIReq := openai.ChatCompletionRequest{
			Model: openai.GPT3Dot5Turbo,
			Messages: []openai.ChatCompletionMessage{
				{
					Role:    openai.ChatMessageRoleUser,
					Content: prompt.Content,
				},
			},
			MaxTokens:        params.MaxTokens,
			Temperature:      params.Temperature,
			TopP:             params.TopP,
			N:                params.N,
			Stream:           false,
			PresencePenalty:  params.PresencePenalty,
			FrequencyPenalty: params.FrequencyPenalty,
		}

		resp, err := openAIClient.CreateChatCompletion(ctx, openAIReq)

		var content string
		if len(resp.Choices) > 0 {
			content = resp.Choices[0].Message.Content
		}

		return dto.CreateCompletionResult{
			Err:     err,
			Object:  resp.Object,
			Created: resp.Created,
			Model:   resp.Model,
			Content: content,
			Usage:   resp.Usage,
		}
	}

	openAIReq := openai.CompletionRequest{
		Model:            openai.GPT3TextDavinci003,
		Temperature:      params.Temperature,
		TopP:             params.TopP,
		N:                params.N,
		Stream:           false,
		MaxTokens:        params.MaxTokens,
		Prompt:           prompt.Content,
		PresencePenalty:  params.PresencePenalty,
		FrequencyPenalty: params.FrequencyPenalty,
	}
	resp, err := openAIClient.CreateCompletion(ctx, openAIReq)

	var content string
	if len(resp.Choices) > 0 {
		content = resp.Choices[0].Text
	}

	return dto.CreateCompletionResult{
		Err:     err,
		Object:  resp.Object,
		Created: resp.Created,
		Model:   resp.Model,
		Content: content,
		Usage:   resp.Usage,
	}
}

func RequestOpenAIInStream(ctx context.Context, req *dto.InteractRequest, prompt chatgpttemplate.Prompt, params dto.OpenAIParams, jobID string, timeout time.Duration) (chan dto.CreateCompletionResult, error) {
	conf := config.Global().ServerConf

	ch := make(chan dto.CreateCompletionResult, 1)

	if conf.EnableStreamMode {
		// TODO: ctx未传入，导致log缺少traceID
		go requestOpenAIInStream(req, prompt, params, jobID, ch, timeout)
	} else {
		// this will and should only be used in AI-Platform, as it can't set a timeout when request openai
		ch <- requestOpenAI(ctx, req, prompt, params)
	}

	return ch, nil
}

func genJobID() string {
	return uuid.New().String()
}

func CreateCompletionToOpenAI(ctx context.Context, prompt chatgpttemplate.Prompt, params dto.OpenAIParams, timeout time.Duration, req *dto.InteractRequest) (string, *dto.CreateCompletionResult, error) {
	tilog.Info(ctx, "openai create completion prompt", zap.Any("prompt", prompt), zap.Any("params", params), zap.Float64("timeout", timeout.Seconds()))
	conf := config.Global().MessageSvcConf
	var result dto.CreateCompletionResult
	start := time.Now()
	jobID := genJobID()

	err := retry.Do(func() error {
		ch, err := RequestOpenAIInStream(ctx, req, prompt, params, jobID, timeout)
		if err != nil {
			tilog.Info(ctx, "openai create completion error", zap.Error(err))
			return errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, err)
		}

		if !req.Async {
			result = <-ch
			if result.Err != nil {
				tilog.Error(ctx, "openai create completion fail", zap.Error(result.Err))
				return errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, result.Err)
			}
			if result.Content == "" {
				tilog.Info(ctx, "openai create completion get result choices empty", zap.Any("result", result))
				return errs.NewBadRequest(chatgpttemplate.DefaultRequestOpenAIError, nil)
			}
		} else {
			// TODO stream error 是否有触发重试
			result.JobID = jobID
		}

		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.Content

	tilog.Info(ctx, "openai create completion get completion", zap.Any("completion", completion))
	// trim prefix string
	if strings.Contains(completion, prefixStr) {
		completion = strings.TrimLeft(completion, "\n")
		completion = strings.TrimSpace(strings.TrimPrefix(completion, prefixStr))
	}
	// trim extra string
	if index := strings.Index(completion, extraStr); index != -1 {
		completion = strings.Trim(completion[:index], "\n")
	}
	tilog.Info(ctx, "final completion", zap.Any("completion", completion))
	return completion, &result, nil
}