// Copyright (C) 2019-2025, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package tree

import (
	"context"

	"golang.org/x/exp/maps"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow/consensus/snowman"
)

// Tree handles the propagation of block acceptance and rejection to inner
// blocks.
//
// The Tree is needed because:
// 1. The consensus engine guarantees that for each verified block, either
// Accept() or Reject() are eventually called, and they are called only once.
// The proposervm must maintain these invariants for the wrapped VM.
// 2. A given inner block may be wrapped into multiple different proposervm
// blocks (e.g. same inner block generated by two validators).
//
// The Tree prevents Accept() and Reject() from being called multiple times on
// the same inner block by:
// 1. tracking inner blocks in a tree-like structure, to be able to easily spot
// siblings
// 2. rejecting an inner block only when one of the siblings is accepted.
// Rejection of a proposervm block does not imply its inner block rejection
// (it may be held by a different proposervm block).
type Tree interface {
	// Add places the block in the tree
	Add(snowman.Block)

	// Get returns the block that was added to this tree whose parent and ID
	// match the provided block. If non-exists, then false will be returned.
	Get(snowman.Block) (snowman.Block, bool)

	// Accept marks the provided block as accepted and rejects every conflicting
	// block.
	Accept(context.Context, snowman.Block) error
}

type tree struct {
	// parentID -> childID -> childBlock
	nodes map[ids.ID]map[ids.ID]snowman.Block
}

func New() Tree {
	return &tree{
		nodes: make(map[ids.ID]map[ids.ID]snowman.Block),
	}
}

func (t *tree) Add(blk snowman.Block) {
	parentID := blk.Parent()
	children, exists := t.nodes[parentID]
	if !exists {
		children = make(map[ids.ID]snowman.Block)
		t.nodes[parentID] = children
	}
	blkID := blk.ID()
	children[blkID] = blk
}

func (t *tree) Get(blk snowman.Block) (snowman.Block, bool) {
	parentID := blk.Parent()
	children := t.nodes[parentID]
	blkID := blk.ID()
	originalBlk, exists := children[blkID]
	return originalBlk, exists
}

func (t *tree) Accept(ctx context.Context, blk snowman.Block) error {
	// accept the provided block
	if err := blk.Accept(ctx); err != nil {
		return err
	}

	// get the siblings of the block
	parentID := blk.Parent()
	children := t.nodes[parentID]
	blkID := blk.ID()
	delete(children, blkID)
	delete(t.nodes, parentID)

	// mark the siblings of the accepted block as rejectable
	childrenToReject := maps.Values(children)

	// reject all the rejectable blocks
	for len(childrenToReject) > 0 {
		i := len(childrenToReject) - 1
		child := childrenToReject[i]
		childrenToReject = childrenToReject[:i]

		// reject the block
		if err := child.Reject(ctx); err != nil {
			return err
		}

		// mark the progeny of this block as being rejectable
		blkID := child.ID()
		children := t.nodes[blkID]
		childrenToReject = append(childrenToReject, maps.Values(children)...)
		delete(t.nodes, blkID)
	}
	return nil
}
