const std = @import("std");

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;

    // Read our input
    var f = try std.fs.cwd().openFile("input", .{});
    defer f.close();
    var contents = try f.readToEndAlloc(alloc, std.math.maxInt(u32));
    defer alloc.free(contents);

    var lit = std.mem.tokenize(contents, "\n");
    var corrupt_score: u32 = 0;
    var completion_scores = std.ArrayList(u64).init(alloc);
    defer completion_scores.deinit();
    var line_id: usize = 0;
    while (lit.next()) |line| {
        //std.log.debug("{any}", .{nodes});
        var tree = try line_to_tree(alloc, line, line_id);
        corrupt_score += tree.corrupt;
        if (tree.corrupt == 0) {
            // Try completion
            var score = try tree.autocomplete();
            std.log.debug("Line {} autocompletion score {}", .{line_id, score});
            try completion_scores.append(score); // score
        }
        line_id += 1;
    }
    std.log.info("[Part 1] Sum of corrupt line scores: {}", .{corrupt_score});
    std.sort.sort(u64, completion_scores.items, {}, comptime std.sort.asc(u64));

    std.log.info(
        "[Part 2] Completion scores: {}",
        .{completion_scores.items[completion_scores.items.len/2]}
    );
}

fn line_to_tree(alloc: *std.mem.Allocator, line: []const u8, id: usize) !Tree {
    var tkzr = Tokenizer.init(line);
    var tkns = std.ArrayList(Token).init(alloc);
    var corrupt_score: u32 = 0;
    while (true) {
        var tkn = tkzr.next();
        if (tkn.tag != .eos) {
            try tkns.append(tkn);
        }
        else {
            break;
        }
    }
    //std.log.debug("{} tokens", .{tkns.items.len});
    var nodes = std.ArrayList(Node).init(alloc);
    try nodes.append(.{
        .index = 0,
        .parent = 0,
        .tag = .root,
        .start_index = 0,
    });
    var current_node: usize = 0;
    for (tkns.items) |tkn, index| {
        switch(tkn.tag) {
            .r_paren,
            .r_brace,
            .r_bracket,
            .r_angle_bracket,
            => {
                if (current_node == 0) {
                    std.log.debug("Closing token before opening token: {}",
                                  .{tkn});
                    break;
                }
                else {
                    //std.log.warn("Current node: {}", .{current_node});
                    nodes.items[current_node].end_index = index;
                    if (nodes.items[current_node].tag != token_tag_to_node_tag(tkns.items[index].tag)) {
                        std.log.warn(
                            "Line {}: mismatched tags: {} != {}",
                            .{id, tkns.items[nodes.items[current_node].start_index].tag,
                              tkns.items[index].tag}
                        );
                        corrupt_score += switch(tkns.items[index].tag) {
                            .r_paren => @as(u32, 3),
                            .r_bracket => @as(u32, 57),
                            .r_brace => @as(u32, 1197),
                            .r_angle_bracket => @as(u32, 25137),
                            else => @as(u32, 0),
                        };
                        break;
                    }
                    current_node = nodes.items[current_node].parent;
                }
            },
            .eos,
            .invalid,
            => unreachable,
            else => {
                // Opens a new node
                var node = Node {
                    .index = nodes.items.len + 1,
                    .parent = current_node,
                    .start_index = index,
                    .tag = switch(tkn.tag) {
                        .l_paren => .paren,
                        .l_brace => .brace,
                        .l_bracket => .bracket,
                        .l_angle_bracket => .angle_bracket,
                        else => unreachable,
                    }
                };
                try nodes.append(node);
                try nodes.items[current_node].children.append(alloc, nodes.items.len - 1);
                current_node = nodes.items.len - 1;
                //std.log.warn("Start new node from token index {} at {}: {any}",
                //            .{index, current_node, node});
            },
        }
    }
    return Tree {
        .nodes = nodes,
        .tokens = tkns,
        .corrupt = corrupt_score,
    };
}

pub const Tree = struct {
    nodes: std.ArrayList(Node),
    tokens: std.ArrayList(Token),
    corrupt: u32 = 0,

    pub fn autocomplete(self: *Tree) anyerror!u64 {
        // I guess we want to traverse the tree, adding in tokens for end_index
        // where necessary
        var node = self.nodes.items[0];
        var score: u64 = 0;
        for (node.children.items) |n| {
            score += try self.complete(&self.nodes.items[n], score);
        }
        return score;
    }

    fn complete(self: *Tree, node: *Node, score: u64) anyerror!u64 {
        var s = score;
        for (node.*.children.items) |n| {
            s += try self.complete(&self.nodes.items[n], score);
        }
        if (node.*.end_index) |ei| {
            return s;
        }
        else {
            try self.tokens.append(.{
                .tag = switch(node.*.tag) {
                    .paren => .r_paren,
                    .brace => .r_brace,
                    .bracket => .r_bracket,
                    .angle_bracket => .r_angle_bracket,
                    else => .invalid,
                },
                .loc = .{.start = 0, .end = 0}, // fake
            });
            node.*.end_index = self.tokens.items.len - 1;
            s *= 5;
            s += switch(node.*.tag) {
                .paren => @as(u64, 1),
                .brace => @as(u64, 3),
                .bracket => @as(u64, 2),
                .angle_bracket => @as(u64, 4),
                else => @as(u64, 0),
            };
        }
        return s;
    }
};

pub const Node = struct {
    index: usize,
    parent: usize,
    children: std.ArrayListUnmanaged(usize) = .{},
    tag: Tag,
    start_index: usize,
    end_index: ?usize = null,
    pub const Tag = enum {
        root,
        paren,
        brace,
        bracket,
        angle_bracket,
        invalid,
    };
};

fn token_tag_to_node_tag(t: Token.Tag) Node.Tag {
    switch(t) {
        .l_paren,
        .r_paren,
        => return .paren,
        .l_brace,
        .r_brace,
        => return .brace,
        .l_bracket,
        .r_bracket,
        => return .bracket,
        .l_angle_bracket,
        .r_angle_bracket,
        => return .angle_bracket,
        else => return .invalid
    }
}

pub const Token = struct {
    tag: Tag,
    loc: Loc,

    pub const Loc = struct {
        start: usize,
        end: usize,
    };
    pub const Tag = enum {
        l_paren,
        r_paren,
        l_brace,
        r_brace,
        l_bracket,
        r_bracket,
        l_angle_bracket,
        r_angle_bracket,
        eos,
        invalid,
    };

    pub fn lexeme(tag: Tag) ?[]const u8 {
        return switch (tag) {
            .l_paren => "(",
            .r_paren => ")",
            .l_brace => "{",
            .r_brace => "}",
            .l_bracket => "[",
            .r_bracket => "]",
            .l_angle_bracket => "<",
            .r_angle_bracket => ">",
            else => "",
        };
    }

};

pub const Tokenizer = struct {
    index: usize,
    buffer: []const u8,

    pub fn init(buffer: []const u8) Tokenizer {
        return Tokenizer{
            .buffer = buffer,
            .index = 0,
        };
    }

    pub fn next(self: *Tokenizer) Token {
        const start_index = self.index;
        var result = Token {
            .tag = .eos,
            .loc = .{
                .start = start_index,
                .end   = undefined,
            },
        };
        while (self.index < self.buffer.len) : (self.index += 1) {
            const c = self.buffer[self.index];
            switch (c) {
                '(' => {
                    result.tag = .l_paren;
                    self.index += 1;
                    break;
                },
                ')' => {
                    result.tag = .r_paren;
                    self.index += 1;
                    break;
                },
                '{' => {
                    result.tag = .l_brace;
                    self.index += 1;
                    break;
                },
                '}' => {
                    result.tag = .r_brace;
                    self.index += 1;
                    break;
                },
                '[' => {
                    result.tag = .l_bracket;
                    self.index += 1;
                    break;
                },
                ']' => {
                    result.tag = .r_bracket;
                    self.index += 1;
                    break;
                },
                '<' => {
                    result.tag = .l_angle_bracket;
                    self.index += 1;
                    break;
                },
                '>' => {
                    result.tag = .r_angle_bracket;
                    self.index += 1;
                    break;
                },
                else => {
                    result.tag = .invalid;
                    self.index += 1;
                    break;
                },
            }
        }
        if (result.tag == .eos) {
            result.loc.start = self.index;
        }
        result.loc.end = self.index;
        return result;
    }
};

test "tokenizer - all" {
    try testTokenize("()[]{}<>", &.{
        .l_paren, .r_paren, .l_bracket, .r_bracket, .l_brace, .r_brace,
        .l_angle_bracket, .r_angle_bracket
    });
}

test "tokenizer - early end" {
    try testTokenize("({", &.{
        .l_paren, .l_brace,
    });
}

test "tokenizer - invalid" {
    try testTokenize("<d", &.{
        .l_angle_bracket, .invalid,
    });
}

fn testTokenize(source: []const u8, expected_tokens: []const Token.Tag) !void {
    var tokenizer = Tokenizer.init(source);
    for (expected_tokens) |expected_token_id| {
        const token = tokenizer.next();
        if (token.tag != expected_token_id) {
            std.debug.panic("expected {s}, found {s}\n", .{ @tagName(expected_token_id), @tagName(token.tag) });
        }
    }
    const last_token = tokenizer.next();
    try std.testing.expect(last_token.tag == .eos);
    try std.testing.expect(last_token.loc.start == source.len);
}

test "ok" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "[<>({}){}[([])<>]]", 1);
    try std.testing.expectEqual(@as(u32, 0), tree.corrupt);
}

test "a" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "{([(<{}[<>[]}>{[]{[(<()>", 1);
    try std.testing.expectEqual(@as(u32, 1197), tree.corrupt);
}

test "b" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "[[<[([]))<([[{}[[()]]]", 1);
    try std.testing.expectEqual(@as(u32, 3), tree.corrupt);
}

test "c" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "[{[{({}]{}}([{[{{{}}([]", 1);
    try std.testing.expectEqual(@as(u32, 57), tree.corrupt);
}

test "d" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "[<(<(<(<{}))><([]([]()", 1);
    try std.testing.expectEqual(@as(u32, 3), tree.corrupt);
}

test "e" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "<{([([[(<>()){}]>(<<{{", 1);
    try std.testing.expectEqual(@as(u32, 25137), tree.corrupt);
}

test "autocomplete simple" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "<", 1);
    try std.testing.expectEqual(@as(u32, 0), tree.corrupt);
    var score = try tree.autocomplete();
    try std.testing.expectEqual(@as(u64, 4), score);
}

test "autocomplete" {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    const alloc = &arena.allocator;
    var tree = try line_to_tree(alloc, "<{([{{}}[<[[[<>{}]]]>[]]", 1);
    try std.testing.expectEqual(@as(u32, 0), tree.corrupt);
    var score = try tree.autocomplete();
    try std.testing.expectEqual(@as(u64, 294), score);
}