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

const HandStrength = enum {
    Nothing,
    HighCard,
    OnePair,
    TwoPair,
    ThreeOfAKind,
    FullHouse,
    FourOfAKind,
    FiveOfAKind,
};

const CARD_VARIANTS = 13;
const CardIndex = u4;

const Hand = [5]CardIndex;

const HandBid = struct {
    hand: Hand = undefined,
    bid: u32
};

const CardCounts = [CARD_VARIANTS]CardIndex;
const IndexArray = std.ArrayList(usize);
const HandBids = std.ArrayList(HandBid);

fn hand_to_str(hand: Hand) [5]u8 {
    var str = [1]u8{0} ** 5;
    for (0..hand.len) |i| {
        str[i] = switch (hand[i]) {
            12 => 'A',
            11 => 'K',
            10 => 'Q',
            9  => 'J',
            8  => 'T',
            7  => '9',
            6  => '8',
            5  => '7',
            4  => '6',
            3  => '5',
            2  => '4',
            1  => '3',
            0  => '2',
            else => unreachable
        };
    }
    return str;
}

fn get_card_index(char: u8) CardIndex {
    return switch (char) {
        'A' => 12,
        'K' => 11,
        'Q' => 10,
        'J' => 9,
        'T' => 8,
        '9' => 7,
        '8' => 6,
        '7' => 5,
        '6' => 4,
        '5' => 3,
        '4' => 2,
        '3' => 1,
        '2' => 0,
        else => unreachable
    };
}

fn parse_input(allocator: Allocator, lines: []const []const u8) !HandBids {
    var hand_bids = HandBids.init(allocator);

    for (lines) |line| {
        var tokens = std.mem.tokenizeScalar(u8, line, ' ');
        const hand_str = tokens.next().?;
        const bid_str = tokens.next().?;
        const bid = try std.fmt.parseInt(u32, bid_str, 10);
        assert(hand_str.len == 5);

        var hand_bid = HandBid{ .bid = bid };
        for (0..5) |i| {
            hand_bid.hand[i] = get_card_index(hand_str[i]);
        }

        try hand_bids.append(hand_bid);
    }

    return hand_bids;
}

fn count_cards(hand: Hand) CardCounts {
    var counts = [1]CardIndex{0} ** CARD_VARIANTS;
    for (hand) |card| {
        counts[card] += 1;
    }
    return counts;
}

fn determine_card_strength(hand: Hand) HandStrength {
    var counts = count_cards(hand);

    var ones_count: u4 = 0;
    var two_count: u4 = 0;
    var three_count: u4 = 0;
    for (counts) |count| {
        switch (count) {
            5 => return .FiveOfAKind,
            4 => return .FourOfAKind,
            3 => three_count += 1,
            2 => two_count += 1,
            1 => ones_count += 1,
            else => continue
        }
    }

    if (ones_count == 5) {
        return .HighCard;
    }
    if (two_count == 2) {
        return .TwoPair;
    }
    if (two_count == 1 and three_count == 1) {
        return .FullHouse;
    }
    if (ones_count == 2 and three_count == 1) {
        return .ThreeOfAKind;
    }
    if (two_count == 1 and ones_count == 3) {
        return .OnePair;
    }

    return .Nothing;
}

fn determine_card_strength_joker(hand: Hand) HandStrength {
    const joker_card = comptime get_card_index('J');

    var jokers = std.BoundedArray(usize, 5).init(0) catch unreachable;
    var card_types = std.BoundedArray(CardIndex, 5).init(0) catch unreachable;
    for (0.., hand) |i, card| {
        if (card == joker_card) {
            jokers.append(i) catch unreachable;
        }
        if (std.mem.indexOfScalar(CardIndex, card_types.slice(), card) == null) {
            card_types.append(card) catch unreachable;
        }
    }

    if (jokers.len == 0) {
        return determine_card_strength(hand);
    }

    var joker_choices_container = [1]usize{0} ** 5;
    var joker_choices = joker_choices_container[0..jokers.len];

    var modified_hand: Hand = undefined;
    @memcpy(&modified_hand, &hand);

    var best_strength = HandStrength.Nothing;
    const combination_count = std.math.powi(usize, card_types.len, jokers.len) catch unreachable;
    for (0..combination_count) |_| {
        for (0..joker_choices.len) |i| {
            modified_hand[jokers.get(i)] = card_types.get(joker_choices[i]);
        }

        const strength = determine_card_strength(modified_hand);

        best_strength = @enumFromInt(@max(@intFromEnum(best_strength), @intFromEnum(strength)));

        joker_choices[0] += 1;
        for (0..(joker_choices.len-1)) |i| {
            if (joker_choices[i] == card_types.len) {
                joker_choices[i] = 0;
                joker_choices[i+1] += 1;
            }
        }
    }

    return best_strength;
}

fn compare_hands(a: *Hand, b: *Hand) bool {
    inline for (0..a.len) |i| {
        if (a[i] != b[i]) {
            return a[i] > b[i];
        }
    }
    return false;
}

fn compare_hands_joker(a: *Hand, b: *Hand) bool {
    const joker_card = comptime get_card_index('J');

    inline for (0..a.len) |i| {
        if (a[i] != b[i]) {
            const card_a = if (a[i] == joker_card) 0 else a[i]+1;
            const card_b = if (b[i] == joker_card) 0 else b[i]+1;
            return card_a > card_b;
        }
    }
    return false;
}

fn sort_same_strength(hands: []Hand, indexes: []usize, cmp: *const fn(*Hand, *Hand) bool) void {
    for (0..(indexes.len-1)) |i| {
        for ((i+1)..indexes.len) |j| {
            const hand_a = &hands[indexes[i]];
            const hand_b = &hands[indexes[j]];
            if (cmp(hand_a, hand_b)) {
                std.mem.swap(usize, &indexes[i], &indexes[j]);
            }
        }
    }
}

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

    var hand_strengths = try allocator.alloc(HandStrength, hand_bids.items.len);
    defer allocator.free(hand_strengths);

    var hands = try allocator.alloc(Hand, hand_bids.items.len);
    defer allocator.free(hands);

    for (0..hand_bids.items.len) |i| {
        hands[i] = hand_bids.items[i].hand;
        hand_strengths[i] = determine_card_strength(hand_bids.items[i].hand);
    }

    var order = try IndexArray.initCapacity(allocator, hand_bids.items.len);
    defer order.deinit();

    var sub_order = try IndexArray.initCapacity(allocator, hand_bids.items.len);
    defer sub_order.deinit();

    inline for (@typeInfo(HandStrength).Enum.fields) |field| {
        sub_order.clearRetainingCapacity();

        const hand_strength: HandStrength = @enumFromInt(field.value);
        for (0..hand_bids.items.len) |idx| {
            if (hand_strengths[idx] == hand_strength) {
                sub_order.appendAssumeCapacity(idx);
            }
        }

        if (sub_order.items.len != 0) {
            sort_same_strength(hands, sub_order.items, compare_hands);
            order.appendSliceAssumeCapacity(sub_order.items);
        }
    }

    var answer: u32 = 0;
    for (0.., order.items) |i, idx| {
        answer += @as(u32, @intCast(i+1)) * hand_bids.items[idx].bid;
    }

    return .{ .uint = answer };
}

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

    var hand_strengths = try allocator.alloc(HandStrength, hand_bids.items.len);
    defer allocator.free(hand_strengths);

    var hands = try allocator.alloc(Hand, hand_bids.items.len);
    defer allocator.free(hands);

    for (0..hand_bids.items.len) |i| {
        hands[i] = hand_bids.items[i].hand;
        hand_strengths[i] = determine_card_strength_joker(hand_bids.items[i].hand);
    }

    var order = try IndexArray.initCapacity(allocator, hand_bids.items.len);
    defer order.deinit();

    var sub_order = try IndexArray.initCapacity(allocator, hand_bids.items.len);
    defer sub_order.deinit();

    inline for (@typeInfo(HandStrength).Enum.fields) |field| {
        sub_order.clearRetainingCapacity();

        const hand_strength: HandStrength = @enumFromInt(field.value);
        for (0..hand_bids.items.len) |idx| {
            if (hand_strengths[idx] == hand_strength) {
                sub_order.appendAssumeCapacity(idx);
            }
        }

        if (sub_order.items.len != 0) {
            sort_same_strength(hands, sub_order.items, compare_hands_joker);
            order.appendSliceAssumeCapacity(sub_order.items);
        }
    }

    var answer: u32 = 0;
    for (0.., order.items) |i, idx| {
        answer += @as(u32, @intCast(i+1)) * hand_bids.items[idx].bid;
    }

    return .{ .uint = answer };
}

const example_input = [_][]const u8{
    "32T3K 765",
    "T55J5 684",
    "KK677 28",
    "KTJJT 220",
    "QQQJA 483"
};

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

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