package schedulers

import (
	"context"
	"testing"
	"time"

	"github.com/pingcap/test-infra/common/model/tcms/taskruns"

	commM "github.com/pingcap/test-infra/common/model"
	"github.com/pingcap/test-infra/services/tps/pkg/common"
	"github.com/stretchr/testify/assert"
)

func newJob(priority int, poolID commM.ID) *common.ArgoJob {
	return &common.ArgoJob{
		PlanExec:     &commM.TestPlanExecution{Priority: commM.ExecPriority(priority), ResourcepoolID: poolID},
		Plan:         &commM.TestPlan{TestPlanBase: commM.TestPlanBase{TestBed: commM.TestBed{ResourcePool: "dummy"}}},
		EPriority:    commM.ExecPriority(priority),
		TestbedSetup: &taskruns.TestbedSetup{Status: taskruns.TestbedSetupStatus{Schedule: &taskruns.Schedule{}}},
	}
}

func TestPriorityScheduler_Filter_Conflict(t *testing.T) {
	s := &PriorityScheduler{}
	jobs := []*common.ArgoJob{newJob(0, 123), newJob(10, 123), newJob(0, 123)}
	expected := []*common.ArgoJob{jobs[1]}
	filtered, err := s.Filter(context.TODO(), jobs)
	assert.NoError(t, err)
	assert.Equal(t, expected, filtered)
}

func TestPriorityScheduler_Filter_NoConflict(t *testing.T) {
	s := &PriorityScheduler{}
	jobs := []*common.ArgoJob{newJob(0, 123), newJob(10, 123), newJob(0, 456)}
	expected := []*common.ArgoJob{jobs[1], jobs[2]}
	filtered, err := s.Filter(context.TODO(), jobs)
	assert.NoError(t, err)
	assert.Equal(t, expected, filtered)
}

func TestPriorityScheduler_Filter_EmptyPool(t *testing.T) {
	s := &PriorityScheduler{}
	jobs := []*common.ArgoJob{newJob(0, 123), newJob(10, 123), newJob(0, 456)}
	expected := []*common.ArgoJob{jobs[1], jobs[2]}
	filtered, err := s.Filter(context.TODO(), jobs)
	assert.NoError(t, err)
	assert.Equal(t, expected, filtered)
}

func TestPriorityScheduler_Filter_MultiConflict(t *testing.T) {
	s := &PriorityScheduler{}
	jobs := []*common.ArgoJob{newJob(0, 123), newJob(10, 123), newJob(0, 456), newJob(5, 456)}
	expected := []*common.ArgoJob{jobs[1], jobs[3]}
	filtered, err := s.Filter(context.TODO(), jobs)
	assert.NoError(t, err)
	assert.Equal(t, expected, filtered)
}

func TestPriorityScheduler_Prefilter_Effective(t *testing.T) {
	createTime := time.Unix(1, 0)
	s := &PriorityScheduler{now: func() time.Time {
		return createTime.Add(49 * time.Hour)
	}}
	job := newJob(int(commM.ExecPriorityDefault), 123)
	job.PlanExec.CreatedAt = commM.Timestamp(createTime.Unix())
	err := s.Prefilter(context.TODO(), job)
	assert.NoError(t, err)
	assert.Equal(t, commM.ExecPriorityDefault+2, job.EPriority)
}

func TestPriorityScheduler_Prefilter_Effective_ZeroCreate(t *testing.T) {
	s := &PriorityScheduler{now: func() time.Time {
		return time.Unix(0, 0).Add(49 * time.Hour)
	}}
	job := newJob(int(commM.ExecPriorityDefault), 123)
	err := s.Prefilter(context.TODO(), job)
	assert.NoError(t, err)
	assert.Equal(t, commM.ExecPriorityDefault, job.EPriority)
}

func TestGetBlockedBy(t *testing.T) {
	jobs := []*common.ArgoJob{
		newJob(0, 123),
		newJob(2, 123),
		newJob(1, 123),
		newJob(1, 123),
		newJob(5, 456)}
	jobs_ := make([]*common.ArgoJob, len(jobs))
	copy(jobs_, jobs)
	blockedBy := getBlockedBy(jobs_)
	expected := map[*common.ArgoJob][]*common.ArgoJob{
		jobs[0]: {jobs[2], jobs[3]},
		jobs[2]: {jobs[1]},
		jobs[3]: {jobs[1]},
	}
	assert.Equal(t, expected, blockedBy)
}

func TestSplitJobs(t *testing.T) {
	jobs := []*common.ArgoJob{
		newJob(0, 123),
		newJob(2, 123),
		newJob(1, 123),
		newJob(1, 123)}
	jobs_ := make([]*common.ArgoJob, len(jobs))
	copy(jobs_, jobs)
	act := splitJobs(jobs_)
	expected := [][]*common.ArgoJob{
		{jobs[1]},
		{jobs[2], jobs[3]},
		{jobs[0]},
	}
	assert.Equal(t, expected, act)
}