package grpc

import (
	"context"
	"fmt"
	"sync"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"dataapps/commons/config"
	dto "dataapps/internal/dto/sqleditor"

	"github.com/tidbcloud/aipower/idl/pbgen/proto/api"
	tierr "github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	tigrpc "github.com/tidbcloud/tiinf/transport/tigrpc/client"
	"go.uber.org/zap"
)

var (
	messageServiceInitOnce sync.Once
	messageServiceCli      api.ChatBotServiceClient
)

const (
	aiOldTips = "We couldn't generate SQL for this instruction, please try another one."
	aiNewTips = "We could not generate SQL for this instruction, please try another one."
)

func getMessageServiceClient() api.ChatBotServiceClient {
	messageServiceInitOnce.Do(func() {
		conn, err := tigrpc.NewClientConn(config.Global().GrpcConfig.MessageServiceAddr, tigrpc.WithDefault())
		if err != nil {
			panic(fmt.Sprintf("dial aipower: %s", err))
		}

		messageServiceCli = api.NewChatBotServiceClient(conn)
	})

	return messageServiceCli
}

var aiApprovedErrCode = map[int64]bool{
	tierr.ErrCodeTooManyRequests:                         true,
	int64(api.ErrCodes_ERR_CODES_GEN_SQL_FAILED):         true,
	int64(api.ErrCodes_ERR_CODES_GUESS_DATABASE_FAILED):  true,
	int64(api.ErrCodes_ERR_CODES_TOO_MANY_TABLES_FAILED): true,
	int64(api.ErrCodes_ERR_CODES_INVALID_PARAMA):         true,
	int64(api.ErrCodes_ERR_CODES_REQUEST_OPENAI_ERROR):   true,
	int64(api.ErrCodes_ERR_CODES_GET_DB_METAS_ERROR):     true,
	int64(api.ErrCodes_ERR_CODES_GET_DB_METAS_EMPTY):     true,
}

func ChatBot(ctx context.Context, req *api.ChatBotReq) (*api.ChatBotRsp, error) {
	tilog.Info(ctx, "rpc/aipower:ChatBot req", zap.Any("req", req))
	startT := time.Now()
	client := getMessageServiceClient()

	rsp, err := client.ChatBot(ctx, req)
	tc := time.Since(startT)
	tilog.Info(ctx, "rpc/aipower:ChatBot rsp", zap.Any("cost", tc.String()), zap.Any("rsp", rsp), zap.Any("err", err))
	if err != nil {
		if status.Code(err) == codes.Canceled || status.Code(err) == codes.DeadlineExceeded {
			// timeout
			return rsp, tierr.NewRequestTimeout(aiNewTips, err)
		} else {
			return rsp, tierr.NewInternalGRPC(aiNewTips, err)
		}
	}
	if rsp.BaseResp.ErrCode != 0 {
		tilog.Error(ctx, "rpc/aipower:ChatBot fail, baseresp errcode not 0", zap.Any("rsp", rsp))
		tiErr := tierr.GetTiErrFromBaseResp(rsp.BaseResp)
		if !aiApprovedErrCode[rsp.BaseResp.ErrCode] {
			tiErr = tierr.NewInternalGRPC(aiNewTips, tiErr)
		}
		rsp.Code = int32(tiErr.BizCode())
		rsp.Completion = tiErr.ErrMsg()
		if rsp.Completion == aiOldTips {
			rsp.Completion = aiNewTips
		}
	}
	return rsp, nil
}

func QueryStreamResponse(ctx context.Context, req *dto.QueryStreamResultReq) (*api.QueryCompletionResp, error) {
	client := getMessageServiceClient()

	apiReq := &api.QueryCompletionReq{
		JobId:   req.JobID,
		BotType: req.BotType,
	}
	return client.QueryCompletionByJobID(ctx, apiReq)
}