const std = @import("std");

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

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

fn parseInput(a: std.mem.Allocator) ![]u8 {
    const input = @embedFile(path);

    var ret = std.ArrayList(u8).init(a);

    for (input) |ch| {
        switch (ch) {
            '0' => try ret.appendSlice("0000"),
            '1' => try ret.appendSlice("0001"),
            '2' => try ret.appendSlice("0010"),
            '3' => try ret.appendSlice("0011"),
            '4' => try ret.appendSlice("0100"),
            '5' => try ret.appendSlice("0101"),
            '6' => try ret.appendSlice("0110"),
            '7' => try ret.appendSlice("0111"),
            '8' => try ret.appendSlice("1000"),
            '9' => try ret.appendSlice("1001"),
            'A' => try ret.appendSlice("1010"),
            'B' => try ret.appendSlice("1011"),
            'C' => try ret.appendSlice("1100"),
            'D' => try ret.appendSlice("1101"),
            'E' => try ret.appendSlice("1110"),
            'F' => try ret.appendSlice("1111"),
            else => break,
        }
    }

    return ret.toOwnedSlice();
}

pub fn first(allocator: ?std.mem.Allocator) !RetType {
    var input = try parseInput(allocator.?);
    defer allocator.?.free(input);

    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    const p = try parsePacket(arena.allocator(), &input);

    return p.sumVersion();
}

const LiteralType = u44;

const Packet = struct {
    version: u3,
    typeID: u3,
    literal: LiteralType,
    subs: std.ArrayList(Packet),
    fn sumVersion(self: @This()) RetType {
        var sum: RetType = self.version;
        for (self.subs.items) |sp| {
            sum += sp.sumVersion();
        }
        return sum;
    }
    fn parseOperator(self: *@This(), a: std.mem.Allocator, in: *Str) anyerror!void {
        in.* = in.*[6..];
        switch (in.*[0]) {
            '0' => {
                const total = try std.fmt.parseUnsigned(usize, in.*[1..16], 2);

                in.* = in.*[16..];
                const initial_length = in.len;

                var pivot: usize = 0;
                while (pivot < total) : (pivot = initial_length - in.len) {
                    var p = try parsePacket(a, in);
                    try self.subs.append(p);
                }
            },
            '1' => {
                const subpackets = try std.fmt.parseUnsigned(usize, in.*[1..12], 2);

                in.* = in.*[12..];

                var counter: usize = 0;
                while (counter < subpackets) : (counter += 1) {
                    var p = try parsePacket(a, in);
                    try self.subs.append(p);
                }
            },
            else => unreachable,
        }
    }
};

fn parsePacket(a: std.mem.Allocator, in: *Str) !Packet {
    var p = Packet{ .version = undefined, .typeID = undefined, .literal = undefined, .subs = std.ArrayList(Packet).init(a) };

    p.version = try std.fmt.parseUnsigned(u3, in.*[0..3], 2);
    p.typeID = try std.fmt.parseUnsigned(u3, in.*[3..6], 2);

    if (p.typeID == 4) {
        p.literal = try parseLiteral(a, in);
    } else {
        try p.parseOperator(a, in);
    }

    return p;
}

fn parseLiteral(a: std.mem.Allocator, in: *Str) !LiteralType {
    in.* = in.*[6..];
    var ret = std.ArrayList(u8).init(a);
    var idx: usize = 0;
    while (idx < in.len) : (idx += 5) {
        try ret.appendSlice(in.*[idx + 1 .. idx + 5]);
        if (in.*[idx] == '0') {
            in.* = in.*[idx + 5 ..];
            break;
        }
    }
    return try std.fmt.parseUnsigned(LiteralType, ret.items, 2);
}

pub fn main() !void {
    var buf: [100_000]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buf);
    var arena = std.heap.ArenaAllocator.init(fba.allocator());
    defer arena.deinit();

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

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

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

test "day16a" {
    try std.testing.expectEqual(@as(RetType, 967), try first(std.testing.allocator));
}