import heapq
import random

from turbofuzzllm.fuzzer.template import Template


class PriorityQueue:
    def __init__(self):
        self.heap = []

    def push(self, item, priority):
        # Push the item into the heap with its priority
        heapq.heappush(self.heap, (-priority, random.random(), item))

    def pop(self) -> Template:
        # Pop the item with the highest priority (lowest number)
        if self.is_empty():
            raise IndexError("pop from empty priority queue")
        return heapq.heappop(self.heap)[-1]

    def peek(self):
        # Peek at the item with the highest priority without removing it
        return self.heap[0][-1] if self.heap else None

    def is_empty(self):
        return len(self.heap) == 0

    def contains(self, item):
        return any(item == heap_item for _, _, heap_item in self.heap)

    def remove(self, template):
        self.heap = [item for item in self.heap if item[-1] != template]
        heapq.heapify(self.heap)

    def size(self):
        return len(self.heap)
