const std = @import("std");

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

const retSize = u24;
const lineSize = getLineSize();

fn first() anyerror!retSize {
    const input = @embedFile(path);
    var lines = std.mem.split(u8, input, "\n");

    var ones: [lineSize]retSize = undefined;
    var zeros: [lineSize]retSize = undefined;

    while (lines.next()) |line| {
        for (line) |bit, idx| {
            if (bit == '0') {
                zeros[idx] += 1;
            } else {
                ones[idx] += 1;
            }
        }
    }

    var gammaRate: [lineSize]u8 = undefined;
    var epsilonRate: [lineSize]u8 = undefined;

    for (ones) |val, i| {
        if (val > zeros[i]) {
            gammaRate[i] = '1';
            epsilonRate[i] = '0';
        } else {
            gammaRate[i] = '0';
            epsilonRate[i] = '1';
        }
    }

    const gr = try std.fmt.parseUnsigned(retSize, &gammaRate, 2);
    const er = try std.fmt.parseUnsigned(retSize, &epsilonRate, 2);

    return gr * er;
}

fn second() anyerror!retSize {
    const input = @embedFile(path);
    var lines = std.mem.tokenize(u8, input, "\n");

    var buffer: [1000 * lineSize * 100]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buffer);
    const allocator = fba.allocator();

    var oxigen: std.BufSet = std.BufSet.init(allocator);
    defer oxigen.deinit();

    var co2: std.BufSet = std.BufSet.init(allocator);
    defer co2.deinit();

    while (lines.next()) |line| {
        try oxigen.insert(line);
        try co2.insert(line);
    }

    const ox = try getOxigenRating(&oxigen);
    const co = try getCO2Rating(&co2);

    return ox * co;
}

fn getOxigenRating(oxigen: *std.BufSet) anyerror!retSize {
    var idx: usize = 0;
    while (oxigen.count() > 1) : (idx += 1) {
        var ones: retSize = 0;
        var zeros: retSize = 0;

        // count ones and zeros
        var it = oxigen.iterator();
        while (it.next()) |line| {
            if (line.*[idx] == '1') {
                ones += 1;
            } else if (line.*[idx] == '0') {
                zeros += 1;
            } else unreachable;
        }

        it = oxigen.iterator(); // reset iterator
        while (it.next()) |line| {
            if (ones >= zeros) {
                if (line.*[idx] != '1') {
                    oxigen.remove(line.*);
                }
            } else {
                if (line.*[idx] != '0') {
                    oxigen.remove(line.*);
                }
            }
        }
    }

    return try std.fmt.parseUnsigned(retSize, oxigen.iterator().next().?.*, 2);
}

fn getCO2Rating(co2: *std.BufSet) anyerror!retSize {
    var idx: usize = 0;
    while (co2.count() > 1) : (idx += 1) {
        var ones: retSize = 0;
        var zeros: retSize = 0;

        // count ones and zeros
        var it = co2.iterator();
        while (it.next()) |line| {
            if (line.*[idx] == '1') {
                ones += 1;
            } else if (line.*[idx] == '0') {
                zeros += 1;
            } else unreachable;
        }

        it = co2.iterator(); // reset iterator
        while (it.next()) |line| {
            if (zeros <= ones) {
                if (line.*[idx] != '0') {
                    co2.remove(line.*);
                }
            } else {
                if (line.*[idx] != '1') {
                    co2.remove(line.*);
                }
            }
        }
    }

    return try std.fmt.parseUnsigned(retSize, co2.iterator().next().?.*, 2);
}

fn getLineSize() usize {
    const input = @embedFile(path);

    var ret: usize = 0;
    for (input) |bit, idx| {
        if (bit == '\n') {
            ret = idx;
            break;
        }
    }

    return ret;
}

pub fn main() anyerror!void {
    var timer = try std.time.Timer.start();

    _ = try first();
    const f = timer.lap() / 1000;
    _ = try second();
    const s = timer.lap() / 1000;

    std.debug.print("Day 3 \t\tfirst: {d}us \t\tsecond: {d}us\n", .{ f, s });
}

test "first" {
    try std.testing.expectEqual(@as(retSize, 3847100), try first(std.testing.allocator));
}

test "second" {
    try std.testing.expectEqual(@as(retSize, 4105235), try second());
}