package main

import (
	"container/heap"
	"fmt"
)

type server struct {
	Idx          int
	Weights      int
	CompleteTime int
}

type baseQueue []*server

func (b baseQueue) Len() int {
	return len(b)
}

func (b baseQueue) Swap(i, j int) {
	b[i], b[j] = b[j], b[i]
}

func (b *baseQueue) Push(x interface{}) {
	*b = append(*b, x.(*server))
}

func (b *baseQueue) Pop() interface{} {
	n := b.Len()
	t := (*b)[n-1]
	*b = (*b)[:n-1]
	return t
}

type freeServerQueue struct {
	baseQueue
}

func (fs *freeServerQueue) Less(i, j int) bool {
	if fs.baseQueue[i].Weights == fs.baseQueue[j].Weights {
		return fs.baseQueue[i].Idx < fs.baseQueue[j].Idx
	}
	return fs.baseQueue[i].Weights < fs.baseQueue[j].Weights
}

type workingServerQueue struct {
	baseQueue
}

func (ws *workingServerQueue) Less(i, j int) bool {
	return ws.baseQueue[i].CompleteTime < ws.baseQueue[j].CompleteTime
}

func assignTasks(servers []int, tasks []int) []int {
	serArr := make([]*server, 0, 0)
	for i, v := range servers {
		serArr = append(serArr, &server{
			Weights:      v,
			Idx:          i,
			CompleteTime: 0,
		})
	}
	freeHeap := &freeServerQueue{
		serArr,
	}
	heap.Init(freeHeap)
	workingHeap := &workingServerQueue{
		make([]*server, 0, 0),
	}
	heap.Init(workingHeap)
	res := make([]int, 0, 0)
	nowTime := 0
	// left - 当前已处理任务(开区间) ， nowTime - 当前可开始任务（闭区间）
	left := 0
	// 当还有剩余任务未处理时，循环继续
	for left < len(tasks) {
		// 从workingQueue中找到当前刚处理完任务的服务器
		for workingHeap.baseQueue.Len() > 0 && workingHeap.baseQueue[0].CompleteTime <= nowTime {
			freeServer := heap.Pop(workingHeap).(*server)
			heap.Push(freeHeap, freeServer)
		}
		// 当有空闲服务器时，且当前任务可以开始时，pop堆顶权重最小服务器进行任务处理
		for freeHeap.Len() > 0 && freeHeap.baseQueue[0].CompleteTime <= nowTime && left <= nowTime && left < len(tasks) {
			delTime := tasks[left]
			tpServer := heap.Pop(freeHeap).(*server)
			res = append(res, tpServer.Idx)
			tpServer.CompleteTime = delTime + nowTime
			heap.Push(workingHeap, tpServer)
			left++
		}
		// 当没有空闲服务器可用时，将时间直接跳到最近的空闲时间
		if freeHeap.Len() == 0 {
			nowTime = workingHeap.baseQueue[0].CompleteTime
		} else {
			nowTime++
		}
	}
	return res
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func main() {
	fmt.Println(assignTasks([]int{3, 3, 2}, []int{1, 2, 3, 2, 1, 2}))
}