// Copyright 2018 The Cockroach Authors.
// Copyright 2021 Andrew Werner.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

package abstract

import (
	"fmt"
	"strings"
	"sync/atomic"
)

// Node represents a node in the tree.
type Node[K, V, A any] struct {
	ref      int32
	count    int16
	aug      A
	keys     [MaxEntries]K
	values   [MaxEntries]V
	children *[MaxEntries + 1]*Node[K, V, A]
}

type interiorNode[K, V, A any] struct {
	Node[K, V, A]
	children [MaxEntries + 1]*Node[K, V, A]
}

func (n *Node[K, V, A]) GetA() *A {
	return &n.aug
}

func (n *Node[K, V, A]) IsLeaf() bool {
	return n.children == nil
}

func (n *Node[K, V, A]) Count() int16 {
	return n.count
}

func (n *Node[K, V, A]) GetKey(i int16) K {
	return n.keys[i]
}

func (n *Node[K, V, A]) GetChild(i int16) *A {
	if !n.IsLeaf() && n.children[i] != nil {
		return &n.children[i].aug
	}
	return nil
}

// mut creates and returns a mutable node reference. If the node is not shared
// with any other trees then it can be modified in place. Otherwise, it must be
// cloned to ensure unique ownership. In this way, we enforce a copy-on-write
// policy which transparently incorporates the idea of local mutations, like
// Clojure's transients or Haskell's ST monad, where nodes are only copied
// during the first time that they are modified between Clone operations.
//
// When a node is cloned, the provided pointer will be redirected to the new
// mutable node.
func mut[K, V, A any](
	np *nodePool[K, V, A],
	n **Node[K, V, A],
) *Node[K, V, A] {
	if atomic.LoadInt32(&(*n).ref) == 1 {
		// Exclusive ownership. Can mutate in place.
		return *n
	}
	// If we do not have unique ownership over the node then we
	// clone it to gain unique ownership. After doing so, we can
	// release our reference to the old node. We pass recursive
	// as true because even though we just observed the node's
	// reference count to be greater than 1, we might be racing
	// with another call to decRef on this node.
	c := (*n).clone(np)
	(*n).decRef(np, true /* recursive */)
	*n = c
	return *n
}

// incRef acquires a reference to the node.
func (n *Node[K, V, A]) incRef() {
	atomic.AddInt32(&n.ref, 1)
}

// decRef releases a reference to the node. If requested, the method
// will recurse into child nodes and decrease their refcounts as well.
func (n *Node[K, V, A]) decRef(
	np *nodePool[K, V, A], recursive bool,
) {
	if atomic.AddInt32(&n.ref, -1) > 0 {
		// Other references remain. Can't free.
		return
	}
	// Clear and release node into memory pool.
	if n.IsLeaf() {
		np.putLeafNode(n)
	} else {
		// Release child references first, if requested.
		if recursive {
			for i := int16(0); i <= n.count; i++ {
				n.children[i].decRef(np, true /* recursive */)
			}
		}
		np.putInteriorNode(n)
	}
}

// clone creates a clone of the receiver with a single reference count.
func (n *Node[K, V, A]) clone(
	np *nodePool[K, V, A],
) *Node[K, V, A] {
	var c *Node[K, V, A]
	if n.IsLeaf() {
		c = np.getLeafNode()
	} else {
		c = np.getInteriorNode()
	}
	// NB: copy field-by-field without touching N.N.ref to avoid
	// triggering the race detector and looking like a data race.
	c.count = n.count
	n.aug = c.aug
	c.keys = n.keys
	if !c.IsLeaf() {
		// Copy children and increase each refcount.
		*c.children = *n.children
		for i := int16(0); i <= c.count; i++ {
			c.children[i].incRef()
		}
	}
	return c
}

func (n *Node[K, V, A]) insertAt(index int, item K, value V, nd *Node[K, V, A]) {
	if index < int(n.count) {
		copy(n.keys[index+1:n.count+1], n.keys[index:n.count])
		copy(n.values[index+1:n.count+1], n.values[index:n.count])
		if !n.IsLeaf() {
			copy(n.children[index+2:n.count+2], n.children[index+1:n.count+1])
		}
	}
	n.keys[index] = item
	n.values[index] = value
	if !n.IsLeaf() {
		n.children[index+1] = nd
	}
	n.count++
}

func (n *Node[K, V, A]) pushBack(item K, value V, nd *Node[K, V, A]) {
	n.keys[n.count] = item
	n.values[n.count] = value
	if !n.IsLeaf() {
		n.children[n.count+1] = nd
	}
	n.count++
}

func (n *Node[K, V, A]) pushFront(item K, value V, nd *Node[K, V, A]) {
	if !n.IsLeaf() {
		copy(n.children[1:n.count+2], n.children[:n.count+1])
		n.children[0] = nd
	}
	copy(n.keys[1:n.count+1], n.keys[:n.count])
	copy(n.values[1:n.count+1], n.values[:n.count])
	n.keys[0] = item
	n.values[0] = value
	n.count++
}

// removeAt removes a value at a given index, pulling all subsequent values
// back.
func (n *Node[K, V, A]) removeAt(index int) (K, V, *Node[K, V, A]) {
	var child *Node[K, V, A]
	if !n.IsLeaf() {
		child = n.children[index+1]
		copy(n.children[index+1:n.count], n.children[index+2:n.count+1])
		n.children[n.count] = nil
	}
	n.count--
	outK := n.keys[index]
	outV := n.values[index]
	copy(n.keys[index:n.count], n.keys[index+1:n.count+1])
	copy(n.values[index:n.count], n.values[index+1:n.count+1])
	var rk K
	var rv V
	n.keys[n.count] = rk
	n.values[n.count] = rv
	return outK, outV, child
}

// popBack removes and returns the last element in the list.
func (n *Node[K, V, A]) popBack() (K, V, *Node[K, V, A]) {
	n.count--
	outK := n.keys[n.count]
	outV := n.values[n.count]
	var rK K
	var rV V
	n.keys[n.count] = rK
	n.values[n.count] = rV
	if n.IsLeaf() {
		return outK, outV, nil
	}
	child := n.children[n.count+1]
	n.children[n.count+1] = nil
	return outK, outV, child
}

// popFront removes and returns the first element in the list.
func (n *Node[K, V, A]) popFront() (K, V, *Node[K, V, A]) {
	n.count--
	var child *Node[K, V, A]
	if !n.IsLeaf() {
		child = n.children[0]
		copy(n.children[:n.count+1], n.children[1:n.count+2])
		n.children[n.count+1] = nil
	}
	outK := n.keys[0]
	outV := n.values[0]
	copy(n.keys[:n.count], n.keys[1:n.count+1])
	copy(n.values[:n.count], n.values[1:n.count+1])
	var rK K
	var rV V
	n.keys[n.count] = rK
	n.values[n.count] = rV
	return outK, outV, child
}

// find returns the index where the given item should be inserted into this
// list. 'found' is true if the item already exists in the list at the given
// index.
func (n *Node[K, V, A]) find(cmp func(K, K) int, item K) (index int, found bool) {
	// Logic copied from sort.Search. Inlining this gave
	// an 11% speedup on BenchmarkBTreeDeleteInsert.
	i, j := 0, int(n.count)
	for i < j {
		h := int(uint(i+j) >> 1) // avoid overflow when computing h
		// i ≤ h < j
		c := cmp(item, n.keys[h])
		if c < 0 {
			j = h
		} else if c > 0 {
			i = h + 1
		} else {
			return h, true
		}
	}
	return i, false
}

// split splits the given node at the given index. The current node shrinks,
// and this function returns the item that existed at that index and a new
// node containing all keys/children after it.
//
// Before:
//
//          +-----------+
//          |   x y z   |
//          +--/-/-\-\--+
//
// After:
//
//          +-----------+
//          |     y     |
//          +----/-\----+
//              /   \
//             v     v
// +-----------+     +-----------+
// |         x |     | z         |
// +-----------+     +-----------+
//
func (n *Node[K, V, A]) split(cfg *config[K, V, A], i int) (K, V, *Node[K, V, A]) {
	outK := n.keys[i]
	outV := n.values[i]
	var next *Node[K, V, A]
	if n.IsLeaf() {
		next = cfg.np.getLeafNode()
	} else {
		next = cfg.np.getInteriorNode()
	}
	next.count = n.count - int16(i+1)
	copy(next.keys[:], n.keys[i+1:n.count])
	copy(next.values[:], n.values[i+1:n.count])
	var rK K
	var rV V
	for j := int16(i); j < n.count; j++ {
		n.keys[j] = rK
		n.values[j] = rV
	}
	if !n.IsLeaf() {
		copy(next.children[:], n.children[i+1:n.count+1])
		for j := int16(i + 1); j <= n.count; j++ {
			n.children[j] = nil
		}
	}
	n.count = int16(i)
	next.update(&cfg.Config)
	n.updateOn(&cfg.Config, Split, outK, next)
	return outK, outV, next
}

func (n *Node[K, V, A]) update(cfg *Config[K, V, A]) bool {
	return n.updateWithMeta(cfg, UpdateInfo[K, A]{})
}

func (n *Node[K, V, A]) updateWithMeta(cfg *Config[K, V, A], md UpdateInfo[K, A]) bool {
	if cfg.Updater == nil {
		return false
	}
	return cfg.Updater.Update(n, md)
}

func (n *Node[K, V, A]) updateOn(cfg *Config[K, V, A], action Action, k K, affected *Node[K, V, A]) bool {
	if cfg.Updater == nil {
		return false
	}
	var a *A
	if affected != nil {
		a = &affected.aug
	}
	return n.updateWithMeta(cfg, UpdateInfo[K, A]{
		Action:        action,
		RelevantKey:   k,
		ModifiedOther: a,
	})
}

// insert inserts an item into the suAugBTree rooted at this node, making sure no
// nodes in the suAugBTree exceed MaxEntries keys. Returns true if an existing item
// was replaced and false if an item was inserted. Also returns whether the
// node's upper bound changes.
func (n *Node[K, V, A]) insert(cfg *config[K, V, A], item K, value V) (replacedK K, replacedV V, replaced, newBound bool) {
	i, found := n.find(cfg.cmp, item)
	if found {
		replacedV = n.values[i]
		replacedK = n.keys[i]
		n.keys[i] = item
		n.values[i] = value
		return replacedK, replacedV, true, false
	}
	if n.IsLeaf() {
		n.insertAt(i, item, value, nil)
		return replacedK, replacedV, false, n.updateOn(&cfg.Config, Insertion, item, nil)
	}
	if n.children[i].count >= MaxEntries {
		splitLK, splitLV, splitNode := mut(cfg.np, &n.children[i]).
			split(cfg, MaxEntries/2)
		n.insertAt(i, splitLK, splitLV, splitNode)
		if c := cfg.cmp(item, n.keys[i]); c < 0 {
			// no change, we want first split node
		} else if c > 0 {
			i++ // we want second split node
		} else {
			// TODO(ajwerner): add something to the augmentation api to
			// deal with replacement.
			replacedV = n.values[i]
			replacedK = n.keys[i]
			n.keys[i] = item
			n.values[i] = value
			return replacedK, replacedV, true, false
		}
	}
	replacedK, replacedV, replaced, newBound =
		mut(cfg.np, &n.children[i]).insert(cfg, item, value)
	if newBound {
		newBound = n.updateOn(&cfg.Config, Insertion, item, nil)
	}
	return replacedK, replacedV, replaced, newBound
}

// removeMax removes and returns the maximum item from the suAugBTree rooted at
// this node.
func (n *Node[K, V, A]) removeMax(cfg *config[K, V, A]) (K, V) {
	if n.IsLeaf() {
		n.count--
		outK := n.keys[n.count]
		outV := n.values[n.count]
		var rK K
		var rV V
		n.keys[n.count] = rK
		n.values[n.count] = rV
		n.updateOn(&cfg.Config, Removal, outK, nil)
		return outK, outV
	}
	// Recurse into max child.
	i := int(n.count)
	if n.children[i].count <= MinEntries {
		// Child not large enough to remove from.
		n.rebalanceOrMerge(cfg, i)
		return n.removeMax(cfg) // redo
	}
	child := mut(cfg.np, &n.children[i])
	outK, outV := child.removeMax(cfg)
	n.updateOn(&cfg.Config, Removal, outK, nil)
	return outK, outV
}

// rebalanceOrMerge grows child 'i' to ensure it has sufficient room to remove
// an item from it while keeping it at or above MinItems.
func (n *Node[K, V, A]) rebalanceOrMerge(
	cfg *config[K, V, A], i int,
) {
	switch {
	case i > 0 && n.children[i-1].count > MinEntries:
		// Rebalance from left sibling.
		//
		//          +-----------+
		//          |     y     |
		//          +----/-\----+
		//              /   \
		//             v     v
		// +-----------+     +-----------+
		// |         x |     |           |
		// +----------\+     +-----------+
		//             \
		//              v
		//              a
		//
		// After:
		//
		//          +-----------+
		//          |     x     |
		//          +----/-\----+
		//              /   \
		//             v     v
		// +-----------+     +-----------+
		// |           |     | y         |
		// +-----------+     +/----------+
		//                   /
		//                  v
		//                  a
		//
		left := mut(cfg.np, &n.children[i-1])
		child := mut(cfg.np, &n.children[i])
		xLaK, xLaV, grandChild := left.popBack()
		yLaK, yLaV := n.keys[i-1], n.values[i-1]
		child.pushFront(yLaK, yLaV, grandChild)
		n.keys[i-1], n.values[i-1] = xLaK, xLaV
		left.updateOn(&cfg.Config, Removal, xLaK, grandChild)
		child.updateOn(&cfg.Config, Insertion, yLaK, grandChild)

	case i < int(n.count) && n.children[i+1].count > MinEntries:
		// Rebalance from right sibling.
		//
		//          +-----------+
		//          |     y     |
		//          +----/-\----+
		//              /   \
		//             v     v
		// +-----------+     +-----------+
		// |           |     | x         |
		// +-----------+     +/----------+
		//                   /
		//                  v
		//                  a
		//
		// After:
		//
		//          +-----------+
		//          |     x     |
		//          +----/-\----+
		//              /   \
		//             v     v
		// +-----------+     +-----------+
		// |         y |     |           |
		// +----------\+     +-----------+
		//             \
		//              v
		//              a
		//
		right := mut(cfg.np, &n.children[i+1])
		child := mut(cfg.np, &n.children[i])
		xLaK, xLaV, grandChild := right.popFront()
		yLaK, yLaV := n.keys[i], n.values[i]
		child.pushBack(yLaK, yLaV, grandChild)
		n.keys[i], n.values[i] = xLaK, xLaV
		right.updateOn(&cfg.Config, Removal, xLaK, grandChild)
		child.updateOn(&cfg.Config, Insertion, yLaK, grandChild)

	default:
		// Merge with either the left or right sibling.
		//
		//          +-----------+
		//          |   u y v   |
		//          +----/-\----+
		//              /   \
		//             v     v
		// +-----------+     +-----------+
		// |         x |     | z         |
		// +-----------+     +-----------+
		//
		// After:
		//
		//          +-----------+
		//          |    u v    |
		//          +-----|-----+
		//                |
		//                v
		//          +-----------+
		//          |   x y z   |
		//          +-----------+
		//
		if i >= int(n.count) {
			i = int(n.count - 1)
		}
		child := mut(cfg.np, &n.children[i])
		// Make mergeChild mutable, bumping the refcounts on its children if necessary.
		_ = mut(cfg.np, &n.children[i+1])
		mergeLaK, mergeLaV, mergeChild := n.removeAt(i)
		child.keys[child.count] = mergeLaK
		child.values[child.count] = mergeLaV
		copy(child.keys[child.count+1:], mergeChild.keys[:mergeChild.count])
		copy(child.values[child.count+1:], mergeChild.values[:mergeChild.count])
		if !child.IsLeaf() {
			copy(child.children[child.count+1:], mergeChild.children[:mergeChild.count+1])
		}
		child.count += mergeChild.count + 1

		child.updateOn(&cfg.Config, Insertion, mergeLaK, mergeChild)
		mergeChild.decRef(cfg.np, false /* recursive */)
	}
}

// remove removes an item from the suAugBTree rooted at this node. Returns the item
// that was removed or nil if no matching item was found. Also returns whether
// the node's upper bound changes.
func (n *Node[K, V, A]) remove(
	cfg *config[K, V, A], item K,
) (outK K, outV V, found, newBound bool) {
	i, found := n.find(cfg.cmp, item)
	if n.IsLeaf() {
		if found {
			outK, outV, _ = n.removeAt(i)
			return outK, outV, true, n.updateOn(&cfg.Config, Removal, outK, nil)
		}
		var rK K
		var rV V
		return rK, rV, false, false
	}
	if n.children[i].count <= MinEntries {
		// Child not large enough to remove from.
		n.rebalanceOrMerge(cfg, i)
		return n.remove(cfg, item) // redo
	}
	child := mut(cfg.np, &n.children[i])
	if found {
		// Replace the item being removed with the max item in our left child.
		outK = n.keys[i]
		outV = n.values[i]
		n.keys[i], n.values[i] = child.removeMax(cfg)
		return outK, outV, true, n.updateOn(&cfg.Config, Removal, outK, nil)
	}
	// Latch is not in this node and child is large enough to remove from.
	outK, outV, found, newBound = child.remove(cfg, item)
	if newBound {
		newBound = n.updateOn(&cfg.Config, Removal, outK, nil)
	}
	return outK, outV, found, newBound
}

func (n *Node[K, V, A]) writeString(b *strings.Builder) {
	if n.IsLeaf() {
		for i := int16(0); i < n.count; i++ {
			if i != 0 {
				b.WriteString(",")
			}
			fmt.Fprintf(b, "%v:%v", n.keys[i], n.values[i])
		}
		return
	}
	for i := int16(0); i <= n.count; i++ {
		b.WriteString("(")
		n.children[i].writeString(b)
		b.WriteString(")")
		if i < n.count {
			fmt.Fprintf(b, "%v:%v", n.keys[i], n.values[i])
		}
	}
}