const aoc = @import("./aoc.zig");
const std = @import("std");
const Allocator = std.mem.Allocator;
const assert = std.debug.assert;

const NodePart = u6;
const NodeLabel = [3]NodePart;
const NodeJunction = struct {
    from : NodeLabel,
    left : NodeLabel,
    right: NodeLabel
};

const Direction = enum { Left, Right };
const DirectionList = std.ArrayList(Direction);
const NodeJunctionList = std.ArrayList(NodeJunction);
const Input = struct {
    instructions: DirectionList,
    nodes: NodeJunctionList,

    fn deinit(self: *const Input) void {
        self.instructions.deinit();
        self.nodes.deinit();
    }
};

const NodeSplit = struct {
    name: NodePart,
    left: NodeLabel,
    right: NodeLabel,
};
const NodeMap       = std.ArrayList(NodeMapBranch);
const NodeMapBranch = struct {
    name: NodePart,
    nodes: std.ArrayList(NodeMapLeaf)
};
const NodeMapLeaf = struct {
    name: NodePart,
    nodes: std.ArrayList(NodeSplit)
};

fn parse_instructions(allocator: Allocator, line: []const u8) !DirectionList {
    var instructions = DirectionList.init(allocator);
    errdefer instructions.deinit();
    for (line) |c| {
        const direction = switch (c) {
            'R' => Direction.Right,
            'L' => Direction.Left,
            else => unreachable
        };
        try instructions.append(direction);
    }
    return instructions;
}

fn parse_char_to_node_part(c: u8) NodePart {
    if (std.ascii.isAlphabetic(c)) {
        assert(std.ascii.isUpper(c));
        return @intCast(c - 'A' + 10);
    } else if (std.ascii.isDigit(c)) {
        return @intCast(c - '0');
    } else {
        unreachable;
    }
}

fn parse_node_label(text: []const u8) NodeLabel {
    assert(text.len == 3);

    var label: NodeLabel = undefined;
    inline for (0..3) |i| {
        label[i] = parse_char_to_node_part(text[i]);
    }
    return label;
}

fn parse_nodes(allocator: Allocator, lines: []const []const u8) !NodeJunctionList {
    var nodes = try NodeJunctionList.initCapacity(allocator, lines.len);
    errdefer nodes.deinit();

    for (lines) |line| {
        const from_str  = line[0..3];
        const left_str  = line[7..10];
        const right_str = line[12..15];
        nodes.appendAssumeCapacity(NodeJunction{
            .from  = parse_node_label(from_str),
            .left  = parse_node_label(left_str),
            .right = parse_node_label(right_str),
        });
    }

    return nodes;
}

fn parse_input(allocator: Allocator, lines: []const []const u8) !Input {
    const instructions = try parse_instructions(allocator, lines[0]);
    errdefer instructions.deinit();

    const nodes = try parse_nodes(allocator, lines[2..]);
    errdefer nodes.deinit();

    return Input{ .instructions = instructions, .nodes = nodes };
}

fn node_map_find_branch(node_map: *NodeMap, branch_name: NodePart) ?*NodeMapBranch {
    for (node_map.items) |*branch| {
        if (branch.name == branch_name) {
            return branch;
        }
    }
    return null;
}

fn node_leaf_find_split(node_leaf: *NodeMapLeaf, name: NodePart) ?*NodeSplit {
    for (node_leaf.nodes.items) |*split| {
        if (split.name == name) {
            return split;
        }
    }
    return null;
}

fn node_branch_find_leaf(node_branch: *NodeMapBranch, leaf_name: NodePart) ?*NodeMapLeaf {
    for (node_branch.nodes.items) |*leaf| {
        if (leaf.name == leaf_name) {
            return leaf;
        }
    }
    return null;
}

fn create_node_map(allocator: Allocator, nodes: NodeJunctionList) !NodeMap {
    var node_map = NodeMap.init(allocator);
    errdefer free_node_map(node_map);

    for (nodes.items) |node_junction| {
        const from = node_junction.from;
        var node_branch_opt = node_map_find_branch(&node_map, from[0]);
        if (node_branch_opt == null) {
            try node_map.append(NodeMapBranch{
                .name = from[0],
                .nodes = std.ArrayList(NodeMapLeaf).init(allocator)
            });
            node_branch_opt = &node_map.items[node_map.items.len-1];
        }
        var node_branch = node_branch_opt.?;

        var node_leaf_opt = node_branch_find_leaf(node_branch, from[1]);
        if (node_leaf_opt == null) {
            try node_branch.nodes.append(NodeMapLeaf{
                .name = from[1],
                .nodes = std.ArrayList(NodeSplit).init(allocator)
            });
            node_leaf_opt = &node_branch.nodes.items[node_branch.nodes.items.len-1];
        }
        var node_leaf = node_leaf_opt.?;

        try node_leaf.nodes.append(NodeSplit{
            .name = from[2],
            .left = node_junction.left,
            .right = node_junction.right,
        });
    }

    return node_map;
}

fn free_node_map(node_map: NodeMap) void {
    for (node_map.items) |branch| {
        for (branch.nodes.items) |leaf| {
            leaf.nodes.deinit();
        }
        branch.nodes.deinit();
    }
    node_map.deinit();
}

fn node_map_lookup(node_map: *NodeMap, from: NodeLabel, direction: Direction) ?NodeLabel {
    const branch_opt = node_map_find_branch(node_map, from[0]);
    if (branch_opt == null) return null;
    const branch = branch_opt.?;

    const leaf_opt = node_branch_find_leaf(branch, from[1]);
    if (leaf_opt == null) return null;
    const leaf = leaf_opt.?;

    const split_opt = node_leaf_find_split(leaf, from[2]);
    if (split_opt == null) return null;
    const split = split_opt.?;

    return switch (direction) {
        .Left => split.left,
        .Right => split.right,
    };
}

pub fn part1(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    const parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var arena = std.heap.ArenaAllocator.init(allocator);
    defer arena.deinit();

    var node_map = try create_node_map(arena.allocator(), parsed.nodes);
    defer free_node_map(node_map);

    var steps: u32 = 0;
    var goal    = comptime parse_node_label("ZZZ");
    var current = comptime parse_node_label("AAA");
    while (!std.mem.eql(NodePart, &current, &goal)) {
        const instruction = parsed.instructions.items[steps % parsed.instructions.items.len];
        current = node_map_lookup(&node_map, current, instruction) orelse @panic("dead end");
        steps += 1;
    }

    return .{ .uint = steps };
}

fn find_steps_needed(node_map: *NodeMap, instructions: []const Direction, start: NodeLabel) u32 {
    var steps: u32 = 0;
    var current = start;
    while (current[2] != comptime parse_char_to_node_part('Z')) {
        const instruction = instructions[steps % instructions.len];
        current = node_map_lookup(node_map, current, instruction) orelse @panic("dead end");
        steps += 1;
    }
    return steps;
}

fn lcm(a: u64, b: u64) u64 {
    return a * b / std.math.gcd(a, b);
}

pub fn part2(input: *aoc.Input) !aoc.Result {
    const allocator = input.allocator;
    const parsed = try parse_input(allocator, input.lines);
    defer parsed.deinit();

    var arena = std.heap.ArenaAllocator.init(allocator);
    defer arena.deinit();

    var node_map = try create_node_map(arena.allocator(), parsed.nodes);
    defer free_node_map(node_map);

    var all_steps = std.ArrayList(u32).init(allocator);
    defer all_steps.deinit();

    for (parsed.nodes.items) |node_junction| {
        if (node_junction.from[2] == comptime parse_char_to_node_part('A')) {
            const steps = find_steps_needed(&node_map, parsed.instructions.items, node_junction.from);
            try all_steps.append(steps);
        }
    }

    assert(all_steps.items.len > 0);
    var answer: u64 = all_steps.items[0];
    for (1..all_steps.items.len) |i| {
        answer = lcm(answer, @intCast(all_steps.items[i]));
    }

    return .{ .uint = answer };
}


test "part 1 example 1" {
    const example_input = [_][]const u8{
        "RL",
        "",
        "AAA = (BBB, CCC)",
        "BBB = (DDD, EEE)",
        "CCC = (ZZZ, GGG)",
        "DDD = (DDD, DDD)",
        "EEE = (EEE, EEE)",
        "GGG = (GGG, GGG)",
        "ZZZ = (ZZZ, ZZZ)",
    };
    try aoc.expectAnswerUInt(part1, 2, &example_input);
}

test "part 1 example 2" {
    const example_input = [_][]const u8{
        "LLR",
        "",
        "AAA = (BBB, BBB)",
        "BBB = (AAA, ZZZ)",
        "ZZZ = (ZZZ, ZZZ)",
    };
    try aoc.expectAnswerUInt(part1, 6, &example_input);
}

test "part 2 example" {
    const example_input = [_][]const u8{
        "LR",
        "",
        "11A = (11B, XXX)",
        "11B = (XXX, 11Z)",
        "11Z = (11B, XXX)",
        "22A = (22B, XXX)",
        "22B = (22C, 22C)",
        "22C = (22Z, 22Z)",
        "22Z = (22B, 22B)",
        "XXX = (XXX, XXX)",
    };
    try aoc.expectAnswerUInt(part2, 6, &example_input);
}