package service

import (
	"context"
	"fmt"

	scalerpb "github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/proto/scaler"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/schema/autoanalyze"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbworker/schema/bgtask"
	"gorm.io/gorm/clause"
)

func (s *ScalerService) RegisterAutoAnalyze(
	ctx context.Context,
	request *scalerpb.RegisterAutoAnalyzeRequest,
) (*scalerpb.Response, error) {
	if request == nil || request.GetHeader() == nil {
		return scalerResponse(ErrBadRequest), nil
	}

	if s.GetConfig().EnableAutoAnalyze && s.GetConfig().EnableSharedWorkerForAutoAnalyze {
		err := s.GetDB().WithContext(ctx).Clauses(clause.OnConflict{
			DoNothing: true,
		}).Create(&bgtask.Task{
			KeyspaceName: request.GetHeader().GetKeyspaceName(),
			TaskKey:      fmt.Sprintf("auto-analyze/%d", request.TaskId),
			GlobalTaskID: int64(request.TaskId),
			TaskType:     "auto-analyze",
			TidbPool:     request.GetHeader().GetTidbPool(),
			ExecID:       "",
			State:        bgtask.StatePending,
		}).Error
		return scalerResponse(err), nil
	}

	err := s.GetDB().WithContext(ctx).Clauses(clause.OnConflict{
		DoNothing: true,
	}).Create(&autoanalyze.Task{
		KeyspaceName: request.GetHeader().GetKeyspaceName(),
		TidbPool:     request.GetHeader().GetTidbPool(),
		TaskID:       request.GetTaskId(),
		State:        autoanalyze.StatePending,
	}).Error
	return scalerResponse(err), nil
}

func (s *ScalerService) RecycleAutoAnalyze(
	ctx context.Context,
	request *scalerpb.RecycleAutoAnalyzeRequest,
) (*scalerpb.Response, error) {
	if request == nil || request.GetHeader() == nil {
		return scalerResponse(ErrBadRequest), nil
	}

	res := s.GetDB().WithContext(ctx).Model(&autoanalyze.Task{}).
		Where("keyspace_name = ?", request.GetHeader().GetKeyspaceName()).
		Where("task_id = ?", request.GetTaskId()).
		Update("state", autoanalyze.StateDone)
	if res.RowsAffected == 0 {
		// not found in autoanalyze task, try to clean bgtask
		res = s.GetDB().WithContext(ctx).Model(&bgtask.Task{}).
			Where("keyspace_name = ?", request.GetHeader().GetKeyspaceName()).
			Where("global_task_id = ?", request.GetTaskId()).
			Where("task_type = ?", "auto-analyze").
			Update("state", bgtask.StateDone)
		return scalerResponse(res.Error), nil
	}
	return scalerResponse(res.Error), nil
}