class_name BehaviorTreeVisualizer
extends Window

const D_TREE_NODE: PackedScene = preload("res://scripts/visualization/d_tree_node.tscn")
@onready var graph_edit: GraphEdit = %GraphEdit

# Input data
var behavior_tree: BehaviorTree
var all_nodes: Array[DTreeNode]         = []
var all_node_names_to_nodes: Dictionary = {}
# Configuration
var x_spacing: int = 300  # Base horizontal distance between parent and child nodes
var y_spacing: int = 200  # Vertical distance between sibling nodes


func build_tree() -> void:
    if not behavior_tree:
        push_error("No behavior tree set.")
        return

    # Reset current visualization
    graph_edit.clear_connections()
    for node in all_nodes:
        node.queue_free()
    all_nodes.clear()
    all_node_names_to_nodes.clear()

    # Build tree starting from the root task
    var root_node: Task = behavior_tree.behavior_tree
    if not root_node:
        push_error("Root behavior tree node is null.")
        return

    # Calculate positions for all nodes
    var positions: Dictionary = calculate_tree_layout(root_node)

    # Create the nodes and connections in the GraphEdit
    for task_name in positions.keys():
        var position   = positions[task_name]
        var task: Task = find_task_by_name(root_node, task_name)
        if task:
            create_node(task, position)

    # Connect the nodes
    for node in all_nodes:
        var task_name: String = node.title
        var task: Task        = find_task_by_name(root_node, task_name)
        if task:
            for child in task.get_children():
                if child is Task:
                    connect_nodes(task.name, child.name)


func calculate_tree_layout(root_task: Task) -> Dictionary:
    """
    Calculates the positions of all nodes in the tree, ensuring no overlaps.
    """
    var positions: Dictionary       = {}
    var layers                      = get_tree_layers(root_task)
    var layer_count                 = layers.size()
    var layer_heights: Array[float] = []

    # Calculate height for each layer based on the number of nodes
    for layer in layers:
        layer_heights.append(float(layer.size() - 1) * y_spacing)

    # Position nodes starting from the deepest layer and working up
    for layer_index in range(layer_count - 1, -1, -1):
        var layer                 = layers[layer_index]
        var layer_y_offset: float = 0.0
        if layer_index < layer_count - 1:
            layer_y_offset = layer_heights[layer_index + 1] / 2.0

        var x: float  = float(layer_index) * x_spacing
        var num_nodes = layer.size()
        var start_y   = -layer_heights[layer_index] / 2.0

        for node_index in range(num_nodes):
            var task = layer[node_index]
            var y    = start_y + float(node_index) * y_spacing + layer_y_offset
            positions[task.name] = Vector2(x, y)

    return positions


func get_tree_layers(root_task: Task):
    """
    Returns a list of layers, where each layer is a list of tasks at the same depth.
    """
    var layers                   = []
    var queue: Array[Dictionary] = [{"task": root_task, "depth": 0}]

    while not queue.is_empty():
        var current_item = queue.pop_front()
        var task         = current_item.task
        var depth        = current_item.depth

        if layers.size() <= depth:
            layers.append([])
        layers[depth].append(task)

        for child in task.get_children():
            if child is Task:
                queue.append({"task": child as Task, "depth": depth + 1})

    return layers


func create_node(task: Task, position: Vector2) -> void:
    """
    Creates a visual node for the given task at the specified position.
    """
    var graph_node: DTreeNode = D_TREE_NODE.instantiate()
    graph_node.title = task.name
    graph_node.position_offset = position
    graph_edit.add_child(graph_node)

    all_nodes.append(graph_node)
    all_node_names_to_nodes[task.name] = graph_node


func connect_nodes(parent_name: String, child_name: String) -> void:
    """
    Connects two nodes in the GraphEdit.
    """
    if all_node_names_to_nodes.has(parent_name) and all_node_names_to_nodes.has(child_name):
        graph_edit.connect_node(
            all_node_names_to_nodes[parent_name].get_name(), 0, # Parent's output port
            all_node_names_to_nodes[child_name].get_name(), 0   # Child's input port
        )


func find_task_by_name(root: Task, name: String) -> Task:
    """
    Finds a task by name within the behavior tree.
    """
    if root.name == name:
        return root
    for child in root.get_children():
        if child is Task:
            var result = find_task_by_name(child as Task, name)
            if result:
                return result
    return null