const std = @import("std");

const path = "data/day12/input.txt";

const RetType = u20;
const Str = []const u8;

const CaveID = u12; // there are 11 caves in input + 1 for the doble cave bit
const CaveInfo = struct {
    ID: CaveID,
    paths: std.ArrayList(Str),
};

const Routes = std.StringHashMap(CaveInfo);

fn parseInput(a: std.mem.Allocator) anyerror!Routes {
    const input = @embedFile(path);
    var lines = std.mem.tokenize(u8, input, "\n");

    var ret = Routes.init(a);

    var counter: u4 = 0;
    while (lines.next()) |line| {
        var from_to = std.mem.tokenize(u8, line, "-");
        const left = from_to.next().?;
        const right = from_to.next().?;

        if (!std.mem.eql(u8, right, "start") and (!std.mem.eql(u8, left, "end"))) {
            var left_array = try ret.getOrPut(left);
            if (!left_array.found_existing) {
                left_array.value_ptr.*.paths = std.ArrayList(Str).init(a);
                left_array.value_ptr.*.ID = @as(CaveID, 1) << counter;
                counter += 1;
            }
            try left_array.value_ptr.paths.append(right);
        }

        if (!std.mem.eql(u8, left, "start") and (!std.mem.eql(u8, right, "end"))) {
            var right_array = try ret.getOrPut(right);
            if (!right_array.found_existing) {
                right_array.value_ptr.*.paths = std.ArrayList(Str).init(a);
                right_array.value_ptr.*.ID = @as(CaveID, 1) << counter;
                counter += 1;
            }
            try right_array.value_ptr.paths.append(left);
        }
    }

    return ret;
}

const CacheKey = struct {
    root: CaveID,
    seen: CaveID,
};

const Cache = std.AutoHashMap(CacheKey, RetType);

pub fn second(allocator: ?std.mem.Allocator) anyerror!RetType {
    var cave = try parseInput(allocator.?);

    var seen: CaveID = 0;

    var cache = Cache.init(allocator.?);

    defer {
        var it = cave.valueIterator();
        while (it.next()) |arr_list| {
            arr_list.paths.deinit();
        }
        cave.deinit();
        cache.deinit();
    }

    return try genRoutes(seen, &cave, "start", &cache);
}

const double_cave = @as(CaveID, 1) << 11;

fn genRoutes(seen: CaveID, routes: *Routes, root: Str, cache: *Cache) anyerror!RetType {
    var sum: RetType = 0;

    const root_item = routes.get(root).?;
    if (cache.get(CacheKey{ .root = root_item.ID, .seen = seen })) |r| {
        return r;
    }

    for (root_item.paths.items) |item| {
        if (std.mem.eql(u8, item, "end")) {
            sum += 1;
            continue;
        }

        const next_item = routes.get(item).?;

        if (seen & next_item.ID == 0) {
            var mark = seen;
            if (root[0] > 96) mark |= root_item.ID;
            sum += try genRoutes(mark, routes, item, cache);
        } else if (seen & double_cave == 0) {
            var mark = seen;
            if (root[0] > 96) mark |= root_item.ID;
            sum += try genRoutes(mark | double_cave, routes, item, cache);
        }
    }

    try cache.put(CacheKey{ .root = root_item.ID, .seen = seen }, sum);

    return sum;
}

pub fn main() anyerror!void {
    var buf: [25_000]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buf);

    var timer = try std.time.Timer.start();
    const ret = try second(fba.allocator());
    const t = timer.lap() / 1000;

    try std.testing.expectEqual(@as(RetType, 134862), ret);

    std.debug.print("Day 12b result: {d} \t\ttime: {d}us\n", .{ ret, t });
}

test "day12b" {
    try std.testing.expectEqual(@as(RetType, 134862), try second(std.testing.allocator));
}