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

const Cube = enum { Red, Green, Blue };
const CubeSet = std.EnumArray(Cube, u32);
const Game = struct {
    sets: [10]CubeSet,
    set_count: u32,
};

fn parse_cube(cube_name: []const u8) ?Cube {
    if (std.mem.eql(u8, cube_name, "blue")) {
        return .Blue;
    } else if (std.mem.eql(u8, cube_name, "red")) {
        return .Red;
    } if (std.mem.eql(u8, cube_name, "green")) {
        return .Green;
    }
    return null;
}

fn parse_game(game: *Game, line: []const u8) void {
    const colon = std.mem.indexOf(u8, line, ":") orelse @panic("Invalid line, missing ':'");

    game.set_count = 0;

    var game_sets = std.mem.splitSequence(u8, line[(colon+2)..], "; ");
    while (game_sets.next()) |game_set_str| {
        var cube_set = &game.sets[game.set_count];
        cube_set.* = CubeSet.initFill(0);
        game.set_count += 1;

        var cubes = std.mem.splitSequence(u8, game_set_str, ", ");
        while (cubes.next()) |cube_count| {
            const space = std.mem.indexOf(u8, cube_count, " ") orelse @panic("Invalid format, expected space");

            const cube = parse_cube(cube_count[(space+1)..]) orelse @panic("Invalid cube name");
            const count = std.fmt.parseInt(u32, cube_count[0..space], 10) catch @panic("Invalid format");
            cube_set.set(cube, count);
        }
    }
}

fn parse_games(allocator: Allocator, lines: []const []const u8) ![]Game {
    var games = try allocator.alloc(Game, lines.len);
    for (0.., lines) |i, line| {
        parse_game(&games[i], line);
    }
    return games;
}

pub fn part1(input: *aoc.Input) !aoc.Result {
    const lines = input.lines;
    const allocator = input.allocator;

    const games = try parse_games(allocator, lines);
    defer allocator.free(games);

    const max_red = 12;
    const max_green = 13;
    const max_blue = 14;
    var sum: u32 = 0;
    for (1.., games) |idx, game| {
        var possible = true;
        for (game.sets[0..game.set_count]) |cube_set| {
            if ((cube_set.get(.Red) > max_red) or (cube_set.get(.Blue) > max_blue) or (cube_set.get(.Green) > max_green)) {
                possible = false;
                break;
            }
        }

        if (possible) {
            sum += @intCast(idx);
        }
    }

    return .{ .uint = sum };
}

pub fn part2(input: *aoc.Input) !aoc.Result {
    const lines = input.lines;
    const allocator = input.allocator;
    const games = try parse_games(allocator, lines);
    defer allocator.free(games);

    var sum: u32 = 0;
    for (games) |game| {
        var max_red: u32   = 0;
        var max_green: u32 = 0;
        var max_blue: u32  = 0;

        for (game.sets[0..game.set_count]) |cube_set| {
            max_red   = @max(max_red, cube_set.get(.Red));
            max_green = @max(max_green, cube_set.get(.Green));
            max_blue  = @max(max_blue, cube_set.get(.Blue));
        }

        sum += max_red*max_green*max_blue;
    }

    return .{ .uint = sum };
}

const example_input = [_][]const u8{
    "Game 1: 3 blue, 4 red; 1 red, 2 green, 6 blue; 2 green",
    "Game 2: 1 blue, 2 green; 3 green, 4 blue, 1 red; 1 green, 1 blue",
    "Game 3: 8 green, 6 blue, 20 red; 5 blue, 4 red, 13 green; 5 green, 1 red",
    "Game 4: 1 green, 3 red, 6 blue; 3 green, 6 red; 3 green, 15 blue, 14 red",
    "Game 5: 6 red, 1 blue, 3 green; 2 blue, 1 red, 2 green"
};

test "part 1 example" {
    try aoc.expectAnswerUInt(part1, 8, &example_input);
}

test "part 2 example" {
    try aoc.expectAnswerUInt(part2, 2286, &example_input);
}