// Package service add by blacknc <blacknc@163.com> at 2023/03/13 16:26:06
package service

import (
	"context"
	"dataservices/commons/config"
	"dataservices/commons/constants"
	"dataservices/commons/grpc"
	"dataservices/commons/http"
	"dataservices/commons/utils"
	dto "dataservices/internal/dto/dataapi"
	"dataservices/internal/entity"
	"errors"
	"strconv"
	"strings"
	"time"

	"github.com/tidbcloud/aipower/idl/pbgen/proto/api"
	"github.com/tidbcloud/dataapps/idl/pbgen/proto/sqleditor"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

const (
	chatBotOKCode     = 0
	chatBotRetryCount = 1
	execSqlRetryCount = 0
)

func Chat2DataGenSQL(ctx context.Context, param *entity.Chat2QueryParam) (*dto.InteractResponse, error) {
	tilog.Info(ctx, "service chat2data start", zap.Any("param", param))
	startT := time.Now()
	resp, err := ChatBotWithParam(ctx, "chat2data", param)
	tc := time.Since(startT)
	tilog.Info(ctx, "service chat2data gen sql", zap.Any("cost", tc.String()), zap.Any("resp", resp), zap.Error(err))
	return resp, err
}

func ChatBotWithParam(ctx context.Context, botType string, param *entity.Chat2QueryParam) (*dto.InteractResponse, error) {
	req := &dto.InteractRequest{
		BotType:               botType,
		OrgID:                 param.OrgID,
		ProjectID:             param.ProjectID,
		ClusterID:             strconv.FormatUint(param.ClusterID, 10),
		Database:              param.Database,
		Question:              param.Instruction,
		PromptTemplate:        param.PromptTemplate,
		PromptTemplateVersion: param.PromptTemplateVersion,
		QuestionNum:           param.QuestionNum,
		Tables:                param.Tables,
		SampleData:            param.SampleData,
		DatabaseSchema:        param.DatabaseSchema,
		Timeout:               param.Timeout,
		AppID:                 param.AppID,
		AppKey:                param.Apikey,
		Endpoint:              param.Endpoint,
	}
	rsp, err := InteractBot(ctx, req, param.Apikey)
	if err != nil {
		tilog.Error(ctx, "call InteractBot failed", zap.Any("req", req), zap.Any("err", err))
		return nil, err
	}
	if rsp.Response == "" {
		tilog.Warn(ctx, "call InteractBot fail, completion is empty", zap.Any("req", req), zap.Any("rsp", rsp))
		return rsp, errs.NewInternalGeneralBusiness("The sql generated by ai is empty", nil)
	}

	tilog.Info(ctx, "call InteractBot succeed", zap.Any("req", req), zap.Any("rsp", rsp))
	return rsp, nil
}

type BotResponse interface {
	GetCompletion() string
	GetSessionId() string
	GetDatabase() string
	GetPrompt() string
	GetPromptVersion() string
	GetTokenUsed() int32
}

func chatBotGrpc(ctx context.Context, req *dto.InteractRequest, createdBy string) (BotResponse, error) {
	clusterID, _ := strconv.ParseUint(req.ClusterID, 10, 64)
	botReq := &api.ChatBotReq{
		BotType:               req.BotType,
		ProjectId:             req.ProjectID,
		ClusterId:             clusterID,
		Database:              req.Database,
		Question:              req.Question,
		Caller:                createdBy,
		PromptTemplate:        req.PromptTemplate,
		PromptTemplateVersion: req.PromptTemplateVersion,
		QuestionNums:          int32(req.QuestionNum),
		Tables:                req.Tables,
		SampleData:            req.SampleData,
		DatabaseSchema:        req.DatabaseSchema,
		OrgId:                 req.OrgID,
		UserType:              "apikey",
		UserId:                createdBy,
		AppId:                 req.AppID,
		AppKey:                req.AppKey,
		Endpoint:              req.Endpoint,
		Source:                constants.AIPowerSourceDataService,
	}
	botRsp, botErr := grpc.ChatBot(ctx, botReq)
	if botErr != nil && errors.Is(botErr, errs.ErrInternalGRPC) {
		for i := 0; i < chatBotRetryCount; i++ {
			botRsp, botErr = grpc.ChatBot(ctx, botReq)
			if botErr == nil || !errors.Is(botErr, errs.ErrInternalGRPC) {
				break
			}
		}
	}
	return botRsp, botErr
}

func chatBotHttp(ctx context.Context, req *dto.InteractRequest, createdBy string) (BotResponse, error) {
	clusterID, _ := strconv.ParseUint(req.ClusterID, 10, 64)
	return http.ChatBot(ctx, req.OrgID, req.ProjectID, &http.ChatBotReq{
		Question:              req.Question,
		BotType:               req.BotType,
		ClusterID:             clusterID,
		Database:              req.Database,
		PromptTemplate:        req.PromptTemplate,
		PromptTemplateVersion: req.PromptTemplateVersion,
		Tables:                req.Tables,
		Caller:                createdBy,
		Timeout:               req.Timeout,
		OrgID:                 req.OrgID,
		AppID:                 req.AppID,
		AppKey:                req.AppKey,
		Endpoint:              req.Endpoint,
		Source:                constants.AIPowerSourceDataService,
	})
}

func InteractBot(ctx context.Context, req *dto.InteractRequest, createdBy string) (*dto.InteractResponse, error) {
	var (
		startTime    = time.Now()
		clusterID, _ = strconv.ParseUint(req.ClusterID, 10, 64)
		botErr       error
		botRsp       BotResponse
	)
	if !config.IsCentralMode() {
		botRsp, botErr = chatBotHttp(ctx, req, createdBy)
	} else {
		botRsp, botErr = chatBotGrpc(ctx, req, createdBy)
	}

	var (
		errorMessage string
		status       int8 = 1
		endTime           = time.Now()
		executeTime       = endTime.Sub(startTime).Milliseconds()
	)

	if botErr != nil {
		status = 2
		errorMessage = botErr.Error()
	}

	// add gen_sql querylog
	go func(ctx context.Context) {
		defer func() {
			if err := recover(); err != nil {
				tilog.Error(ctx, "InteractBot add gen_sql querylog panic", zap.Any("err", err))
			}
		}()
		botType := req.BotType
		if botType == "" {
			botType = "chat2sql"
		}
		promptVersion := botRsp.GetPromptVersion()
		if promptVersion == "" {
			promptVersion = botType + "-default"
		}
		sqlContent := botRsp.GetCompletion()
		// add querylog
		oplogParam := sqleditor.CreateOpLogReq{
			OrgId:         req.OrgID,
			ProjectId:     req.ProjectID,
			ClusterId:     clusterID,
			SessionId:     req.SessionID,
			Chat2QuerySid: botRsp.GetSessionId(),
			SqlFileId:     req.SqlFileID,
			EventType:     "gen_sql",
			UserQuestion:  req.Question,
			SqlContent:    sqlContent,
			StartTime:     startTime.Format("2006-01-02 15:04:05"),
			EndTime:       endTime.Format("2006-01-02 15:04:05"),
			Duration:      executeTime,
			Status:        int32(status),
			ErrorMessage:  errorMessage,
			CreatedBy:     createdBy,
			PromptVersion: promptVersion,
			Databases:     botRsp.GetDatabase(),
			Channel:       botType,
		}
		if err := createOpLog(ctx, &oplogParam); err != nil {
			tilog.Error(ctx, "CreateOplog failed", zap.Error(err), zap.String("event_type", "gen_sql"))
		}

		oplogParam.EventType = "sql_accept"
		oplogParam.IsAccept = true
		if err := createOpLog(ctx, &oplogParam); err != nil {
			tilog.Error(ctx, "CreateOplog failed", zap.Error(err), zap.String("event_type", "sql_accept"))
		}
	}(utils.TraceContext(ctx))

	if botErr != nil {
		tilog.Error(ctx, "InteractBot.ChatBot failed", zap.Any("err", botErr), zap.Any("botReq", req))
		return nil, botErr
	}

	return &dto.InteractResponse{
		SessionID:     botRsp.GetSessionId(),
		Question:      req.Question,
		Prompt:        botRsp.GetPrompt(),
		PromptVersion: botRsp.GetPromptVersion(),
		Response:      botRsp.GetCompletion(),
		StartTime:     startTime.Format("2006-01-02 15:04:05"),
		EndTime:       endTime.Format("2006-01-02 15:04:05"),
		StartTimeMS:   startTime.UnixMilli(),
		EndTimeMS:     endTime.UnixMilli(),
		ExecuteTime:   strconv.FormatInt(executeTime, 10),
		TokenUsed:     botRsp.GetTokenUsed(),
		Database:      botRsp.GetDatabase(),
	}, nil
}

func AddExecOplog(ctx context.Context, param *entity.Chat2QueryParam, sql string, execRes *dataapi.ExecSqlRespV2, err error) {
	// add exec_sql querylog
	go func(ctx context.Context) {
		defer func() {
			if err := recover(); err != nil {
				tilog.Error(ctx, "executeSql add exec_sql querylog panic", zap.Any("err", err))
			}
		}()
		var (
			status       int8 = 1
			errorMessage string
		)
		if err != nil {
			status = 2
			errorMessage = err.Error()
		}

		var sqlRes *dataapi.ExecStmtResult
		if len(execRes.GetStmtResults()) > 0 {
			sqlRes = execRes.GetStmtResults()[0]
		}
		executeTime := sqlRes.GetEndMs() - sqlRes.GetStartMs()
		startTime := time.UnixMilli(sqlRes.GetStartMs())
		endTime := time.UnixMilli(sqlRes.GetEndMs())
		if err := createOpLog(ctx, &sqleditor.CreateOpLogReq{
			OrgId:           param.OrgID,
			ProjectId:       param.ProjectID,
			ClusterId:       param.ClusterID,
			Chat2QuerySid:   param.Chat2QuerySid,
			EventType:       "exec_sql",
			SqlContent:      sql,
			StartTime:       startTime.Format("2006-01-02 15:04:05"),
			EndTime:         endTime.Format("2006-01-02 15:04:05"),
			Duration:        executeTime,
			Status:          int32(status),
			ErrorMessage:    errorMessage,
			CreatedBy:       param.Apikey,
			CreatedAt:       time.Now().Unix(),
			Databases:       strings.Join(sqlRes.GetDatabases(), ","),
			AffectedRows:    sqlRes.GetRowNum(),
			AffectedColumns: int64(len(sqlRes.GetColumns())),
		}); err != nil {
			tilog.Error(ctx, "CreateOplog failed", zap.Error(err), zap.String("event_type", "exec_sql"))
		}
	}(utils.TraceContext(ctx))
}

func createOpLog(ctx context.Context, req *sqleditor.CreateOpLogReq) error {
	if !config.IsCentralMode() {
		httpReq := &http.CreateOpLogReq{
			OrgId:           req.OrgId,
			ProjectId:       req.ProjectId,
			ClusterId:       req.ClusterId,
			Chat2QuerySid:   req.Chat2QuerySid,
			EventType:       req.EventType,
			UserQuestion:    req.UserQuestion,
			SqlContent:      req.SqlContent,
			StartTime:       req.StartTime,
			EndTime:         req.EndTime,
			Duration:        req.Duration,
			IsAccept:        req.IsAccept,
			Status:          req.Status,
			ErrorMessage:    req.ErrorMessage,
			CreatedBy:       req.CreatedBy,
			CreatedAt:       req.CreatedAt,
			Databases:       req.Databases,
			AffectedRows:    req.AffectedRows,
			AffectedColumns: req.AffectedColumns,
			PromptVersion:   req.PromptVersion,
			Channel:         req.Channel,
		}
		return http.CreateOpLog(ctx, httpReq)
	}
	return grpc.CreateOpLog(ctx, req)
}