package core

import (
	"container/heap"
	"context"
)

// If using rockDB no need priority queue anymore
// implement priority queue from golang
// ref: https://pkg.go.dev/container/heap#example-package-PriorityQueue

type timerServiceKey struct{}

// InjectTimerService inject the timer into Context
func InjectTimerService(ctx context.Context, ts TimerService) context.Context {
	return context.WithValue(ctx, timerServiceKey{}, ts)
}

// ExtractTimerService get the timer for user from context
func ExtractTimerService(ctx context.Context) (TimerService, bool) {
	ts, ok := ctx.Value(timerServiceKey{}).(TimerService)
	return ts, ok
}

type Timer struct {
	// Timestamp is the time when Timer is triggered
	Timestamp int64
	// Key is for who User involve
	Key string
}

// Priority Queue (Min Heap). Smallest timestamp will be on top

type timerHeap []Timer

func (h *timerHeap) Len() int           { return len(*h) }
func (h *timerHeap) Less(i, j int) bool { return (*h)[i].Timestamp < (*h)[j].Timestamp }
func (h *timerHeap) Swap(i, j int) {
	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func (h *timerHeap) Push(x any) {
	*h = append(*h, x.(Timer))
}

func (h *timerHeap) Pop() any {
	old := *h
	n := len(old)
	last := old[n-1]
	*h = old[:n-1]
	return last
}

type TimerService interface {
	Register(time int64, key string)
	CurrentWatermark() int64
}

type InternalTimerService interface {
	TimerService
	AdvanceWatermark(watermark int64) []Timer
}

type timerService struct {
	queue            timerHeap
	currentWatermark int64
}

func NewTimerService() InternalTimerService {
	return &timerService{
		queue:            make(timerHeap, 0),
		currentWatermark: -1,
	}
}

func (t *timerService) Register(time int64, key string) {
	heap.Push(&t.queue, Timer{
		Timestamp: time,
		Key:       key,
	})
}

func (t *timerService) CurrentWatermark() int64 {
	return t.currentWatermark
}

func (t *timerService) AdvanceWatermark(watermark int64) []Timer {
	if watermark < t.currentWatermark {
		// no need to advance when get smaller
		return nil
	}

	t.currentWatermark = watermark

	var timers []Timer
	for t.queue.Len() > 0 {
		earliestTimer := t.queue[0]

		if earliestTimer.Timestamp > watermark {
			break
		}

		pop := heap.Pop(&t.queue).(Timer)
		timers = append(timers, pop)
	}

	return timers
}