package aipowersvc

import (
	"context"

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

	"github.com/tidbcloud/aitestsuit/internal/entity"
)

var (
	aipowerClient aipower.ChatBotServiceClient
)

func InitAIPower(addr string) error {
	conn, err := tigrpc.NewClientConn(addr, tigrpc.WithDefault())
	if err != nil {
		return err
	}

	aipowerClient = aipower.NewChatBotServiceClient(conn)

	return nil
}

func GenSQLByTestCase(ctx context.Context, projectID, clusterID uint64, model string, testCase entity.TestCase, promptTemplate entity.PromptTemplate) (*aipower.ChatBotRsp, error) {
	req := aipower.ChatBotReq{
		ProjectId:             projectID,
		ClusterId:             clusterID,
		Database:              testCase.Database,
		Question:              testCase.Question,
		Caller:                "ai-platform",
		BotType:               model,
		PromptTemplate:        promptTemplate.Content,
		PromptTemplateVersion: promptTemplate.Version,
		OrgId:                 30061,
	}
	resp, err := aipowerClient.ChatBot(ctx, &req)

	tilog.Info(ctx, "request aipower with payload", zap.Any("payload", &req), zap.Any("resp", resp), zap.Error(err))

	return resp, err
}

func SavePromptTemplate(ctx context.Context, aiModel, deployment string, promptTemplate entity.PromptTemplate) (*aipower.SavePromptResp, error) {
	req := aipower.SavePromptReq{Prompt: &aipower.Prompt{
		AiModel:    aiModel,
		Deployment: deployment,
		Version:    promptTemplate.Version,
		Content:    promptTemplate.Content,
	}}

	return aipowerClient.SavePrompt(ctx, &req)
}

func GetCurrentPromptTemplate(ctx context.Context, aiModel, deployment string) (*aipower.Prompt, error) {
	req := aipower.GetLatestPromptReq{
		AiModel:    aiModel,
		Deployment: deployment,
	}

	resp, err := aipowerClient.GetLatestPrompt(ctx, &req)
	if err != nil {
		return nil, err
	}

	return resp.Prompt, nil
}

func GetAllPromptTemplate(ctx context.Context, aiModel, deployment string) ([]*aipower.Prompt, error) {
	req := aipower.GetAllPromptByAIModelReq{
		AiModel:    aiModel,
		Deployment: deployment,
	}

	resp, err := aipowerClient.GetAllPromptByAIModel(ctx, &req)
	if err != nil {
		return nil, err
	}

	var items []*aipower.Prompt
	for _, i := range resp.Prompts {
		items = append(items, i)
	}

	return items, nil
}