package service

import (
	"context"

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

func (s *ScalerService) GetBgTaskConfig(
	ctx context.Context,
	request *scalerpb.GetBgTaskConfigRequest,
) (*scalerpb.GetBgTaskConfigResponse, error) {
	// TODO: implement after server refactor.
	return &scalerpb.GetBgTaskConfigResponse{Error: scalerError(ErrNotImplemented)}, nil
}

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

	err := s.GetDB().WithContext(ctx).Clauses(clause.OnConflict{
		DoNothing: true,
	}).Create(&bgtask.Task{
		KeyspaceName: request.GetHeader().GetKeyspaceName(),
		TaskKey:      request.GetTaskKey(),
		GlobalTaskID: request.GetGlobalTaskId(),
		SubTaskID:    request.GetSubTaskId(),
		TaskType:     request.GetTaskType(),
		TidbPool:     request.GetHeader().GetTidbPool(),
		ExecID:       request.GetExecId(),
		State:        bgtask.StatePending,
	}).Error
	return scalerResponse(err), nil
}

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

	txn := s.GetDB().WithContext(ctx).Model(&bgtask.Task{}).
		Where("keyspace_name = ?", request.GetHeader().GetKeyspaceName()).
		Where("global_task_id = ?", request.GetGlobalTaskId())
	if request.GetSubTaskId() != 0 {
		txn = txn.Where("sub_task_id = ?", request.GetSubTaskId())
	}
	err := txn.Update("state", bgtask.StateDone).Error
	return scalerResponse(err), nil
}

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

	m := make(map[string][]int64)
	for i := range request.GetSubTaskIds() {
		execID, subTaskID := request.GetExecIds()[i], request.GetSubTaskIds()[i]
		m[execID] = append(m[execID], subTaskID)
	}

	err := s.GetDB().WithContext(ctx).Model(&bgtask.Task{}).Transaction(func(txn *gorm.DB) error {
		for execID, subTaskIDs := range m {
			err := txn.Model(&bgtask.Task{}).
				Where("keyspace_name = ?", request.GetHeader().GetKeyspaceName()).
				Where("global_task_id = ?", request.GetGlobalTaskId()).
				Where("sub_task_id IN ?", subTaskIDs).
				Update("exec_id", execID).Error
			if err != nil {
				return err
			}
		}
		return nil
	})

	return scalerResponse(err), nil
}