package service

import (
	"context"
	"encoding/json"
	"strconv"
	"strings"
	"sync"
	"time"

	"dataapps/commons/config"
	"dataapps/commons/grpc"
	dto "dataapps/internal/dto/sqleditor"
	entity "dataapps/internal/entity/sqleditor"
	"dataapps/internal/service/util"

	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"

	tierr "github.com/tidbcloud/tiinf/errs"

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

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

const defaultUnknownAIErrorMsg = "No SQL could be generated based on your question, please try another one."
const defaultRunSQLErrorMsg = "No data is returned based on AI-generated SQL, please try another question."
const defaultNotSelectSQLMsg = "Please try not to update the current dataset."

func Data2ColumnMatching(ctx context.Context, param *dto.Chat2QueryParam) ([]dto.ColumnDescription, string, error) {
	tilog.Info(ctx, "service data2ColumnMatching start", zap.Any("param", param))
	startT := time.Now()
	resp, err := ChatBotWithParam(ctx, "data2columnMatching", param)
	tc := time.Since(startT)
	tilog.Info(ctx, "service data2ColumnMatching end", zap.String("cost", tc.String()), zap.Any("resp", resp))
	if err != nil {
		tilog.Error(ctx, "data2ColumnMatching error", zap.Error(err))
		return nil, "", tierr.NewInternalGRPC("failed to get data type and description", err)
	}
	columnDescription := []dto.ColumnDescription{}
	conf := config.Global().ServerConfig

	if !(conf.EnableStreaming && param.Async) {
		err = json.Unmarshal([]byte(resp.Response), &columnDescription)
		if err != nil {
			tilog.Error(ctx, "Parse columnDescription failed.", zap.Error(err))
			return nil, "", tierr.NewBadRequest(defaultUnknownAIErrorMsg, err)
		}
	}
	return columnDescription, resp.JobID, nil
}

func Chat2Question(ctx context.Context, param *dto.Chat2QueryParam) ([]dto.Chat2QuestionResp, string, error) {
	tilog.Info(ctx, "service chat2question start", zap.Any("param", param))
	startT := time.Now()
	resp, err := ChatBotWithParam(ctx, "chat2question", param)
	tc := time.Since(startT)
	tilog.Info(ctx, "service chat2question gen questions", zap.Any("cost", tc.String()), zap.Any("resp", resp), zap.Any("err", err))
	if err != nil {
		return nil, "", tierr.NewInternalGRPC("failed to gen questions", err)
	}
	questions := []dto.Chat2QuestionResp{}
	conf := config.Global().ServerConfig

	if !(conf.EnableStreaming && param.Async) {
		err = json.Unmarshal([]byte(resp.Response), &questions)
		if err != nil {
			tilog.Error(ctx, "Parse Chat2ChartInfo failed. "+err.Error())
			return nil, "", tierr.NewBadRequest(defaultUnknownAIErrorMsg, err)
		}
		// todo: add switch
		if conf.GenQuestionConcurrently {
			tilog.Warn(ctx, "unfiltered questions", zap.Any("unfilteredQuestions", questions))
			questions = questionsFilter(ctx, param, questions)
			tilog.Warn(ctx, "filtered questions", zap.Any("filteredQuestions", questions))
		}
	}

	return questions, resp.JobID, nil
}

func questionsFilter(ctx context.Context, param *dto.Chat2QueryParam, questions []dto.Chat2QuestionResp) []dto.Chat2QuestionResp {
	chanFilteredQuestions := make(chan dto.Chat2QuestionResp)
	var wg sync.WaitGroup
	for _, questionInfo := range questions {
		wg.Add(1)
		go questionValidate(ctx, param, questionInfo, &wg, chanFilteredQuestions)
	}
	// close the channel when all goroutines have completed
	go func() {
		wg.Wait()
		close(chanFilteredQuestions)
	}()
	var filteredQuestions []dto.Chat2QuestionResp
	for question := range chanFilteredQuestions {
		filteredQuestions = append(filteredQuestions, question)
	}
	return filteredQuestions
}

func questionValidate(ctx context.Context, param *dto.Chat2QueryParam, questionInfo dto.Chat2QuestionResp, wg *sync.WaitGroup, filteredQuestion chan dto.Chat2QuestionResp) {
	newChat2QueryParam := &dto.Chat2QueryParam{
		OrgID:                 param.OrgID,
		ProjectID:             param.ProjectID,
		ClusterID:             param.ClusterID,
		Database:              param.Database,
		Instruction:           questionInfo.Question,
		Apikey:                param.Apikey,
		PromptTemplate:        param.PromptTemplate,
		PromptTemplateVersion: param.PromptTemplateVersion,
		ShowInternalInfo:      param.ShowInternalInfo,
		Chat2QuerySid:         param.Chat2QuerySid,
		Tables:                param.Tables,
	}
	result, err := Chat2Chart(ctx, newChat2QueryParam)
	if err != nil {
		tilog.Warn(ctx, "question validate failed", zap.String("question", questionInfo.Question), zap.Error(err))
	} else if len(result.Rows) == 0 {
		tilog.Warn(ctx, "question validate failed", zap.String("question", questionInfo.Question), zap.Int("rows", 0))
	} else {
		// Has query data return
		filteredQuestion <- questionInfo
	}
	defer wg.Done()
}

func Chat2Chart(ctx context.Context, param *dto.Chat2QueryParam) (*dto.Chat2QueryResult, error) {
	tilog.Info(ctx, "service chat2chart start", zap.Any("param", param))
	startT := time.Now()
	resp, err := ChatBotWithParam(ctx, "chat2chart", param)
	tc := time.Since(startT)
	tilog.Info(ctx, "service chat2chart gen sql", zap.Any("cost", tc.String()), zap.Any("resp", resp), zap.Any("err", err))
	if err != nil {
		return nil, err
	}
	chat2ChartInfo := dto.Chat2ChartInfo{}
	res := resp.Response
	err = json.Unmarshal([]byte(res), &chat2ChartInfo)
	if err != nil {
		tilog.Error(ctx, "Parse Chat2ChartInfo failed. "+err.Error())
		return nil, tierr.NewBadRequest(defaultUnknownAIErrorMsg, err)
	}

	stmtType, err := GetStmtType(ctx, chat2ChartInfo.Sql)
	if err != nil {
		ret := &dto.Chat2QueryResult{
			GenSQL:         chat2ChartInfo.Sql,
			CallOpenAITime: tc.String(),
		}
		tilog.Error(ctx, "GetStmtType failed", zap.Error(err), zap.String("sql", chat2ChartInfo.Sql))
		errMsg := "Parse sql failed, please check if your sql is correct"
		return ret, tierr.NewBadRequest(errMsg, nil)
	}
	if stmtType != dataapi.StmtType_Select {
		ret := &dto.Chat2QueryResult{
			GenSQL:         chat2ChartInfo.Sql,
			CallOpenAITime: tc.String(),
		}
		tilog.Error(ctx, "StmtType is invalid", zap.String("sql", chat2ChartInfo.Sql), zap.String("type", stmtType.String()))
		return ret, tierr.NewBadRequest(defaultNotSelectSQLMsg, nil)
	}

	param.Chat2QuerySid = resp.SessionID
	dataRes, err := executeSql(ctx, param, chat2ChartInfo.Sql)
	tilog.Info(ctx, "service chat2data exec sql", zap.Any("dataRes", dataRes), zap.Any("err", err))
	if err != nil {
		ret := &dto.Chat2QueryResult{
			GenSQL:         chat2ChartInfo.Sql,
			CallOpenAITime: tc.String(),
		}
		tilog.Error(ctx, "ExecuteSql failed. "+err.Error())
		return ret, tierr.NewBadRequest(defaultRunSQLErrorMsg, err)
	}

	chatRet := generateChatResult(chat2ChartInfo.Sql, dataRes, tc, resp, false)
	chatRet.ChartInfo = &chat2ChartInfo.ChartInfo
	tilog.Info(ctx, "service chat2data gen chat result", zap.Any("chatRet", chatRet))
	return chatRet, nil
}

func Chat2Data(ctx context.Context, param *dto.Chat2QueryParam) (*dto.Chat2QueryResult, 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))
	if err != nil {
		return nil, err
	}

	sql := resp.Response
	// stmtType, err := GetStmtType(ctx, sql)
	// if err != nil {
	// 	tilog.Error(ctx, "GetStmtType failed", zap.Error(err), zap.String("sql", sql))
	// 	return nil, bizerr.Restore(bizerr.ErrCodeGeneralBusiness, "Parse sql failed, please check if your sql is correct")
	// }
	// if stmtType != dataapi.StmtType_Select {
	// 	ret := &dto.Chat2QueryResult{
	// 		GenSQL:         sql,
	// 		CallOpenAITime: tc.String(),
	// 	}
	// 	tilog.Error(ctx, "StmtType is invalid", zap.String("sql", sql), zap.String("type", stmtType.String()))
	// 	return ret, tierr.NewBadRequest(defaultNotSelectSQLMsg, nil)
	// }
	param.Chat2QuerySid = resp.SessionID
	dataRes, err := executeSql(ctx, param, sql)
	tilog.Info(ctx, "service chat2data exec sql", zap.Any("dataRes", dataRes), zap.Any("err", err))
	if err != nil {
		ret := &dto.Chat2QueryResult{
			GenSQL:         sql,
			CallOpenAITime: tc.String(),
		}
		return ret, err
	}

	chatRet := generateChatResult(sql, dataRes, tc, resp, param.ShowInternalInfo)
	tilog.Info(ctx, "service chat2data gen chat result", zap.Any("chatRet", chatRet))

	return chatRet, nil
}

func generateChatResult(sql string, dataRes *entity.DataRes, duration time.Duration, response *dto.InteractResponse, showInternalInfo bool) *dto.Chat2QueryResult {
	chatRet := dto.Chat2QueryResult{
		GenSQL:         sql,
		CallOpenAITime: duration.String(),
		StartTime:      dataRes.StartTime,
		EndTime:        dataRes.EndTime,
		ExecuteTime:    dataRes.ExecuteTime,
		Query:          dataRes.Query,
		RowCount:       dataRes.RowCount,
		Rows:           dataRes.Rows,
		Columns:        dataRes.Columns,
		Limit:          dataRes.Limit,
	}

	if showInternalInfo {
		chatRet.TokenUsed = int(response.TokenUsed)
		chatRet.FullPrompt = response.Prompt
	}

	return &chatRet
}

func executeSql(ctx context.Context, param *dto.Chat2QueryParam, sql string) (*entity.DataRes, error) {

	dataRes := &entity.DataRes{}
	dataRes.Limit = LIMIT
	sqlResp, err := grpc.ExecSqlDirectly(ctx, &dataapi.BasicEndpointReq{
		ProjectId: param.ProjectID,
		ClusterId: param.ClusterID,
		Database:  param.Database,
		Sql:       sql,
		LimitRows: uint64(LIMIT),
	})
	if err != nil {
		return nil, err
	}
	dataRes = convertDataAPIResult(ctx, sql, sqlResp, dataRes)

	// 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()
		}
		executeTime, _ := strconv.ParseInt(dataRes.ExecuteTime, 10, 64)
		if err := CreateOplog(ctx, entity.SqlEditorOplog{
			OrgId:           param.OrgID,
			ProjectId:       param.ProjectID,
			ClusterId:       param.ClusterID,
			Chat2QuerySId:   param.Chat2QuerySid,
			EventType:       string(entity.EventTypeExecSql),
			SqlContent:      sql,
			StartTime:       dataRes.StartTime,
			EndTime:         dataRes.EndTime,
			Duration:        executeTime,
			Status:          status,
			ErrorMessage:    errorMessage,
			CreatedBy:       param.Apikey,
			CreatedAt:       time.Now(),
			Databases:       strings.Join(dataRes.Databases, ","),
			AffectedRows:    dataRes.RowCount,
			AffectedColumns: int64(len(dataRes.Columns)),
		}); err != nil {
			tilog.Error(ctx, "CreateOplog failed", zap.Error(err), zap.String("event_type", string(entity.EventTypeExecSql)))
		}
	}(util.TraceContext(ctx))

	if err != nil {
		tilog.Error(ctx, "exec sql failed", zap.Any("param", param), zap.Any("sql", sql), zap.Any("err", err))
		return nil, err
	}

	return dataRes, nil
}