package lattice

func (t *Tree) Walk(branch Branch, key Key) (nodeIdx uint32, pos int, found bool) {
	idx := t.Roots[branch]
	for {
		node := t.getNode(idx)
		k := node.KeyCount()
		p := 0
		for p < k && keyLess(node.Keys[p], key) {
			p++
		}
		if p < k && keyEqual(key, node.Keys[p]) {
			return idx, p, true
		}
		if node.IsLeaf() {
			return idx, p, false
		}
		child := node.Children[p]
		if child == NullIndex {
			return idx, p, false
		}
		idx = child
	}
}

func (t *Tree) Insert(branch Branch, key Key, rec Record) uint32 {
	root := t.Roots[branch]
	if t.getNode(root).KeyCount() == C {
		t.splitRoot(branch)
	}
	rIdx := t.allocRecord()
	r := t.getRecord(rIdx)
	*r = rec
	t.markRecDirty(rIdx)
	t.RecKey[rIdx] = key
	return t.insertNonFull(t.Roots[branch], branch, key, rIdx)
}

// InsertRec inserts rec at key in branch and returns the record index (not
// the node index). Use this when you need to index side-data by record index.
func (t *Tree) InsertRec(branch Branch, key Key, rec Record) uint32 {
	root := t.Roots[branch]
	if t.getNode(root).KeyCount() == C {
		t.splitRoot(branch)
	}
	rIdx := t.allocRecord()
	r := t.getRecord(rIdx)
	*r = rec
	// Always reset links to NullRec after copy.
	r.Link[0] = NullRec
	r.Link[1] = NullRec
	t.markRecDirty(rIdx)
	t.RecKey[rIdx] = key
	t.insertNonFull(t.Roots[branch], branch, key, rIdx)
	return rIdx
}

func (t *Tree) insertNonFull(idx uint32, branch Branch, key Key, recIdx uint32) uint32 {
	for {
		node := t.getNode(idx)
		k := node.KeyCount()
		p := 0
		for p < k && keyLess(node.Keys[p], key) {
			p++
		}
		if p < k && keyEqual(key, node.Keys[p]) {
			return idx
		}
		if node.IsLeaf() {
			for i := k; i > p; i-- {
				node.Keys[i] = node.Keys[i-1]
				node.RecPtrs[i] = node.RecPtrs[i-1]
			}
			node.Keys[p] = key
			node.RecPtrs[p] = recIdx
			node.SetCount(k + 1)
			node.SetBranch(branch)
			node.Mult = 1
			t.markNodeDirty(idx)
			return idx
		}
		childIdx := node.Children[p]
		if t.getNode(childIdx).KeyCount() == C {
			median := t.splitChild(idx, p)
			node = t.getNode(idx)
			if keyEqual(key, median) {
				for pp := 0; pp < node.KeyCount(); pp++ {
					if keyEqual(node.Keys[pp], key) {
						return idx
					}
				}
			}
			if keyLess(median, key) {
				p++
			}
			childIdx = t.getNode(idx).Children[p]
		}
		idx = childIdx
	}
}

func (t *Tree) splitChild(parentIdx uint32, childPos int) Key {
	fullIdx := t.getNode(parentIdx).Children[childPos]
	full := *t.getNode(fullIdx)

	// For C=8: median is at index C/2=4.
	// Left node gets keys [0..3] (4 keys), right gets keys [5..7] (3 keys).
	const mid = C / 2
	median := full.Keys[mid]
	medianRec := full.RecPtrs[mid]

	rightIdx := t.allocNode()
	right := t.getNode(rightIdx)
	rightCount := C - mid - 1 // keys after the median
	for i := 0; i < rightCount; i++ {
		right.Keys[i] = full.Keys[mid+1+i]
		right.RecPtrs[i] = full.RecPtrs[mid+1+i]
	}
	right.SetCount(rightCount)
	right.SetBranch(full.GetBranch())
	right.Mult = full.Mult
	if !full.IsLeaf() {
		for i := 0; i <= rightCount; i++ {
			right.Children[i] = full.Children[mid+1+i]
		}
	}
	t.markNodeDirty(rightIdx)

	left := t.getNode(fullIdx)
	leftCount := mid
	for i := 0; i < leftCount; i++ {
		left.Keys[i] = full.Keys[i]
		left.RecPtrs[i] = full.RecPtrs[i]
	}
	for i := leftCount; i < C; i++ {
		left.Keys[i] = KeyZero
		left.RecPtrs[i] = NullRec
	}
	left.SetCount(leftCount)
	if !full.IsLeaf() {
		for i := 0; i <= leftCount; i++ {
			left.Children[i] = full.Children[i]
		}
		for i := leftCount + 1; i <= C; i++ {
			left.Children[i] = NullIndex
		}
	} else {
		for i := range left.Children {
			left.Children[i] = NullIndex
		}
	}
	t.markNodeDirty(fullIdx)

	parent := t.getNode(parentIdx)
	pk := parent.KeyCount()
	for i := pk; i > childPos; i-- {
		parent.Keys[i] = parent.Keys[i-1]
		parent.RecPtrs[i] = parent.RecPtrs[i-1]
	}
	parent.Keys[childPos] = median
	parent.RecPtrs[childPos] = medianRec
	for i := pk + 1; i > childPos+1; i-- {
		parent.Children[i] = parent.Children[i-1]
	}
	parent.Children[childPos+1] = rightIdx
	parent.SetCount(pk + 1)
	t.markNodeDirty(parentIdx)

	return median
}

func (t *Tree) splitRoot(branch Branch) {
	rootIdx := t.Roots[branch]
	leftIdx := t.allocNode()
	*t.getNode(leftIdx) = *t.getNode(rootIdx)
	t.markNodeDirty(leftIdx)

	root := t.getNode(rootIdx)
	*root = Node{}
	initNodeChildren(root)
	root.Children[0] = leftIdx
	root.SetBranch(branch)
	root.Mult = 1
	t.markNodeDirty(rootIdx)

	t.splitChild(rootIdx, 0)
}

func (t *Tree) Lookup(branch Branch, key Key) *Record {
	idx, pos, found := t.Walk(branch, key)
	if !found {
		return nil
	}
	node := t.getNode(idx)
	if node.IsDeleted() {
		return nil
	}
	recIdx := node.RecPtrs[pos]
	if recIdx == NullRec {
		return nil
	}
	return t.getRecord(recIdx)
}

func (t *Tree) Delete(branch Branch, key Key) bool {
	idx, pos, found := t.Walk(branch, key)
	if !found {
		return false
	}
	node := t.getNode(idx)
	if node.IsDeleted() {
		return false
	}
	recIdx := node.RecPtrs[pos]
	if recIdx != NullRec {
		rec := t.getRecord(recIdx)
		for li := 0; li < 2; li++ {
			partnerRI := rec.Link[li]
			if partnerRI == NullRec {
				continue
			}
			pRec := t.getRecord(partnerRI)
			for pi := 0; pi < 2; pi++ {
				if pRec.Link[pi] == recIdx {
					pRec.Link[pi] = NullRec
					t.markRecDirty(partnerRI)
				}
			}
		}
		rec.Link[0] = NullRec
		rec.Link[1] = NullRec
		t.markRecDirty(recIdx)
	}
	delete(t.RecKey, recIdx)
	node.SetDeleted()
	node.Children[0] = t.FreeHead
	t.FreeHead = idx
	t.FreeCount++
	t.markNodeDirty(idx)
	return true
}

func (t *Tree) InsertTriple(nounKey, verbKey, modKey Key, nRec, vRec, mRec Record) {
	t.Insert(Bnoun, nounKey, nRec)
	t.Insert(Bverb, verbKey, vRec)
	t.Insert(Bmodifier, modKey, mRec)
	nri := t.LookupRecIdx(Bnoun, nounKey)
	vri := t.LookupRecIdx(Bverb, verbKey)
	mri := t.LookupRecIdx(Bmodifier, modKey)
	nr := t.getRecord(nri)
	nr.Link = [2]uint32{vri, mri}
	nr.Branch = uint8(Bnoun)
	t.markRecDirty(nri)
	vr := t.getRecord(vri)
	vr.Link = [2]uint32{nri, mri}
	vr.Branch = uint8(Bverb)
	t.markRecDirty(vri)
	mr := t.getRecord(mri)
	mr.Link = [2]uint32{nri, vri}
	mr.Branch = uint8(Bmodifier)
	t.markRecDirty(mri)
}