package grpc

import (
	"context"
	"dataservices/commons/config"
	"fmt"
	"github.com/pkg/errors"
	"strings"
	"sync"
	"time"

	"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 (
	aipowerInitOnce sync.Once
	aipowerCli      api.ChatBotServiceClient
)

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

func getAipowerClient() api.ChatBotServiceClient {
	aipowerInitOnce.Do(func() {
		conn, err := tigrpc.NewClientConn(config.Global().GrpcConf.AipowerAddr, tigrpc.WithDefault())
		if err != nil {
			panic(fmt.Sprintf("dial aipower: %s", err))
		}

		aipowerCli = api.NewChatBotServiceClient(conn)
	})

	return aipowerCli
}

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 := getAipowerClient()

	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 strings.Contains(errors.Errorf("%+v", err).Error(), "context deadline") {
			return rsp, tierr.NewInternalGRPC("Execution time exceeds timeout!", err)
		}
		return nil, tierr.NewInternalGRPC(aiTips, err)
	}
	if rsp.BaseResp.ErrCode != 0 {
		tiErr := tierr.GetTiErrFromBaseResp(rsp.BaseResp)
		if aiApprovedErrCode[rsp.BaseResp.ErrCode] {
			return nil, tiErr
		} else {
			return nil, tierr.NewInternalGRPC(aiTips, tiErr)
		}
	}
	return rsp, nil
}