package alg

import "strings"

type IndexKey interface {
	Compare(other IndexKey) int
}

type SortOrder struct {
	Field string
	Desc  bool
}

type IndexEntry struct {
	Key   IndexKey
	Value interface{}
}

const BTreeMinDegree = 2

type BTreeNode struct {
	Keys   []IndexKey
	Values []interface{} // 叶子节点存储数据
	Childs []*BTreeNode  // 非叶子节点存储子节点
	IsLeaf bool
	Next   *BTreeNode // 叶子节点之间的链表指针
	Prev   *BTreeNode
}

type BTree struct {
	Root *BTreeNode
}

func NewBTreeNode(leaf bool) *BTreeNode {
	return &BTreeNode{
		Keys:   make([]IndexKey, 0),
		Values: make([]interface{}, 0),
		Childs: make([]*BTreeNode, 0),
		IsLeaf: leaf,
	}
}

func NewBTree() *BTree {
	return &BTree{
		Root: NewBTreeNode(true),
	}
}

func (bt *BTree) Insert(key IndexKey, value interface{}) {
	root := bt.Root
	// 如果根节点已满，需要分裂
	if len(root.Keys) == 2*BTreeMinDegree-1 {
		newRoot := NewBTreeNode(false)
		newRoot.Childs = append(newRoot.Childs, root)
		splitChild(newRoot, 0)
		bt.Root = newRoot
	}
	insertNonFull(bt.Root, key, value)
}

// 分裂子节点
func splitChild(parent *BTreeNode, index int) {
	minDegree := BTreeMinDegree
	fullNode := parent.Childs[index]
	// 创建新节点
	newNode := NewBTreeNode(fullNode.IsLeaf)
	// 中间键位置
	mid := minDegree - 1

	// 拆分键和值
	newNode.Keys = fullNode.Keys[mid:]
	fullNode.Keys = fullNode.Keys[:mid]

	if !fullNode.IsLeaf {
		newNode.Childs = fullNode.Childs[mid:]
		fullNode.Childs = fullNode.Childs[:mid]
	} else {
		// 叶子节点需要维护链表
		newNode.Values = fullNode.Values[mid:]
		fullNode.Values = fullNode.Values[:mid]
		newNode.Next = fullNode.Next
		newNode.Prev = fullNode
		if fullNode.Next != nil {
			fullNode.Next.Prev = newNode
		}
		fullNode.Next = newNode
	}

	// 将中间键提升到父节点
	parent.Keys = append(parent.Keys[:index+1], append([]IndexKey{fullNode.Keys[mid]}, parent.Keys[index+1:]...)...)
	fullNode.Keys = fullNode.Keys[:mid]
}

// 插入非满节点
func insertNonFull(node *BTreeNode, key IndexKey, value interface{}) {
	i := len(node.Keys) - 1

	if node.IsLeaf {
		// 在叶子节点插入
		for i >= 0 && key.Compare(node.Keys[i]) < 0 {
			i--
		}
		// 插入键值
		if i >= 0 {
			node.Keys = append(node.Keys[:i+1], append([]IndexKey{key}, node.Keys[i+1:]...)...)
			node.Values = append(node.Values[:i+1], append([]interface{}{value}, node.Values[i+1:]...)...)
		} else {
			node.Keys = append([]IndexKey{key}, node.Keys...)
			node.Values = append([]interface{}{value}, node.Values...)
		}
	} else {
		// 找到子节点
		for i >= 0 && key.Compare(node.Keys[i]) < 0 {
			i--
		}
		i++

		// 如果子节点已满，需要分裂
		child := node.Childs[i]
		if len(child.Keys) == 2*BTreeMinDegree-1 {
			splitChild(node, i)
			if key.Compare(node.Keys[i]) > 0 {
				i++ // 选择正确的位置
			}
		}
		insertNonFull(node.Childs[i], key, value)
	}
}

func (bt *BTree) Search(key IndexKey) (interface{}, bool) {
	node := bt.Root
	for node != nil {
		// 在当前节点查找位置
		i := 0
		for i < len(node.Keys) && key.Compare(node.Keys[i]) > 0 {
			i++
		}

		if node.IsLeaf {
			// 在叶子节点中查找精确匹配
			if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 {
				return node.Values[i], true
			}
			return nil, false
		} else {
			// 继续查找子节点
			node = node.Childs[i]
		}
	}
	return nil, false
}

func (bt *BTree) Delete(key IndexKey) {
	deleteFromNode(bt.Root, key)
}

// 从节点删除指定键
func deleteFromNode(node *BTreeNode, key IndexKey) {
	minDegree := BTreeMinDegree
	// 寻找键的位置
	i := 0
	for i < len(node.Keys) && key.Compare(node.Keys[i]) > 0 {
		i++
	}

	if node.IsLeaf {
		// 叶子节点直接删除
		if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 {
			// 删除键值对
			node.Keys = append(node.Keys[:i], node.Keys[i+1:]...)
			node.Values = append(node.Values[:i], node.Values[i+1:]...)
		}
	} else {
		// 非叶子节点
		if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 {
			// 找到键，需要替换为后继值
			child := node.Childs[i]
			sibling := node.Childs[i+1]

			if len(child.Keys) >= minDegree {
				// 前驱节点存在，替换为前驱
				predecessor := getPredecessor(child)
				node.Keys[i] = predecessor
				deleteFromNode(child, predecessor)
			} else if len(sibling.Keys) >= minDegree {
				// 后继节点存在，替换为后继
				successor := getSuccessor(sibling)
				node.Keys[i] = successor
				deleteFromNode(sibling, successor)
			} else {
				// 合并节点
				mergeNodes(node, i)
				deleteFromNode(child, key)
			}
		} else {
			// 继续删除子节点
			child := node.Childs[i]
			if len(child.Keys) >= minDegree {
				// 子节点足够大，继续删除
				deleteFromNode(child, key)
			} else {
				// 需要重新平衡
				if i > 0 && len(node.Childs[i-1].Keys) >= minDegree {
					// 从左兄弟借
					borrowFromLeft(node, i)
				} else if i < len(node.Childs)-1 && len(node.Childs[i+1].Keys) >= minDegree {
					// 从右兄弟借
					borrowFromRight(node, i)
				} else {
					// 合并节点
					if i > 0 {
						mergeNodes(node, i-1)
					} else {
						mergeNodes(node, i)
					}
				}
				deleteFromNode(child, key)
			}
		}
	}
}

// 获取前驱节点
func getPredecessor(node *BTreeNode) IndexKey {
	for !node.IsLeaf {
		node = node.Childs[len(node.Childs)-1]
	}
	return node.Keys[len(node.Keys)-1]
}

// 获取后继节点
func getSuccessor(node *BTreeNode) IndexKey {
	for !node.IsLeaf {
		node = node.Childs[0]
	}
	return node.Keys[0]
}

// 合并节点
func mergeNodes(parent *BTreeNode, index int) {
	left := parent.Childs[index]
	right := parent.Childs[index+1]

	// 将左节点和右节点合并
	left.Keys = append(left.Keys, parent.Keys[index])
	left.Keys = append(left.Keys, right.Keys...)

	if left.IsLeaf {
		left.Values = append(left.Values, right.Values...)
		left.Next = right.Next
		if right.Next != nil {
			right.Next.Prev = left
		}
	} else {
		left.Childs = append(left.Childs, right.Childs...)
	}

	// 更新父节点
	parent.Keys = append(parent.Keys[:index], parent.Keys[index+1:]...)
	parent.Childs = append(parent.Childs[:index+1], parent.Childs[index+2:]...)
}

// 从左兄弟借
func borrowFromLeft(parent *BTreeNode, index int) {
	left := parent.Childs[index-1]
	current := parent.Childs[index]

	// 移动键
	current.Keys = append([]IndexKey{parent.Keys[index-1]}, current.Keys...)
	if !left.IsLeaf {
		current.Childs = append([]*BTreeNode{left.Childs[len(left.Childs)-1]}, current.Childs...)
		left.Childs = left.Childs[:len(left.Childs)-1]
	}

	// 更新父节点
	parent.Keys[index-1] = left.Keys[len(left.Keys)-1]
	left.Keys = left.Keys[:len(left.Keys)-1]
}

// 从右兄弟借
func borrowFromRight(parent *BTreeNode, index int) {
	current := parent.Childs[index]
	right := parent.Childs[index+1]

	// 移动键
	current.Keys = append(current.Keys, parent.Keys[index])
	parent.Keys[index] = right.Keys[0]

	if !current.IsLeaf {
		current.Childs = append(current.Childs, right.Childs[0])
		right.Childs = right.Childs[1:]
	}

	right.Keys = right.Keys[1:]
}

func (bt *BTree) Traverse(cb func(key IndexKey) bool) []interface{} {
	result := make([]interface{}, 0)
	// 找到最左边的叶子节点
	node := bt.Root
	for node != nil && !node.IsLeaf {
		node = node.Childs[0]
	}

	// 遍历所有叶子节点
	for node != nil {
		for i, key := range node.Keys {
			if cb(key) {
				result = append(result, node.Values[i])
			} else {
				return result
			}
		}
		node = node.Next
	}
	return result
}

type MultiFieldKey struct {
	Fields   []string
	SortDirs []bool // true 表示降序
	Values   map[string]interface{}
}

func (m MultiFieldKey) Compare(other IndexKey) int {
	o := other.(MultiFieldKey)
	for i, field := range m.Fields {
		v1 := m.Values[field]
		v2 := o.Values[field]

		var cmp int
		switch v1.(type) {
		case string:
			cmp = strings.Compare(v1.(string), v2.(string))
		case float64:
			if v1.(float64) < v2.(float64) {
				cmp = -1
			} else if v1.(float64) > v2.(float64) {
				cmp = 1
			} else {
				cmp = 0
			}
		default:
			panic("Unsupported type")
		}

		if cmp != 0 {
			if m.SortDirs[i] {
				return -cmp
			}
			return cmp
		}
	}
	return 0
}