package container

import (
	"cmp"
	"container/heap"
)

// 堆封装
type heapQueue[T any, C cmp.Ordered] struct {
	queue   []T
	compare func(T) C
}

func (h *heapQueue[T, C]) Less(i, j int) bool {
	return h.compare(h.queue[i]) < h.compare(h.queue[j])
}

func (h *heapQueue[T, C]) Swap(i, j int) {
	h.queue[i], h.queue[j] = h.queue[j], h.queue[i]
}

func (h *heapQueue[T, C]) Len() int {
	return len(h.queue)
}

func (h *heapQueue[T, C]) Pop() (v any) {
	h.queue, v = h.queue[:h.Len()-1], h.queue[h.Len()-1]
	return
}

func (h *heapQueue[T, C]) Push(v any) {
	h.queue = append(h.queue, v.(T))
}

// Heap 封装堆操作
type Heap[T any, C cmp.Ordered] struct {
	queue *heapQueue[T, C]
}

func NewHeap[T any, C cmp.Ordered](compare func(T) C) *Heap[T, C] {
	return &Heap[T, C]{
		queue: &heapQueue[T, C]{compare: compare},
	}
}

func (h *Heap[T, C]) Push(v T) {
	heap.Push(h.queue, v)
}

func (h *Heap[T, C]) Pop() T {
	return heap.Pop(h.queue).(T)
}

func (h *Heap[T, C]) Peek() (v T, ok bool) {
	if h.queue.Len() == 0 {
		return v, false
	}
	return h.queue.queue[0], true
}