package sqleditor

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"sync"
	"time"

	"github.com/tidbcloud/tiinf/transport/http/tigin"

	"dataapps/commons/config"
	"dataapps/commons/grpc"
	"dataapps/consts"
	dto "dataapps/internal/dto/sqleditor"
	service "dataapps/internal/service/sqleditor"
	"dataapps/internal/types"

	tierr "github.com/tidbcloud/tiinf/errs"
	"go.uber.org/zap"

	"github.com/gin-gonic/gin"
	"github.com/tidbcloud/tiinf/monitor/tilog"
)

const maxInstructionLength = 1024

var (
	streamEnabledAIModels = map[string]bool{
		consts.AIModelChat2SQL:       true,
		consts.AIModelColumnMatching: true,
		consts.AIModelChat2Question:  true,
	}
)

func Data2ColumnMatching(c *gin.Context) {
	ctx := c.Request.Context()
	startT := time.Now()
	param, err := parseChat2QueryParam(c, false, false)
	tilog.Info(ctx, "Data2ColumnMatching start", zap.Any("param", param))
	if err != nil {
		tilog.Error(ctx, "Data2ColumnMatching parse error", zap.Any("err", err))
		res := &dto.BasicRes{
			Code:    int(err.BizCode()),
			Message: err.ErrMsg(),
		}
		c.JSON(http.StatusBadRequest, res)
		return
	}
	columnDescription, jobID, e := service.Data2ColumnMatching(ctx, param)
	if e != nil {
		tilog.Info(ctx, "Data2ColumnMatching request error", zap.Any("err", e), zap.Any("param", param))
		tigin.AbortWithError(c, e)
		return
	}
	result := dto.Data2ColumnMatchingResult{
		BasicRes: dto.BasicRes{
			Code:    http.StatusOK,
			Message: "ok",
		},
		Columns: columnDescription,
		JobID:   jobID,
	}
	tc := time.Since(startT)
	result.CostTime = tc.String()
	c.JSON(http.StatusOK, result)
}

func Chat2Question(c *gin.Context) {
	ctx := c.Request.Context()
	startT := time.Now()
	param, err := parseChat2QueryParam(c, false, true)
	tilog.Info(ctx, "Chat2Question start", zap.Any("param", param), zap.Any("err", err))
	if err != nil {
		res := &dto.BasicRes{
			Code:    int(err.BizCode()),
			Message: err.ErrMsg(),
		}
		c.JSON(http.StatusBadRequest, res)
		return
	}

	conf := config.Global().ServerConfig
	param.Instruction = "chat2question"
	if param.QuestionNum < 2 {
		param.QuestionNum = 2
	}
	if param.QuestionNum > 5 {
		param.QuestionNum = 5
	}
	var (
		questions []dto.Chat2QuestionResp
		jobID     string
		e         error
	)
	// todo: add switch and concurrencyNum
	// 1. if we enable streaming in server side, and the request use a async mode
	// 2. if we don't generate question concurrently
	if (conf.EnableStreaming && param.Async) || !conf.GenQuestionConcurrently {
		questions, jobID, e = service.Chat2Question(ctx, param)
		if e != nil {
			tilog.Info(ctx, "Chat2Question error", zap.Any("err", e), zap.Any("param", param))
			tigin.AbortWithError(c, e)
			return
		}
	} else {
		questions, e = chat2QuestionConcurrency(ctx, param, 3)
		if e != nil {
			tilog.Info(ctx, "Concurrent Chat2Question error ", zap.Any("err", e), zap.Any("param", param))
			tigin.AbortWithError(c, e)
			return
		}
	}

	result := dto.Chat2QuestionResult{
		BasicRes: dto.BasicRes{
			Code:    http.StatusOK,
			Message: "ok",
		},
		Questions: questions,
		JobID:     jobID,
	}

	tc := time.Since(startT)
	result.CostTime = tc.String()

	tilog.Info(ctx, "Chat2Question end", zap.Any("cost", result.CostTime), zap.Any("result", result), zap.Any("param", param))

	c.JSON(http.StatusOK, result)
}

func chat2QuestionConcurrency(ctx context.Context, param *dto.Chat2QueryParam, concurrencyNum int) ([]dto.Chat2QuestionResp, error) {
	chanQuestionArray := make(chan []dto.Chat2QuestionResp)
	var wg sync.WaitGroup
	for i := 0; i < concurrencyNum; i++ {
		wg.Add(1)
		go chat2QuestionRoutine(ctx, param, &wg, chanQuestionArray)
	}
	// close the channel when all goroutines have completed
	go func() {
		wg.Wait()
		close(chanQuestionArray)
	}()
	// filter question
	var questionsReturn []dto.Chat2QuestionResp
	for questionArray := range chanQuestionArray {
		tilog.Info(ctx, "question array return by goroutine", zap.Any("questionArray", questionArray))
		if len(questionsReturn) >= param.QuestionNum {
			break
		}
		for _, questionInfo := range questionArray {
			if len(questionsReturn) >= param.QuestionNum {
				break
			}
			// default not exist
			isExisted := false
			for _, questionExisted := range questionsReturn {
				if questionInfo.Question == questionExisted.Question || questionInfo.QuestionKeyword == questionExisted.QuestionKeyword {
					isExisted = true
					break
				}
			}
			if !isExisted {
				questionsReturn = append(questionsReturn, questionInfo)
			}
		}
	}
	if len(questionsReturn) == 0 {
		return questionsReturn, tierr.NewBadRequest("No Question Generated", errors.New("no Question Generated"))
	}
	return questionsReturn, nil
}

func chat2QuestionRoutine(ctx context.Context, param *dto.Chat2QueryParam, wg *sync.WaitGroup, filteredQuestion chan []dto.Chat2QuestionResp) {
	questions, _, err := service.Chat2Question(ctx, param)
	if err == nil {
		filteredQuestion <- questions
	} else {
		tilog.Warn(ctx, "chat2Question error for concurrency", zap.Error(err))
	}
	defer wg.Done()
}

func Chat2Chart(c *gin.Context) {
	ctx := c.Request.Context()
	startT := time.Now()
	param, err := parseChat2QueryParam(c, true, true)
	tilog.Info(ctx, "Chat2Chart start", zap.Any("param", param), zap.Any("err", err))
	if err != nil {
		res := &dto.BasicRes{
			Code:    int(err.BizCode()),
			Message: err.ErrMsg(),
		}
		c.JSON(http.StatusBadRequest, res)
		return
	}

	result, e := service.Chat2Chart(ctx, param)
	if e != nil {
		tilog.Error(ctx, "Chat2Chart error", zap.Any("err", e), zap.Any("param", param))
		if result != nil {
			result.Code = int(tierr.ErrCodeInternalServerError)
			result.Message = e.Error()
			c.JSON(http.StatusOK, result)
			return
		}
		res := &dto.BasicRes{
			Code:    int(tierr.ErrCodeInternalServerError),
			Message: e.Error(),
		}
		c.JSON(http.StatusOK, res)
		return
	}

	result.Code = http.StatusOK
	result.Message = "ok"

	tc := time.Since(startT)
	result.CostTime = tc.String()

	tilog.Info(ctx, "Chat2Chart end", zap.Any("cost", result.CostTime), zap.Any("result", result), zap.Any("param", param))

	c.JSON(http.StatusOK, result)
}

func InternalChat2Data(c *gin.Context) {
	ctx := c.Request.Context()
	var req types.InternalChat2DataReq
	if err := c.ShouldBindJSON(&req); err != nil {
		tigin.AbortWithError(c, tierr.NewBadRequest(err.Error(), nil))
		return
	}

	startT := time.Now()
	param := &dto.Chat2QueryParam{
		ProjectID:             req.ProjectID,
		ClusterID:             req.ClusterID,
		Database:              req.Database,
		Instruction:           req.Instruction,
		PromptTemplate:        req.PromptTemplate,
		PromptTemplateVersion: req.PromptTemplateVersion,
		Apikey:                "internal",
		ShowInternalInfo:      true,
	}

	result, e := service.Chat2Data(ctx, param)
	if e != nil {
		tilog.Info(ctx, "Chat2Data error", zap.Any("err", e), zap.Any("param", param))
		if result != nil {
			result.Code = int(tierr.ErrCodeInternalServerError)
			result.Message = e.Error()
			c.JSON(http.StatusInternalServerError, result)
			return
		}
		res := &dto.BasicRes{
			Code:    int(tierr.ErrCodeInternalServerError),
			Message: e.Error(),
		}
		c.JSON(http.StatusInternalServerError, res)
		return
	}

	result.Code = http.StatusOK
	result.Message = "ok"

	tc := time.Since(startT)
	result.CostTime = tc.String()

	tilog.Info(ctx, "Chat2Data end", zap.Any("cost", result.CostTime), zap.Any("result", result), zap.Any("param", param))

	c.JSON(http.StatusOK, result)
}

func Chat2Data(c *gin.Context) {
	ctx := c.Request.Context()
	startT := time.Now()
	param, err := parseChat2QueryParam(c, true, true)
	tilog.Info(ctx, "Chat2Data start", zap.Any("param", param), zap.Any("err", err))
	if err != nil {
		res := &dto.BasicRes{
			Code:    int(err.BizCode()),
			Message: err.ErrMsg(),
		}
		c.JSON(http.StatusBadRequest, res)
		return
	}

	result, e := service.Chat2Data(ctx, param)
	if e != nil {
		tilog.Info(ctx, "Chat2Data error", zap.Any("err", e), zap.Any("param", param))
		if result != nil {
			result.Code = int(tierr.ErrCodeInternalServerError)
			result.Message = e.Error()
			c.JSON(http.StatusInternalServerError, result)
			return
		}
		res := &dto.BasicRes{
			Code:    int(tierr.ErrCodeInternalServerError),
			Message: e.Error(),
		}
		c.JSON(http.StatusInternalServerError, res)
		return
	}

	result.Code = http.StatusOK
	result.Message = "ok"

	tc := time.Since(startT)
	result.CostTime = tc.String()

	tilog.Info(ctx, "Chat2Data end", zap.Any("cost", result.CostTime), zap.Any("result", result), zap.Any("param", param))

	c.JSON(http.StatusOK, result)
}

func getAuthInfo(c *gin.Context) (*dto.Chat2QueryAuthInfo, tierr.Error) {
	authInfoAny, ok := c.Get(consts.CONTEXT_CURRENT_API_KEY)
	if !ok {
		msg := "not found auth info"
		return nil, tierr.NewBadRequest(msg, errors.New(msg))
	}
	authInfo, ok := authInfoAny.(*dto.Chat2QueryAuthInfo)
	if !ok {
		msg := "auth info invalid"
		return nil, tierr.NewBadRequest(msg, errors.New(msg))
	}

	return authInfo, nil
}

func parseChat2QueryParam(c *gin.Context, needInstruction bool, needDatabaseSchema bool) (*dto.Chat2QueryParam, tierr.Error) {
	authInfo, err := getAuthInfo(c)
	if err != nil {
		return nil, err
	}

	req := &dto.Chat2QueryReq{}
	if err := c.ShouldBindJSON(req); err != nil {
		tilog.Error(c.Request.Context(), "invalid params", zap.Any("req", req), zap.Error(err))
		msg := "request param invalid"
		return nil, tierr.NewBadRequest(msg, errors.New(msg))
	}
	if needDatabaseSchema {
		if req.ClusterID == 0 {
			msg := "cluster_id invalid"
			return nil, tierr.NewBadRequest(msg, errors.New(msg))
		}
		database := req.Database
		if database == "" {
			msg := "database invalid"
			return nil, tierr.NewBadRequest(msg, errors.New(msg))
		}
	}
	if needInstruction {
		if req.Instruction == "" {
			msg := "instruction invalid"
			return nil, tierr.NewBadRequest(msg, errors.New(msg))
		}
		if len(req.Instruction) > maxInstructionLength {
			msg := "instruction is too long"
			return nil, tierr.NewBadRequest(msg, errors.New(msg))
		}
	}
	// new added param
	var (
		sampleData     string
		databaseSchema string
	)
	if len(req.SampleData) != 0 {
		jsonBytes, err := json.Marshal(req.SampleData)
		if err == nil {
			tilog.Error(c.Request.Context(), "marshal req.SampleData error", zap.Error(err))
			sampleData = string(jsonBytes)
		}
	}
	if len(req.DatabaseSchema) != 0 {
		jsonBytes, err := json.Marshal(req.DatabaseSchema)
		tilog.Error(c.Request.Context(), "marshal req.DatabaseSchema error", zap.Error(err))
		if err == nil {
			databaseSchema = string(jsonBytes)
		}
	}
	param := &dto.Chat2QueryParam{
		OrgID:          authInfo.OrgID,
		ProjectID:      authInfo.ProjectID,
		Apikey:         authInfo.ApiKey,
		Instruction:    req.Instruction,
		ClusterID:      req.ClusterID,
		Database:       req.Database,
		QuestionNum:    req.QuestionNum,
		Tables:         req.Tables,
		SampleData:     sampleData,
		DatabaseSchema: databaseSchema,
		Async:          req.Async,
	}

	return param, nil
}

func QueryStreamResultByJobID(c *gin.Context) {
	var req dto.QueryStreamResultReq
	if err := c.ShouldBindQuery(&req); err != nil {
		tigin.AbortWithError(c, tierr.NewBadRequest("invalid request", nil))
		return
	}

	ctx := c.Request.Context()
	tilog.Info(ctx, "request params", zap.Any("req", req))
	if !streamEnabledAIModels[req.BotType] {
		tigin.AbortWithError(c, tierr.NewBadRequest("invalid bot_type", nil))
		return
	}

	resp, err := grpc.QueryStreamResponse(ctx, &req)
	tilog.Info(ctx, "query streaming response", zap.Any("req", req), zap.Any("resp", resp), zap.Error(err))
	if err != nil {
		tigin.AbortWithError(c, err)
		return
	}

	// parse completion
	var result dto.QueryStreamResultResp
	switch req.BotType {
	case consts.AIModelColumnMatching:
		columnDescription := []dto.ColumnDescription{}
		err = json.Unmarshal([]byte(resp.Completion), &columnDescription)
		if err != nil {
			tigin.AbortWithError(c, tierr.NewInternalHTTP("failed to unmarshal", err))
			return
		}

		result.Columns = columnDescription
	case consts.AIModelChat2Question:
		questions := []dto.Chat2QuestionResp{}
		err = json.Unmarshal([]byte(resp.Completion), &questions)
		if err != nil {
			tigin.AbortWithError(c, tierr.NewInternalHTTP("failed to unmarshal", err))
			return
		}

		result.Questions = questions
	case consts.AIModelChat2SQL:
		result.SQL = resp.Completion
	}

	c.JSON(http.StatusOK, gin.H{
		"result":  result,
		"message": resp.Message,
		"status":  resp.Status,
	})
}