const std = @import("std");
const Allocator = std.mem.Allocator;

pub fn Tree(comptime T: type) type {
    return struct {
        // const Self = @This();

        root: Node,

        const Children = std.ArrayList(Node);
        const Node = struct {
            val: T,
            parent: ?*Node,
            index: usize,
            children: std.ArrayList(Node),

            pub fn init(val: T, parent: ?*Node, index: usize) Node {
                return .{
                    .val = val,
                    .children = .empty,
                    .parent = parent,
                    .index = index,
                };
            }

            pub fn append(self: *Node, allocator: Allocator, val: T) !void {
                const child = Node.init(val, self, self.children.items.len);
                try self.children.append(allocator, child);
            }
        };

        pub fn init(val: T) @This() {
            return .{ .root = .init(val, null, 0) };
        }

        pub const DepthFirstIterator = struct {
            const State = enum {
                GoDeeper,
                GoBroader,
            };
            tree: *Tree(T),
            current: ?*Node,
            state: State,

            pub fn init(tree: *Tree(T)) DepthFirstIterator {
                return DepthFirstIterator{
                    .tree = tree,
                    .current = &tree.root,
                    .state = State.GoDeeper,
                };
            }

            pub fn next(it: *DepthFirstIterator) ?*Node {
                // State machine
                while (it.current) |current| {
                    switch (it.state) {
                        State.GoDeeper => {
                            // Follow child node until deepest possible level
                            if (current.children.items.len > 0) {
                                it.current = &current.children.items[0];
                            } else {
                                it.state = State.GoBroader;
                                return current;
                            }
                        },
                        State.GoBroader => {
                            if (current.parent) |parent| {
                                if (parent.children.items.len - 1 > (current.index)) {
                                    it.current = &parent.children.items[current.index + 1];
                                    it.state = .GoDeeper;
                                } else {
                                    it.current = current.parent;
                                    return current.parent;
                                }
                            } else {
                                return null;
                            }
                        },
                    }
                }
                return null;
            }

            pub fn reset(it: *DepthFirstIterator) void {
                it.current = it.tree.root;
            }
        };

        pub fn depthFirstIterator(tree: *@This()) DepthFirstIterator {
            return DepthFirstIterator.init(tree);
        }

        pub fn deinit(self: *@This(), allocator: Allocator) void {
            var iterator = self.depthFirstIterator();
            while (iterator.next()) |node| {
                node.children.deinit(allocator);
            }
        }
    };
}

pub fn main() !void {
    var debug_allocator = std.heap.DebugAllocator(.{}).init;
    defer _ = debug_allocator.deinit();
    const allocator = debug_allocator.allocator();

    const UTree = Tree(u16);

    var tree = UTree.init(10);
    defer tree.deinit(allocator);

    try tree.root.append(allocator, 20);
    try tree.root.append(allocator, 500);

    var iterator = tree.depthFirstIterator();
    std.debug.print("\n", .{});
    while (iterator.next()) |node| {
        var accumulator: u32 = node.val;
        var current = node;
        while (current.parent) |parent| {
            accumulator *= parent.val;
            current = parent;
        }
        std.debug.print("Val: {} Acc: {}\n", .{ node.val, accumulator });
    }
}

test "Tree" {
    const allocator = std.testing.allocator;

    const UTree = Tree(u16);

    var tree = UTree.init(10);
    defer tree.deinit(allocator);

    try tree.root.append(allocator, 20);
    try tree.root.append(allocator, 500);

    var iterator = tree.depthFirstIterator();
    std.debug.print("\n", .{});
    while (iterator.next()) |node| {
        var accumulator: u32 = node.val;
        var current = node;
        while (current.parent) |parent| {
            accumulator *= parent.val;
            current = parent;
        }
        std.debug.print("Val: {} Acc: {}\n", .{ node.val, accumulator });
    }
}
