package xslices

import (
	"sync"
)

func NewLRU[K comparable, T any](capacity int) *lru[K, T] {
	return &lru[K, T]{
		items:    make(map[K]*lruNode[K, T], capacity),
		capacity: capacity,
	}
}

type lruNode[K comparable, T any] struct {
	key   K
	value T

	next *lruNode[K, T]
	prev *lruNode[K, T]
}

type lru[K comparable, T any] struct {
	items    map[K]*lruNode[K, T]
	length   int
	capacity int

	first *lruNode[K, T]
	last  *lruNode[K, T]

	mu sync.RWMutex
}

func (l *lru[K, T]) Push(key K, value T) {
	l.mu.Lock()
	defer l.mu.Unlock()

	node, ok := l.items[key]
	if ok {
		l.bumpUnsafe(node)
		return
	}

	node = &lruNode[K, T]{
		key:   key,
		value: value,
		next:  l.first,
	}

	if l.first != nil {
		l.first.prev = node
	}
	if l.last == nil {
		l.last = node
	}

	l.first = node
	l.items[key] = node

	if l.length == l.capacity && l.last != nil {
		deletedNode := l.last
		delete(l.items, deletedNode.key)

		l.last = l.last.prev
		return
	}

	l.length++
}

func (l *lru[K, T]) Get(key K) (T, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	node, ok := l.items[key]
	if !ok {
		var t T
		return t, false
	}

	out := node.value
	l.bumpUnsafe(node)

	return out, true
}

func (l *lru[K, T]) bumpUnsafe(node *lruNode[K, T]) {
	if l.first == node {
		return
	}

	if node.next != nil {
		node.next.prev = node.prev
	}
	if node.prev != nil {
		node.prev.next = node.next
	}

	node.next = l.first
	l.first.prev = node

	l.first = node
	node.prev = nil
}