const std = @import("std");

const Score = enum { FiveOfAKind, FourOfAKind, FullHouse, ThreeOfAKind, TwoPair, OnePair, HighCard };
const Hand = struct { cards: std.ArrayList(u8), score: Score, bid: u16 };

pub fn main() !void {
    const content = @embedFile("data.txt");
    var lines = std.mem.split(u8, content, "\n");

    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    var hands = std.ArrayList(Hand).init(allocator);
    defer hands.deinit();

    while (lines.next()) |l| {
        if (l.len == 0) { // TODO: really gotta figure this out
            break;
        }

        var cards = std.ArrayList(u8).init(allocator);
        for (0..5) |i| {
            try cards.append(cardValue(l[i]));
        }

        const bid = try std.fmt.parseInt(u16, l[6..], 10);
        const hand = Hand{ .cards = cards, .score = try findScore(cards), .bid = bid };

        try hands.append(hand);
    }

    // sort hands by score
    var sortedHands = try hands.toOwnedSlice();
    std.sort.insertion(Hand, sortedHands, {}, compareHands);

    // sum up winnings
    var winnings: u32 = 0;
    for (sortedHands, 1..) |hand, i| {
        winnings += hand.bid * @as(u32, @intCast(i));
    }

    std.debug.print("Part 1 {?d}\n", .{winnings});

    // Part 2
    var hands2 = std.ArrayList(Hand).init(allocator);
    defer hands2.deinit();

    // Augment the scores
    for (sortedHands) |hand| {
        try hands2.append(try augmentHand(hand, allocator));
    }

    // sort hands by score
    var sortedHands2 = try hands2.toOwnedSlice();
    std.sort.insertion(Hand, sortedHands2, {}, compareHands);

    // sum up winnings
    var winnings2: u32 = 0;
    for (sortedHands2, 1..) |hand, i| {
        winnings2 += hand.bid * @as(u32, @intCast(i));
    }

    std.debug.print("Part 2 {?d}\n", .{winnings2});
}

// for part 2, update the score now that J is wild
fn augmentHand(hand: Hand, allocator: std.mem.Allocator) !Hand {
    var numJokers: u8 = 0;
    for (hand.cards.items) |c| {
        if (c == 11) {
            numJokers += 1;
        }
    }

    var cards = std.ArrayList(u8).init(allocator);
    for (hand.cards.items) |card| {
        const newCard = if (card == 11) 1 else card;
        try cards.append(newCard);
    }

    if (numJokers == 0) {
        return Hand{ .cards = hand.cards, .score = hand.score, .bid = hand.bid };
    }

    var score = switch (hand.score) {
        Score.FiveOfAKind => Score.FiveOfAKind,
        Score.FourOfAKind => Score.FiveOfAKind,
        Score.FullHouse => Score.FiveOfAKind,
        Score.ThreeOfAKind => Score.FourOfAKind,
        Score.TwoPair => switch (numJokers) {
            2 => Score.FourOfAKind,
            else => Score.FullHouse,
        },
        Score.OnePair => Score.ThreeOfAKind,
        Score.HighCard => Score.OnePair,
    };

    return Hand{ .cards = cards, .score = score, .bid = hand.bid };
}

fn compareHands(_: void, a: Hand, b: Hand) bool {
    const aScore = @intFromEnum(a.score);
    const bScore = @intFromEnum(b.score);
    if (aScore > bScore) {
        return true;
    } else if (aScore < bScore) {
        return false;
    } else {
        for (0..5) |i| {
            const aVal = a.cards.items[i];
            const bVal = b.cards.items[i];
            if (aVal == bVal) {
                continue;
            }
            return aVal < bVal;
        }
    }
    return true;
}

// Convert A, K, Q, etc. to 14, 13, 12, etc.
fn cardValue(c: u8) u8 {
    return switch (c) {
        'A' => 14,
        'K' => 13,
        'Q' => 12,
        'J' => 11,
        'T' => 10,
        else => c - 48,
    };
}

fn findScore(cards: std.ArrayList(u8)) !Score {
    // make a copy and sort it
    var copy = try cards.clone();
    std.sort.insertion(u8, copy.items, {}, comptime std.sort.asc(u8));
    var c = copy.items; // so much typing

    // I thought about using a map of some kind to count pairs.. but this isn't too bad
    if (std.mem.eql(u8, c[0..4], c[1..5])) {
        return Score.FiveOfAKind;
    } else if ((c[0] == c[1] and c[1] == c[2] and c[2] == c[3]) or (c[1] == c[2] and c[2] == c[3] and c[3] == c[4])) {
        return Score.FourOfAKind;
    } else if ((c[0] == c[1] and c[1] == c[2] and c[3] == c[4]) or (c[2] == c[3] and c[3] == c[4] and c[0] == c[1])) {
        return Score.FullHouse;
    } else if ((c[0] == c[1] and c[1] == c[2]) or (c[1] == c[2] and c[2] == c[3]) or (c[2] == c[3] and c[3] == c[4])) {
        return Score.ThreeOfAKind;
    } else if ((c[0] == c[1] and c[2] == c[3]) or (c[0] == c[1] and c[3] == c[4]) or (c[1] == c[2] and c[3] == c[4])) {
        return Score.TwoPair;
    } else if (c[0] == c[1] or c[1] == c[2] or c[2] == c[3] or c[3] == c[4]) {
        return Score.OnePair;
    } else {
        return Score.HighCard;
    }
}