package main

import (
	"slices"

	"github.com/pkg/errors"
)

func DecomposeTree(tree map[string]*WeightedNode, threshold uint16) ([]SubTree, error) {
	var res []SubTree
	// Рабочий цикл
	for len(tree) > 0 {
		// Индексируем tree по weight и сортируем индекс
		index := GetTreeIndex(tree)
		// Ищем ветку подходящую под лимит
		branchID, weight := GetBranch(tree, index, threshold)

		// Определяем путь к найденной ветке
		path, err := GetNodePath(tree, branchID)
		if err != nil {
			return nil, errors.Wrap(err, "[GetNodePath] ")
		}
		// Вычитаем ветку из веса предков
		RecalculateParents(tree, branchID, weight)

		// Вырезаем её из среза вместе со всеми потомками
		branch := CutBranch(tree, branchID)
		// Добавляем найденную ветку к res
		res = append(res, SubTree{
			path: path,
			tree: branch,
		})
	}
	return res, nil
}

type TreeIndex struct {
	id     string
	weight uint16
}

// GetTreeIndex индексирует дерево по weight и сортирует индекс по уменьшению weight
func GetTreeIndex(tree map[string]*WeightedNode) []TreeIndex {
	var index = make([]TreeIndex, len(tree))
	i := 0
	for k, v := range tree {
		index[i] = TreeIndex{
			id:     k,
			weight: v.weight,
		}
		i += 1
	}
	slices.SortFunc(index, func(a, b TreeIndex) int {
		return int(b.weight) - int(a.weight)
	})
	return index
}

// GetBranch возвращает первую найденную ветвь дерева, суммарный вес которой меньше threshold
// Для работы необходим актуальный индекс мапы, отсортированный по убыванию
func GetBranch(tree map[string]*WeightedNode, index []TreeIndex, threshold uint16) (string, uint16) {
	var branchID string
	for _, v := range index {
		if v.weight < threshold {
			branchID = v.id
			break
		}
	}
	return tree[branchID].node.ID, tree[branchID].weight
}

// GetNodePath возвращает полный путь ноды от корня
func GetNodePath(tree map[string]*WeightedNode, nodeID string) (string, error) {
	if _, ok := tree[nodeID]; !ok {
		return "", errors.New("node not exists")
	}
	if tree[nodeID].parent == nil {
		return "/", nil
	}
	nodeID = tree[nodeID].parent.ID
	path := tree[nodeID].node.Name
	for {
		if tree[nodeID].parent.Name == "/" {
			path = "/" + path
			break
		}
		nodeID = tree[nodeID].parent.ID
		path = tree[nodeID].node.Name + "/" + path
	}
	return path, nil
}

// RecalculateParents вычитает вес вырезанной ветки из всех родительских веток
func RecalculateParents(tree map[string]*WeightedNode, currentID string, nodeWeight uint16) {
	if tree[currentID].parent == nil {
		return
	}
	currentID = tree[currentID].parent.ID
	for {
		tree[currentID].weight -= nodeWeight
		if tree[currentID].parent == nil {
			break
		}
		currentID = tree[currentID].parent.ID
	}
}

// CutBranch удаляет из дерева ветку со всеми потомками, а так же удаляет ссылку на себя из родительской ноды
func CutBranch(tree map[string]*WeightedNode, nodeID string) *Node {
	// Удаление ссылки на ноду из родительской ноды
	if tree[nodeID].parent != nil {
		parentID := tree[nodeID].parent.ID
		k := slices.IndexFunc(tree[parentID].children, func(node *Node) bool {
			return node.ID == nodeID
		})
		if k > 0 {
			tree[parentID].children = append(tree[parentID].children[:k], tree[parentID].children[k+1:]...)
		}
		m := slices.IndexFunc(tree[parentID].node.Children, func(node *Node) bool {
			return node.ID == nodeID
		})
		if m > 0 {
			tree[parentID].children = append(tree[parentID].node.Children[:m], tree[parentID].node.Children[m+1:]...)
		}
	}

	node := tree[nodeID].node
	// Удаление ветви вместе со всеми потомками
	stack := []string{nodeID}
	for len(stack) > 0 {
		current := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		if _, ok := tree[current]; ok {
			for _, v := range tree[current].children {
				stack = append(stack, v.ID)
			}
			delete(tree, current)
		}
	}
	return node
}