package resource

import (
	"context"
	"fmt"
	"math/rand"
	"strings"

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

type ResourcePoolScheduler struct {
	pools ResourcePoolMgr
}

func NewResourcePoolScheduler() *ResourcePoolScheduler {
	return &ResourcePoolScheduler{pools: NewTcmsResourcePoolMgr()}
}

func (s *ResourcePoolScheduler) Name() string {
	return "ResourcePoolScheduler"
}

func (s *ResourcePoolScheduler) Refresh(ctx context.Context) error {
	klog.Infof("before refresh pools: %s", s.pools)
	err := s.pools.Refresh(ctx)
	if err != nil {
		return err
	}
	klog.Infof("after refresh pools: %s", s.pools)
	return nil
}

func (s *ResourcePoolScheduler) Prefilter(ctx context.Context, job *common.ArgoJob) error {
	testbed := toITestbed(&job.Plan.TestBed)
	testbed.AttachEstimatedRequest()

	pool, err := s.pools.SetRespool4PlanExec(job.PlanExec, &testbed)
	if err != nil {
		return err
	}

	enough := pool.CheckResourceEnough(testbed.Request())
	if enough.Status == ResCheckStatusError {
		return fmt.Errorf("%s", enough.Message)
	}
	err = s.AttachTestbedName(ctx, pool, job)
	if err != nil {
		return err
	}
	if job.TestbedSetup != nil {
		job.TestbedSetup.Status.Schedule.TestbedSpec = &job.Plan.TestBed
	}
	return nil
}

const MaxNamingTrial = 5

func (s ResourcePoolScheduler) AttachTestbedName(ctx context.Context, pool ResourcePoolInterface, job *common.ArgoJob) error {
	pool = pool.Copy()
	for i := 0; i < MaxNamingTrial; i++ {
		name := s.GenerateName(job, rand.Intn(1000))
		err := pool.Assume(NewJobContext(ctx, job), name, PhysicalResource{})
		if err != nil {
			if strings.HasPrefix(err.Error(), "conflict testbed name") {
				continue
			} else {
				return err
			}
		}
		job.Plan.TestBed.Name = name
		return nil
	}
	return fmt.Errorf("naming testbed reach MaxNamingTrial")
}

func (s ResourcePoolScheduler) GenerateName(job *common.ArgoJob, postfix int) string {
	testbed := &job.Plan.TestBed
	if testbed.Name != "" {
		return testbed.Name
	}
	suffix := fmt.Sprintf("tps-%d-%d-%d", job.PlanExec.ID, 1, postfix)
	if strings.HasSuffix(testbed.GenerateName, "-") {
		return testbed.GenerateName + suffix
	} else {
		return testbed.GenerateName + "-" + suffix
	}
}

func (s *ResourcePoolScheduler) Filter(ctx context.Context, jobs []*common.ArgoJob) ([]*common.ArgoJob, error) {
	pools := s.pools.Copy()

	result := make([]*common.ArgoJob, 0)
	for _, job := range jobs {
		if job.Plan == nil {
			continue
		}
		testbeds := toITestbeds(&job.Plan.TestBed)
		jobCtx := NewJobContext(ctx, job)
		if enough := pools.CheckResourceEnough(jobCtx, testbeds); enough == ResCheckEnough {
			err := pools.Assume(jobCtx, testbeds)
			if err != nil {
				job.Message = err.Error()
				klog.Error(err)
				continue
			}
			result = append(result, job)
		} else {
			job.TestbedSetup.Status.Schedule.ScheduleReason.Message = enough.Message
			klog.Infof("exec_id: %d, %s",
				job.ID(),
				enough.Message)
		}
	}
	klog.Infof("after resource scheduler: %v", common.GetJobExecID(result))
	return result, nil
}

func (s *ResourcePoolScheduler) Reserve(ctx context.Context, job common.ArgoJob) error {
	return s.pools.Assume(NewJobContext(ctx, &job), toITestbeds(&job.Plan.TestBed))
}

func GetJobTestbeds(job *common.ArgoJob) []Testbed {
	return toITestbeds(&job.Plan.TestBed)
}