package gnarlring

import "io"

// TreeNode holds the key material and state for one tree node.
// Leaf nodes hold only NTRU keypairs. Internal nodes hold NTRU keypairs
// and manage child epochs.
type TreeNode struct {
	PK *NTRUPublicKey
	SK *NTRUPrivateKey

	// Children is a set of child commitments (from lower tree levels).
	// For leaf nodes, Children is nil.
	Children *AggregatedCommitment

	// Epoch counter for this node's group state.
	Epoch uint64
}

// NewLeafNode creates a leaf tree node.
func NewLeafNode() *TreeNode {
	pk, sk := NTRUKeyGen()
	return &TreeNode{PK: pk, SK: sk}
}

// NewInternalNode creates an internal tree node that coordinates children.
func NewInternalNode() *TreeNode {
	pk, sk := NTRUKeyGen()
	return &TreeNode{
		PK:       pk,
		SK:       sk,
		Children: NewAggregatedCommitment(),
	}
}

// AddChild inserts a child commitment into this node.
func (n *TreeNode) AddChild(index int, childPK *NTRUPublicKey, rng io.Reader) bool {
	if n.Children == nil {
		return false
	}
	w := NewChildCommitment(index, childPK, rng)
	return n.Children.Add(w)
}

// SignChildCommitments produces an NTRU signature binding all child
// commitments under this node's public key. Returns the epoch signature.
// The node acts as a coordinator for its children.
func (n *TreeNode) SignChildCommitments(msg []byte) *NTRUSignature {
	if n.Children == nil || !n.Children.IsComplete() {
		return nil
	}
	target := n.Children.Target(msg, n.Epoch)
	return NTRUSignTarget(n.SK, target, nil)
}

// VerifyChildCommitments verifies the coordinator's signature on child
// commitments under this node's public key.
func (n *TreeNode) VerifyChildCommitments(msg []byte, sig *NTRUSignature) bool {
	if n.Children == nil || !n.Children.IsComplete() {
		return false
	}
	target := n.Children.Target(msg, n.Epoch)
	return NTRUVerifyTarget(n.PK, target, sig)
}

// RecursiveTree is a multi-level tree of NTRU keypairs.
// Level 0: root coordinator. Level 1..h: sub-coordinators. Level h+1: leaves.
type RecursiveTree struct {
	Root *TreeNode
	Levels [][]*TreeNode // Levels[0] = [root], Levels[1] = [27 coordinators], ...
}

// BuildRecursiveTree creates a tree of depth h with full branching factor N.
// depth=1: root + 27 leaves (total 28 nodes)
// depth=2: root + 27 coordinators + 27*27 = 729 leaves
func BuildRecursiveTree(depth int) *RecursiveTree {
	if depth < 1 {
		return nil
	}
	tree := &RecursiveTree{
		Levels: make([][]*TreeNode, depth+1),
	}

	// Level 0: root.
	root := NewInternalNode()
	tree.Root = root
	tree.Levels[0] = []*TreeNode{root}

	// Build each level.
	parents := []*TreeNode{root}
	for d := 1; d <= depth; d++ {
		numNodes := powN(d)
		nodes := make([]*TreeNode, numNodes)
		for i := 0; i < numNodes; i++ {
			if d == depth {
				// Leaf level.
				nodes[i] = NewLeafNode()
			} else {
				nodes[i] = NewInternalNode()
			}
		}
		tree.Levels[d] = nodes

		// Wire children: each parent at level d-1 gets N children from level d.
		if d > 0 {
			childIdx := 0
			for _, parent := range parents {
				for j := 0; j < N && childIdx < len(nodes); j++ {
					parent.AddChild(j, nodes[childIdx].PK, nil)
					childIdx++
				}
			}
		}

		parents = nodes
	}

	return tree
}

// SignRoot signs the root epoch with all transitive child commitments.
// Each internal node signs its children's commitments; the root signs
// its children (level-1 coordinators' commitments).
func (t *RecursiveTree) SignRoot(msg []byte) *NTRUSignature {
	// Sign bottom-up: each level signs for its children.
	for d := len(t.Levels) - 2; d >= 0; d-- {
		for _, node := range t.Levels[d] {
			if node.Children == nil {
				continue
			}
			// Sign this node's children.
			node.SignChildCommitments(msg)
		}
	}
	return t.Root.SignChildCommitments(msg)
}

// VerifyRoot verifies the root signature recursively down the tree.
func (t *RecursiveTree) VerifyRoot(msg []byte, rootSig *NTRUSignature) bool {
	return t.Root.VerifyChildCommitments(msg, rootSig)
}

func powN(exp int) int {
	r := 1
	for i := 0; i < exp; i++ {
		r *= N
	}
	return r
}