package schedulers

import (
	"context"
	"fmt"
	"sort"
	"time"

	commM "github.com/pingcap/test-infra/common/model"

	"github.com/pingcap/test-infra/services/tps/pkg/common"
)

type PriorityScheduler struct {
	now func() time.Time
}

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

const DailyPriorityQueueCnt = 14 // 2 week

func (s *PriorityScheduler) Prefilter(ctx context.Context, job *common.ArgoJob) error {
	planExec := job.PlanExec
	job.EPriority = job.PlanExec.Priority
	if s.isDailyJob(job) {
		if planExec.CreatedAt == 0 {
			return nil
		}
		waitTime := s.now().Sub(planExec.CreatedAt.Time())
		waitedDays := int(waitTime / (24 * time.Hour))
		if waitedDays > DailyPriorityQueueCnt {
			waitedDays = DailyPriorityQueueCnt
		}
		job.EPriority += commM.ExecPriority(waitedDays)
	}
	return nil
}

func (s *PriorityScheduler) isDailyJob(job *common.ArgoJob) bool {
	// right now  only daily run use ExecPriorityDefault
	// todo fix this later
	return job.PlanExec.Priority == commM.ExecPriorityDefault
}

func (s *PriorityScheduler) Filter(ctx context.Context, jobs []*common.ArgoJob) ([]*common.ArgoJob, error) {
	blockedByMap := getBlockedBy(jobs)
	var ret []*common.ArgoJob
	for _, job := range jobs {
		if blockedBy := blockedByMap[job]; len(blockedBy) != 0 {
			fillUnScheduleReason(job, blockedBy)
		} else {
			ret = append(ret, job)
		}
	}
	return ret, nil
}

func fillUnScheduleReason(job *common.ArgoJob, blockedBy []*common.ArgoJob) *common.ArgoJob {
	scheduleReason := &job.TestbedSetup.Status.Schedule.ScheduleReason
	scheduleReason.Message = fmt.Sprintf("blocked by high priority jobs, current priority: %d", job.EPriority)
	scheduleReason.RelatedSetupIds = nil
	for _, argoJob := range blockedBy {
		scheduleReason.RelatedSetupIds = append(scheduleReason.RelatedSetupIds, argoJob.TestbedSetup.MetaData.ID)
	}
	return job
}

func getBlockedBy(jobs []*common.ArgoJob) map[*common.ArgoJob][]*common.ArgoJob {
	queues := poolQueues{
		poolQueues:   map[commM.ID]*JobQueue{},
		blockedByMap: map[*common.ArgoJob][]*common.ArgoJob{},
	}
	for _, job := range jobs {
		queues.add(job)
	}
	queues.collectLowPriorityJobs()
	return queues.blockedByMap
}

type poolQueues struct {
	poolQueues   map[commM.ID]*JobQueue
	blockedByMap map[*common.ArgoJob][]*common.ArgoJob
}

func (c *poolQueues) add(job *common.ArgoJob) {
	c.getQueue(job.PlanExec.ResourcepoolID).add(job)
}

func (c *poolQueues) isLowPriority(job *common.ArgoJob) bool {
	if _, ok := c.blockedByMap[job]; ok {
		return true
	}
	return false
}

func (c *poolQueues) getQueue(queue commM.ID) *JobQueue {
	if _, ok := c.poolQueues[queue]; !ok {
		c.poolQueues[queue] = &JobQueue{}
	}
	return c.poolQueues[queue]
}

func (c *poolQueues) collectLowPriorityJobs() {
	for _, queue := range c.poolQueues {
		segments := splitJobs(*queue)
		var prev []*common.ArgoJob
		for _, segment := range segments {
			for _, job := range segment {
				if prev != nil {
					c.blockedByMap[job] = prev
				}
			}
			prev = segment
		}
	}
}

type JobQueue []*common.ArgoJob

func (q *JobQueue) add(job ...*common.ArgoJob) {
	*q = append(*q, job...)
}

func splitJobs(jobs []*common.ArgoJob) (rt [][]*common.ArgoJob) {
	jobs = sortByPriority(jobs)
	if len(jobs) == 0 {
		return rt
	}
	var current = getPriority(jobs[0])
	var prev = 0
	for i, job := range jobs {
		p := getPriority(job)
		if p < current {
			rt = append(rt, jobs[prev:i])
			prev = i
			current = p
		} else if p > current {
			panic("jobs must be sorted by Priority")
		}
	}
	rt = append(rt, jobs[prev:])
	return rt
}

func sortByPriority(jobs []*common.ArgoJob) []*common.ArgoJob {
	sort.SliceStable(jobs, func(i, j int) bool {
		return getPriority(jobs[j]) < getPriority(jobs[i])
	})
	return jobs
}

func getPriority(job *common.ArgoJob) int {
	return int(job.EPriority)
}