package mls

// MLS tree math (RFC 9420 Appendix C).
// Array-based complete balanced binary tree representation.

// numLeaves exposes operations on a tree with a given number of leaves.
// RFC 9420 §7 bounds trees to 2^32 leaves, so width is always representable in uint32.
type numLeaves uint32

func numLeavesFromWidth(w uint32) (v numLeaves) {
	if w == 0 {
		return 0
	}
	return numLeaves((w-1)/2 + 1)
}

// width returns the number of nodes (minimum array length).
func (n *numLeaves) width() (nv uint32) {
	if *n == 0 {
		return 0
	}
	return 2*(uint32(*n)-1) + 1
}

// root returns the index of the root node.
func (n *numLeaves) root() (v nodeIndex) {
	return nodeIndex((1 << log2(n.width())) - 1)
}

// parent returns the parent node index. Returns false for the root.
func (n *numLeaves) parent(x nodeIndex) (v nodeIndex, ok bool) {
	if x == n.root() {
		return 0, false
	}
	lvl := nodeIndex(x.level())
	b := (x >> (lvl + 1)) & 1
	p := (x | (1 << lvl)) ^ (b << (lvl + 1))
	return p, true
}

// sibling returns the other child of the node's parent.
func (n *numLeaves) sibling(x nodeIndex) (v nodeIndex, ok2 bool) {
	p, ok := n.parent(x)
	if !ok {
		return 0, false
	}
	if x < p {
		return p.right()
	} else {
		return p.left()
	}
}

// directPath returns the path from x to the root (excluding x, excluding root).
func (n *numLeaves) directPath(x nodeIndex) (ss []nodeIndex) {
	var path []nodeIndex
	for {
		p, ok := n.parent(x)
		if !ok {
			break
		}
		path = push(path, p)
		x = p
	}
	return path
}

// copath returns the copath (siblings along the direct path).
func (n *numLeaves) copath(x nodeIndex) (ss []nodeIndex) {
	path := n.directPath(x)
	if len(path) == 0 {
		return nil
	}
	path = []nodeIndex{x} | path
	path = path[:len(path)-1]

	var copath []nodeIndex
	for _, y := range path {
		s, ok := n.sibling(y)
		if !ok {
			panic("unreachable")
		}
		copath = push(copath, s)
	}
	return copath
}

// --- nodeIndex ---

// nodeIndex addresses any node in the tree (leaf or parent).
// Internally 2*leafIndex for leaves, odd for parents; RFC 9420 keeps this in u32.
type nodeIndex uint32

func (x *nodeIndex) isLeaf() (ok bool) {
	return (*x)%2 == 0
}

func (x *nodeIndex) leafIndex() (v leafIndex, ok bool) {
	if !x.isLeaf() {
		return 0, false
	}
	return leafIndex(*x) >> 1, true
}

func (x *nodeIndex) left() (v nodeIndex, ok bool) {
	lvl := x.level()
	if lvl == 0 {
		return 0, false
	}
	return (*x) ^ (1 << (nodeIndex(lvl) - 1)), true
}

func (x *nodeIndex) right() (v nodeIndex, ok bool) {
	lvl := x.level()
	if lvl == 0 {
		return 0, false
	}
	return (*x) ^ (3 << (nodeIndex(lvl) - 1)), true
}

func (x *nodeIndex) children() (left, right nodeIndex, ok bool) {
	l, hasLeft := x.left()
	if !hasLeft {
		return 0, 0, false
	}
	r, _ := x.right()
	return l, r, true
}

func (x *nodeIndex) level() (n uint32) {
	if (*x)&1 == 0 {
		return 0
	}
	lvl := uint32(0)
	for ((*x)>>lvl)&1 == 1 {
		lvl++
	}
	return lvl
}

// commonAncestor returns the lowest node in both direct paths.
func commonAncestor(x, y nodeIndex) (v nodeIndex) {
	lx, ly := x.level()+1, y.level()+1
	if lx <= ly && x>>ly == y>>ly {
		return y
	} else if ly <= lx && x>>lx == y>>lx {
		return x
	}

	xn, yn := x, y
	k := uint32(0)
	for xn != yn {
		xn, yn = xn>>1, yn>>1
		k++
	}
	return (xn << k) + (1 << (k - 1)) - 1
}

// --- leafIndex ---

// leafIndex addresses a member leaf in the ratchet tree.
// Wire-encoded as uint32 (RFC 9420 §7).
type leafIndex uint32

func (li *leafIndex) nodeIndex() (v nodeIndex) {
	return nodeIndex(*li * 2)
}

// --- Helpers ---

func log2(x uint32) (n uint32) {
	if x == 0 {
		return 0
	}
	k := uint32(0)
	for x>>k > 0 {
		k++
	}
	return k - 1
}

func isPowerOf2(x uint32) (ok bool) {
	return x != 0 && x&(x-1) == 0
}