package container

import (
	"golang.org/x/exp/constraints"
)

type HeapItem[T any, V constraints.Integer] struct {
	Value    T
	Priority V
}

// Heap represents a generic heap data structure.
type Heap[T any, V constraints.Integer] struct {
	elements  []HeapItem[T, V]
	isMinHeap bool
}

// NewHeap creates a new heap with the specified ordering.
// The isMinHeap parameter determines whether the heap is a min heap or a max heap.
// The heap is initially empty.
func NewHeap[T any, V constraints.Integer](isMinHeap bool) *Heap[T, V] {
	return &Heap[T, V]{elements: []HeapItem[T, V]{}, isMinHeap: isMinHeap}
}

// Push adds an element to the heap.
func (h *Heap[T, V]) Push(element T, priority V) {
	h.elements = append(h.elements, HeapItem[T, V]{element, priority})
	h.heapifyUp()
}

// Pop removes and returns the root element from the heap.
func (h *Heap[T, V]) Pop() T {
	if len(h.elements) == 0 {
		var zero T // Return zero value of T
		return zero
	}

	root := h.elements[0]
	lastIndex := len(h.elements) - 1
	h.elements[0] = h.elements[lastIndex]
	h.elements = h.elements[:lastIndex]

	h.heapifyDown()
	return root.Value
}

// Peek returns the root element without removing it.
func (h *Heap[T, V]) Peek() T {
	if len(h.elements) == 0 {
		var zero T // Return zero value of T
		return zero
	}
	return h.elements[0].Value
}

// heapifyUp adjusts the heap after adding a new element.
func (h *Heap[T, V]) heapifyUp() {
	index := len(h.elements) - 1
	for index > 0 {
		parentIndex := (index - 1) / 2
		if h.compare(h.elements[index].Priority, h.elements[parentIndex].Priority) {
			h.elements[index], h.elements[parentIndex] = h.elements[parentIndex], h.elements[index]
			index = parentIndex
		} else {
			break
		}
	}
}

// heapifyDown adjusts the heap after removing the root element.
func (h *Heap[T, V]) heapifyDown() {
	index := 0
	lastIndex := len(h.elements) - 1
	for index < lastIndex {
		leftChildIndex := 2*index + 1
		rightChildIndex := 2*index + 2

		var childIndex int
		if leftChildIndex <= lastIndex {
			childIndex = leftChildIndex
			if rightChildIndex <= lastIndex && h.compare(h.elements[rightChildIndex].Priority, h.elements[leftChildIndex].Priority) {
				childIndex = rightChildIndex
			}

			if h.compare(h.elements[childIndex].Priority, h.elements[index].Priority) {
				h.elements[index], h.elements[childIndex] = h.elements[childIndex], h.elements[index]
				index = childIndex
			} else {
				break
			}
		} else {
			break
		}
	}
}

// compare compares two elements based on the heap type.
func (h *Heap[T, V]) compare(x, y V) bool {
	if h.isMinHeap {
		return x < y
	}
	return x > y
}