package search

import (
	"fmt"
	"codeberg.org/snonux/algorithms/ds"
)

type linkColor bool

const (
	red   linkColor = true
	black linkColor = false
)

type rbNode[K,V ds.Number] struct {
	key      K
	val      V
	color    linkColor
	capacity int
	left     *rbNode[K,V]
	right    *rbNode[K,V]
	// Just mark a node as deleted if deleted. Not fully implemented in lecture.
	deleted bool
}

func (n *rbNode[K,V]) String() string {
	recurse := func(n *rbNode[K,V]) string {
		if n == nil {
			return ""
		}
		return fmt.Sprintf("\n%s", n.String())
	}

	color := "red"
	if n.color == black {
		color = "black"
	}
	return fmt.Sprintf("rbNode[K,V]{%v;%v:%v,%s,%d,%s,%s}",
		n.deleted,
		n.key,
		n.val,
		color,
		n.capacity,
		recurse(n.left),
		recurse(n.right),
	)
}

func (n *rbNode[K,V]) isRed() bool {
	if n == nil {
		return false
	}
	return n.color == red
}

func (n *rbNode[K,V]) Capacity() int {
	if n == nil {
		return 0
	}
	return n.capacity
}

type RedBlackBST[K,V ds.Number] struct {
	root *rbNode[K,V]
	size int
}

func NewRedBlackBST[K,V ds.Number]() *RedBlackBST[K,V] {
	return &RedBlackBST[K,V]{}
}

func (t *RedBlackBST[K,V]) String() string {
	if t.Empty() {
		return fmt.Sprintf("RedBlackBST[K,V]{%d:%d}", t.Size(), t.Capacity())
	}

	return fmt.Sprintf("RedBlackBST[K,V]{%d:%d;%s}", t.Size(), t.Capacity(), t.root)
}

func (t *RedBlackBST[K,V]) Size() int {
	return t.size
}

func (t *RedBlackBST[K,V]) Empty() bool {
	return t.Size() == 0
}

func (t *RedBlackBST[K,V]) Capacity() int {
	if t.root == nil {
		return 0
	}
	return t.root.capacity
}

func (t *RedBlackBST[K,V]) Put(key K, val V) {
	t.root = t.put(t.root, key, val)
	t.root.color = black
}

func (t *RedBlackBST[K,V]) put(n *rbNode[K,V], key K, val V) *rbNode[K,V] {
	if n == nil {
		t.size++
		return &rbNode[K,V]{key, val, red, 1, nil, nil, false}
	}

	switch {
	case key < n.key:
		n.left = t.put(n.left, key, val)
	case key > n.key:
		n.right = t.put(n.right, key, val)
	default:
		if n.deleted {
			n.deleted = false
			t.size++
		}
		n.val = val
	}

	switch {
	case n.right.isRed() && !n.left.isRed():
		n = t.rotateLeft(n)
	case n.left != nil && n.left.isRed() && n.left.left.isRed():
		n = t.rotateRight(n)
	case n.left.isRed() && n.right.isRed():
		t.flipColors(n)
	}

	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()
	return n
}

func (t *RedBlackBST[K,V]) Get(key K) (V, error) {
	return t.get(t.root, key)
}

func (t *RedBlackBST[K,V]) get(n *rbNode[K,V], key K) (V, error) {
	if n == nil {
		return 0, NotFound
	}

	switch {
	case key < n.key:
		return t.get(n.left, key)
	case key > n.key:
		return t.get(n.right, key)
	default:
		if n.deleted {
			return 0, NotFound
		}
		return n.val, nil
	}
}

func (t *RedBlackBST[K,V]) Del(key K) (V, error) {
	return t.del(t.root, key)
}

func (t *RedBlackBST[K,V]) del(n *rbNode[K,V], key K) (V, error) {
	if n == nil {
		return 0, NotFound
	}

	switch {
	case key < n.key:
		return t.del(n.left, key)
	case key > n.key:
		return t.del(n.right, key)
	default:
		if n.deleted {
			return 0, NotFound
		}
		t.size--
		n.deleted = true
		val := n.val
		n.val = 0 
		return val, nil
	}
}

func (t *RedBlackBST[K,V]) rotateLeft(n *rbNode[K,V]) *rbNode[K,V] {
	x := n.right
	n.right = x.left
	x.left = n

	x.color = n.color
	n.color = red

	x.capacity = n.capacity
	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()

	return x
}

func (t *RedBlackBST[K,V]) rotateRight(n *rbNode[K,V]) *rbNode[K,V] {
	x := n.left
	n.left = x.right
	x.right = n

	x.color = n.color
	n.color = red

	x.capacity = n.capacity
	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()

	return x
}

func (t *RedBlackBST[K,V]) flipColors(n *rbNode[K,V]) {
	n.color = red
	n.left.color = black
	n.right.color = black
}