package schedulers

import (
	"context"

	"github.com/pingcap/test-infra/services/tps/pkg/common"
	"k8s.io/klog/v2"
)

type ChainScheduler struct {
	refreshSchedulers   []RefreshScheduler
	prefilterSchedulers []PrefilterScheduler
	filterSchedulers    []FilterScheduler
	reserveSchedulers   []ReserveScheduler
}

func (c *ChainScheduler) Name() string {
	return "ChainScheduler"
}

func (c *ChainScheduler) Prefilter(ctx context.Context, job *common.ArgoJob) error {
	for i, scheduler := range c.prefilterSchedulers {
		err := scheduler.Prefilter(ctx, job)
		if err != nil {
			klog.Errorf("scheduler %s(%d) raise error at Prefilter: %s", scheduler.Name(), i, err.Error())
			job.TestbedSetup.Status.Schedule.ScheduleReason.SchedulerName = scheduler.Name()
			return err
		}
	}
	return nil
}

func (c *ChainScheduler) Refresh(ctx context.Context) error {
	for i, scheduler := range c.refreshSchedulers {
		err := scheduler.Refresh(ctx)
		if err != nil {
			klog.Errorf("scheduler %s(%d) raise error at Refresh: %s", scheduler.Name(), i, err.Error())
			return err
		}
	}
	return nil
}

func (c *ChainScheduler) Reserve(ctx context.Context, job common.ArgoJob) (err error) {
	for i, scheduler := range c.reserveSchedulers {
		err = scheduler.Reserve(ctx, job)
		if err != nil {
			klog.Errorf("scheduler %s(%d) raise error at Reserve: %s", scheduler.Name(), i, err.Error())
			return err
		}
	}
	return nil
}

func (c *ChainScheduler) Filter(ctx context.Context, jobs []*common.ArgoJob) ([]*common.ArgoJob, error) {
	var current []*common.ArgoJob = jobs
	var err error
	for i, scheduler := range c.filterSchedulers {
		if len(current) == 0 {
			break
		}
		last := current
		current, err = scheduler.Filter(ctx, last)
		if err != nil {
			klog.Errorf("scheduler %s(%d) raise error at Filter: %s", scheduler.Name(), i, err.Error())
			return current, err
		}
		survivor := map[*common.ArgoJob]struct{}{}
		for _, job := range current {
			survivor[job] = struct{}{}
		}
		for _, job := range last {
			if _, survive := survivor[job]; !survive {
				job.TestbedSetup.Status.Schedule.ScheduleReason.SchedulerName = scheduler.Name()
			}
		}
		klog.Infof("after  %s(%d) scheduler: %v", scheduler.Name(), i, common.GetJobExecID(current))
	}
	return current, nil
}