package scheduler

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

	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/constant"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/entity"
	"github.com/tidbcloud/aws-shared-provider/pkg/auditlog/metric"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/logid"

	"go.uber.org/zap"
)

const (
	taskNumberEachWorker = 2
)

// Dispatch dispatch audit log tasks
func (s *Scheduler) Dispatch() {
	ctx := logid.WithLogid(context.Background(), logid.GenLogid())
	logger := log.WithContextL(ctx)

	logger.Info("start to run dispatch task")

	tasks, err := s.auditLogTaskRepo.List(ctx, &entity.ListAuditLogTaskOptions{
		States: []entity.AuditLogTaskState{entity.AuditLogTaskStatePending, entity.AuditLogTaskStateFailed},
	})

	if err != nil {
		logger.Error("failed to list audit log tasks", zap.Error(err))
		metric.AuditLogScheduleDispatchFailedGauge.
			WithLabelValues(metric.DispatchTypeGlobal, "list audit log schedules error").Inc()
		return
	}

	var wg sync.WaitGroup
	// Process tasks in chunks of taskNumberEachWorker
	for i := 0; i < len(tasks); i += taskNumberEachWorker {
		end := min(i+taskNumberEachWorker, len(tasks))
		wg.Add(1)
		go s.createAuditLogWorker(ctx, tasks[i:end], &wg)
	}

	wg.Wait()
	logger.Info("run dispatch task finish")
}

func (s *Scheduler) createAuditLogWorker(ctx context.Context, tasks []*entity.AuditLogTask, wg *sync.WaitGroup) {
	defer wg.Done()

	taskIds := make([]string, 0, len(tasks))
	for _, task := range tasks {
		taskIds = append(taskIds, task.TaskId)
	}
	logger := log.WithContextL(ctx).With(zap.Any("task_ids", taskIds))
	if len(tasks) == 0 {
		logger.Warn("taskIds is empty, skip create audit log worker")
		return
	}

	// build	worker
	worker := s.auditLogWorker.DeepCopy()

	// add timestamp to avoid worker name conflict
	timestamp := time.Now().Unix()
	worker.Name = fmt.Sprintf("%s-%d", tasks[0].TaskId, timestamp)
	worker.Namespace = s.config.WorkerNs
	// image and args
	args := []string{
		"--bucket-name=" + s.config.BucketName,
		"--region=" + s.config.Region,
		"--rds-secret-arn=" + s.config.RdsSecretArn,
		"--env=" + s.config.Env,
		"--o11y-bucket-name=" + s.config.O11yBucketName,
		"--worker-name=" + worker.Name,
		"--cloud-provider-type=" + s.config.CloudProviderType,
		"--rds-secret-region=" + s.config.RdsSecretRegion,
	}
	for _, taskId := range taskIds {
		args = append(args, "--task-ids="+taskId)
	}
	containers := worker.Spec.Template.Spec.Containers
	for i := range containers {
		if containers[i].Name == constant.AuditLogWorkerContainerName {
			containers[i].Image = s.config.Image
			containers[i].Args = args
			break
		}
	}
	// lables
	worker.Labels = map[string]string{
		types.LabelJobType: constant.AuditLogWorkerLabelValue,
	}
	// service account
	worker.Spec.Template.Spec.ServiceAccountName = constant.AuditLogWorkerServiceAccountName

	logger.Info("start to create audit log worker", zap.String("worker_name", worker.Name))

	err := s.txProvider.Transaction(ctx, func(ctx context.Context) error {
		err := s.auditLogTaskRepo.BatchUpdateState(ctx, entity.AuditLogTaskStateRunning, taskIds)
		if err != nil {
			logger.Error("failed to batch update audit log task", zap.Error(err))
			return err
		}
		err = s.manager.GetClient().Create(ctx, worker)
		if err != nil {
			logger.Error("failed to create audit log worker", zap.Error(err))
			return err
		}
		return nil
	})

	if err != nil {
		logger.Error("failed to create audit log worker", zap.Error(err))
		metric.AuditLogScheduleDispatchFailedGauge.WithLabelValues(metric.DispatchTypeRegional, err.Error()).Inc()
	}
}