package mdb

import (
	"sync"
	"sync/atomic"

	"github.com/google/btree"
)

type BTreeIndex[T any] struct {
	c       *Collection[T]
	name    string
	modLock sync.Mutex
	bt      atomic.Value // *btree.BTreeG[*T]
	getID   func(*T) uint64
	less    func(*T, *T) bool
	include func(*T) bool
}

func NewBTreeIndex[T any](
	c *Collection[T],
	name string,
	less func(*T, *T) bool,
	include func(*T) bool,
) *BTreeIndex[T] {

	t := &BTreeIndex[T]{
		c:       c,
		name:    name,
		getID:   c.getID,
		less:    less,
		include: include,
	}

	btree := btree.NewG(32, less)
	t.bt.Store(btree)
	c.indices = append(c.indices, t)

	return t
}

func (t *BTreeIndex[T]) load(m map[uint64]*T) error {
	btree := btree.NewG(64, t.less)
	t.bt.Store(btree)
	for _, item := range m {
		if t.include == nil || t.include(item) {
			if x, _ := btree.ReplaceOrInsert(item); x != nil {
				return ErrDuplicate
			}
		}
	}
	return nil
}

// ----------------------------------------------------------------------------

func (t *BTreeIndex[T]) Ascend() *BTreeIterator[T] {
	iter := newBTreeIterator[T]()
	go func() {
		t.btree().Ascend(iter.each)
		iter.done()
	}()
	return iter
}

func (t *BTreeIndex[T]) AscendAfter(pivot T) *BTreeIterator[T] {
	iter := newBTreeIterator[T]()
	go func() {
		t.btree().AscendGreaterOrEqual(&pivot, iter.each)
		iter.done()
	}()
	return iter
}

func (t *BTreeIndex[T]) Descend() *BTreeIterator[T] {
	iter := newBTreeIterator[T]()
	go func() {
		t.btree().Descend(iter.each)
		iter.done()
	}()
	return iter
}

func (t *BTreeIndex[T]) DescendAfter(pivot T) *BTreeIterator[T] {
	iter := newBTreeIterator[T]()
	go func() {
		t.btree().DescendLessOrEqual(&pivot, iter.each)
		iter.done()
	}()
	return iter
}

func (t *BTreeIndex[T]) Get(item T) (T, bool) {
	ptr, ok := t.btree().Get(&item)
	if !ok {
		return item, false
	}
	return *ptr, true
}

func (t *BTreeIndex[T]) Min() (item T, ok bool) {
	if ptr, ok := t.btree().Min(); ok {
		return *ptr, ok
	}
	return item, false
}

func (t *BTreeIndex[T]) Max() (item T, ok bool) {
	if ptr, ok := t.btree().Max(); ok {
		return *ptr, ok
	}
	return item, false
}

func (t *BTreeIndex[T]) Len() int {
	return t.btree().Len()
}

// ----------------------------------------------------------------------------

func (t *BTreeIndex[T]) insert(item *T) {
	if t.include == nil || t.include(item) {
		t.modify(func(bt *btree.BTreeG[*T]) {
			if _, ok := bt.ReplaceOrInsert(item); ok {
				panic("Insert replaced an item.")
			}
		})
	}
}

func (t *BTreeIndex[T]) update(old, new *T) {
	t.modify(func(bt *btree.BTreeG[*T]) {
		if t.include == nil || t.include(old) {
			bt.Delete(old)
		}
		if t.include == nil || t.include(new) {
			if _, ok := bt.ReplaceOrInsert(new); ok {
				panic("Update insert replaced an item.")
			}
		}
	})
}

func (t *BTreeIndex[T]) delete(item *T) {
	if t.include == nil || t.include(item) {
		t.modify(func(bt *btree.BTreeG[*T]) {
			bt.Delete(item)
		})
	}
}

// ----------------------------------------------------------------------------

func (t *BTreeIndex[T]) btree() *btree.BTreeG[*T] {
	return t.bt.Load().(*btree.BTreeG[*T])
}

func (t *BTreeIndex[T]) modify(mod func(clone *btree.BTreeG[*T])) {
	t.modLock.Lock()
	defer t.modLock.Unlock()
	clone := t.btree().Clone()
	mod(clone)
	t.bt.Store(clone)
}