package shs

import (
	"iter"

	"github.com/liyue201/gostl/ds/rbtree"
	"github.com/liyue201/gostl/utils/comparator"
)

var _ Index = (*RBTreeIndex)(nil)

type RBTreeIndex rbtree.RbTree[string, uint64]

func NewRBTreeIndex() *RBTreeIndex {
	return (*RBTreeIndex)(rbtree.New[string, uint64](comparator.StringComparator))
}

func (R *RBTreeIndex) Get(key []byte) (uint64, error) {
	t := (*rbtree.RbTree[string, uint64])(R)
	return t.Find(bytesToString(key))
}

func (R *RBTreeIndex) Set(key []byte, value uint64) error {
	t := (*rbtree.RbTree[string, uint64])(R)
	t.Insert(bytesToString(key), value)
	return nil
}

func (R *RBTreeIndex) Delete(key []byte) error {
	t := (*rbtree.RbTree[string, uint64])(R)
	node := t.FindNode(bytesToString(key))
	if node == nil {
		return ErrKeyNotFound
	}
	t.Delete(node)
	return nil
}

func (R *RBTreeIndex) Iterator() iter.Seq2[[]byte, uint64] {
	return func(yield func([]byte, uint64) bool) {
		t := (*rbtree.RbTree[string, uint64])(R)
		t.Traversal(func(key string, value uint64) bool {
			if !yield(stringToBytes(key), value) {
				return false
			}
			return true
		})
	}
}