const std = @import("std");
const print = std.debug.print;
const assert = std.debug.assert;
const ArrayList = std.ArrayList;
const HashMap = std.HashMap;
const mem = std.mem;

const fin = mem.trim(u8, @embedFile("./input.txt"), &std.ascii.whitespace);
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
const allocator = gpa.allocator();

const Record = struct {
    springs: []const u8,
    arrangement: []const u64,
};

// Credit: https://github.com/RestrictedPower/Advent-of-Code-2023/blob/main/Day12/main.py
fn evaluate(springs: []const u8, arrangement: []const u64) u64 {
    const DPSIZE = 200;
    var dp = [_][DPSIZE][2]u64{[_][2]u64{[_]u64{0} ** 2} ** DPSIZE} ** DPSIZE;

    dp[0][0][0] = 1;

    for (springs, 1..) |c, i| {
        if (c == '.' or c == '?') {
            for (0..arrangement.len + 1) |j| {
                dp[i][j][0] = dp[i - 1][j][0] + dp[i - 1][j][1];
            }
        }

        for (arrangement, 1..) |n, j| {
            const begidx: u64 = if (i >= n) i - n else 0;
            const segment = springs[begidx..i];
            const possible = mem.indexOfScalar(u8, segment, '.') == null;

            if (possible and i >= n) {
                dp[i][j][1] = dp[begidx][j - 1][0];
            }
        }
    }

    return dp[springs.len][arrangement.len][0] +
        dp[springs.len][arrangement.len][1];
}

fn part1(records: []const Record) void {
    var ans: u64 = 0;
    for (records) |record| {
        ans += evaluate(record.springs, record.arrangement);
    }
    print("{d}\n", .{ans});
}

fn part2(records: []const Record) !void {
    var ans: u64 = 0;
    for (records) |record| {
        const unfolded_springs = try allocator.alloc(u8, record.springs.len * 5 + 4);
        _ = try std.fmt.bufPrint(unfolded_springs, "{s}?{s}?{s}?{s}?{s}", .{
            record.springs,
            record.springs,
            record.springs,
            record.springs,
            record.springs,
        });

        var unfolded_arrangement = ArrayList(u64).init(allocator);
        for (0..5) |_| {
            try unfolded_arrangement.appendSlice(record.arrangement);
        }
        ans += evaluate(unfolded_springs, unfolded_arrangement.items);
    }
    print("{d}\n", .{ans});
}

pub fn solve(records: []const Record) !void {
    part1(records);
    try part2(records);
}

pub fn main() !void {
    var splitLines = mem.splitScalar(u8, fin, '\n');
    var records = ArrayList(Record).init(allocator);
    defer records.deinit();

    while (splitLines.next()) |line| {
        var toks = mem.tokenizeAny(u8, line, " ,");

        const springs = toks.next().?;
        var arrangement = ArrayList(u64).init(allocator);

        while (toks.next()) |tok| {
            try arrangement.append(try std.fmt.parseInt(u64, tok, 10));
        }

        try records.append(.{
            .springs = springs,
            .arrangement = arrangement.items,
        });
    }
    try solve(records.items);
}