package api

import (
	"context"
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"github.com/spf13/viper"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.buf.build/protocolbuffers/go/linuxgit/tidbcloud/base"
	"go.uber.org/zap"

	"github.com/tidbcloud/aipower/commons/chatgpttemplate"
	"github.com/tidbcloud/aipower/commons/config"
	"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/bean/entity"
	"github.com/tidbcloud/aipower/internal/dao"
	"github.com/tidbcloud/aipower/internal/service/cachesvc"
	"github.com/tidbcloud/aipower/internal/service/openaisvc"
	"github.com/tidbcloud/aipower/pkg/ratelimit"
)

const maxQuestionLength = 1024 * 5

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

type chatBotService struct {
	api.UnsafeChatBotServiceServer
}

func NewChatBotService() api.ChatBotServiceServer {
	return &chatBotService{}
}

func (c *chatBotService) checkIfHackyQuestion(ctx context.Context, req *api.ChatBotReq) error {
	if !config.Global().ServerConf.HackyQuestionEnabled {
		return nil
	}

	question := strings.Trim(req.Question, " ")

	tilog.Info(ctx, "hacky question check enabled", zap.String("question", question))
	// hacky question, which is used for testing purpose
	switch question {
	case consts.HackyQuestionForRateLimit:
		msg := fmt.Sprintf(consts.RateLimitErrorMsgTemplate, 100, config.Global().ServerConf.OrgRateLimitOverSupportURL)
		return errs.NewTooManyRequests(msg, nil)
	case consts.HackyQuestionNoSQLGenerated:
		return errs.NewBadRequest(chatgpttemplate.DefaultNoSqlGeneratedMsg, nil)
	case consts.HackyQuestionAIPowerGoesDown:
		return errs.NewInternalDB("database connection invalid(fake for testing purpose)", nil)
	default:
		return nil
	}
}

func getOrgMaxRequestUpperLimit(ctx context.Context, orgID uint64) int {
	conf := config.Global().ServerConf
	m := viper.GetStringMapString(config.RateLimitWhitelist)
	tilog.Info(ctx, "read org rate limit whitelist from oneconfig", zap.Any("whitelist", m))

	maxCountStr, inWhitelist := m[fmt.Sprintf("%d", orgID)]
	upperLimit := conf.OrgRateLimitPerDay
	if inWhitelist {
		var err error
		upperLimit, err = strconv.Atoi(maxCountStr)
		if err != nil {
			tilog.Error(ctx, "failed to parse whitelist from oneconfig", zap.Error(err))
			return conf.OrgRateLimitPerDay
		}
	}

	return upperLimit
}

func (c *chatBotService) ChatBot(ctx context.Context, req *api.ChatBotReq) (*api.ChatBotRsp, error) {
	tilog.Info(ctx, "api chat bot req", zap.Any("req", req))
	if req.Caller == "" {
		tilog.Info(ctx, "param invalid for ChatBot", zap.Any("req", req))
		return nil, errs.NewBadRequest("param invalid", nil)
	}

	conf := config.Global().ServerConf

	if err := c.checkIfHackyQuestion(ctx, req); err != nil {
		tilog.Info(ctx, "hacky question matched", zap.String("question", req.Question), zap.Error(err))
		return nil, err
	}

	if conf.EnableRateLimit {
		count, err := ratelimit.ByOrgIDPerDay(req.OrgId)
		if err != nil {
			tilog.Error(ctx, "failed to check rate limit", zap.Error(err))
			return nil, errs.NewInternalGRPC("failed to check rate limit", err)
		}

		if req.OrgId == 0 {
			tilog.Error(ctx, "org id did not pass in request", zap.Any("request", req))
		} else {
			upperLimit := getOrgMaxRequestUpperLimit(ctx, req.OrgId)

			if count > upperLimit {
				msg := fmt.Sprintf(consts.RateLimitErrorMsgTemplate, upperLimit, config.Global().ServerConf.OrgRateLimitOverSupportURL)
				tilog.Info(ctx, "requests exceeds the threshold", zap.Uint64("orgID", req.OrgId), zap.Int("count", count), zap.Int("upperLimit", upperLimit))
				return nil, errs.NewTooManyRequests(msg, nil)
			}
		}
	}

	if len(req.Question) > maxQuestionLength {
		return nil, errs.NewBadRequest("question is too large", nil)
	}

	caller := req.Caller
	param := &dto.InteractRequest{
		BotType:               req.BotType,
		SessionID:             req.SessionId,
		Question:              req.Question,
		Database:              req.Database,
		ProjectID:             req.ProjectId,
		ClusterID:             req.ClusterId,
		PromptTemplate:        req.PromptTemplate,
		PromptTemplateVersion: req.PromptTemplateVersion,
		QuestionNums:          int(req.QuestionNums),
		Tables:                req.Tables,
		Model:                 req.Model,
		SampleData:            req.SampleData,
		DatabaseSchema:        req.DatabaseSchema,
		Async:                 req.Async,
		Caller:                req.Caller,
		Timeout:               int(req.Timeout),
		OrgID:                 req.OrgId,
		AppID:                 req.AppId,
		AppKey:                req.AppKey,
		Endpoint:              req.Endpoint,
		Source:                req.Source,
	}

	result, err := openaisvc.InteractBot(ctx, param, caller)

	rsp := &api.ChatBotRsp{
		BaseResp: &base.BaseResp{
			ErrCode: 0,
			ErrMsg:  "ok",
		},
		SessionId:  result.SessionID,
		Completion: result.Response,
		TokenUsed:  int32(result.TokenUsed),
		Database:   result.Database,
		Duration:   result.ExecuteTime,
		JobId:      result.JobID,
	}

	if conf.Deployment == consts.DeploymentOffline {
		rsp.Prompt = result.Prompt
		rsp.PromptVersion = result.PromptVersion
	}

	if err != nil {
		tilog.Info(ctx, "api chat bot err", zap.Any("'rsp", rsp), zap.Any("'err", err))
		return rsp, err
	}

	tilog.Info(ctx, "api chat bot success", zap.Any("'req", req), zap.Any("'rsp", rsp))

	return rsp, nil

}

func (c *chatBotService) GetLatestPrompt(ctx context.Context, req *api.GetLatestPromptReq) (*api.GetLatestPromptResp, error) {
	prompt, err := dao.GetLatestPromptByAIModelAndDeploy(req.AiModel, req.Deployment)
	if err != nil {
		return nil, errs.WrapGORMError(err, "failed to get latest prompt")
	}

	resp := &api.GetLatestPromptResp{
		BaseResp: &base.BaseResp{
			ErrCode: 0,
			ErrMsg:  "ok",
		},
		Prompt: dto.ConvertPromptTemplate(prompt),
	}

	return resp, nil
}

func (c *chatBotService) GetAllPromptByAIModel(ctx context.Context, req *api.GetAllPromptByAIModelReq) (*api.GetAllPromptByAIModelResp, error) {
	prompts, err := dao.GetAllPromptsByAIModelAndDeploy(req.AiModel, req.Deployment)
	if err != nil {
		return nil, errs.WrapGORMError(err, "failed to get all prompts")
	}

	// unique by aimodel + deployment
	aiModelMapper := map[string]entity.PromptTemplate{}
	for _, i := range prompts {
		key := fmt.Sprintf("%s-%s", i.AIModel, i.Deployment)
		aiModelMapper[key] = i
	}

	// get latest prompt by aimodel + deployment
	latestAIModelMapper := map[string]*entity.PromptTemplate{}
	for key, prompt := range aiModelMapper {
		latestPrompt, err := dao.GetLatestPromptByAIModelAndDeploy(prompt.AIModel, prompt.Deployment)
		if err == nil {
			latestAIModelMapper[key] = latestPrompt
		}
	}

	var items []*api.Prompt
	for _, i := range prompts {
		item := dto.ConvertPromptTemplate(&i)

		// check if it is latest one by checking version
		key := fmt.Sprintf("%s-%s", i.AIModel, i.Deployment)
		if latestPrompt, ok := latestAIModelMapper[key]; ok {
			if latestPrompt.Version == i.Version {
				item.IsLatest = true
			}
		}

		items = append(items, item)
	}

	resp := &api.GetAllPromptByAIModelResp{
		BaseResp: &base.BaseResp{
			ErrCode: 0,
			ErrMsg:  "ok",
		},
		Prompts: items,
	}

	return resp, nil
}

func (c *chatBotService) SavePrompt(ctx context.Context, req *api.SavePromptReq) (*api.SavePromptResp, error) {
	err := dao.SavePrompt(req.Prompt.AiModel, req.Prompt.Deployment, req.Prompt.Version, req.Prompt.Content)
	if err != nil {
		return nil, errs.WrapGORMError(err, "failed to save prompt")
	}

	resp := &api.SavePromptResp{
		BaseResp: &base.BaseResp{
			ErrCode: 0,
			ErrMsg:  "ok",
		},
	}

	return resp, nil
}

func (c *chatBotService) QueryCompletionByJobID(ctx context.Context, req *api.QueryCompletionReq) (*api.QueryCompletionResp, error) {
	if req.JobId == "" {
		return nil, errs.NewBadRequest("job_id is required", nil)
	}
	if !streamEnabledAIModels[req.BotType] {
		return nil, errs.NewBadRequest("invalid bot_type", nil)
	}

	fullCompletion, err := cachesvc.ReadCompletionResultByJobID(ctx, req.JobId)
	if err != nil {
		tilog.Info(ctx, "failed to query completion from redis", zap.Any("req", req), zap.Error(err))
		return nil, err
	}

	var resp api.QueryCompletionResp
	if err = json.Unmarshal([]byte(fullCompletion), &resp); err != nil {
		return nil, errs.NewInternalRedis("failed to read completion", err)
	}

	if resp.Status != consts.AIPowerStreamStatusDone {
		resp.Message = ""
	}

	return &resp, nil
}