//brainf
const std = @import("std");
const io = std.io;
const fs = std.fs;
const mem = std.mem;
const process = std.process;
const testing = std.testing;
const ArrayList = std.ArrayList;

const Level1 = union(enum) {
    // think about it... i hope
    inc: i8,
    shift: i16,
    // u8 because why would we be reading to/writing from the same cell 256 times in a row
    // well, might as well use the extra space (union is bigger than its biggest member)
    in: u16,
    out: u16,

    // u24s because i arbitrarily chose 2^24 as the limit
    // maybe u32 would be more efficient idk this union should fit in a u32
    // hm we can probably go way lower bc level1 is so much more compact than prog_unfiltered
    // shouldn't be necessary
    loop_start: u24,
    loop_end: u24,

    zero,

    const Self = @This();

    fn pointless(self: Self) bool {
        return switch (self) {
            .inc, .shift => |n| n == 0,
            .in, .out => |n| n == 0,
            else => false,
        };
    }

    fn nothing() Self {
        return .{ .inc = 0 };
    }

    // the default fmt is human and machine readable..... but idc  xd lol    pog
    pub fn format(
        self: Self,
        comptime fmt: []const u8,
        options: std.fmt.FormatOptions,
        writer: anytype,
    ) !void {
        _ = fmt;
        _ = options;

        switch (self) {
            .inc => |n| try writer.print("INC {d: >4}", .{n}),
            .shift => |n| try writer.print("SHF {d: >4}", .{n}),
            .in => |n| try writer.print("INP {d: >4}", .{n}),
            .out => |n| try writer.print("OUT {d: >4}", .{n}),
            .loop_start => |n| try writer.print("JEZ {d:0>4}", .{n}),
            .loop_end => |n| try writer.print("JNZ {d:0>4}", .{n}),
            .zero => try writer.print("ZER     ", .{}),
        }
    }
};

// mandelbrot.b has 120 instances of [-] which can be replaced with a memory[mp] = 0
// dang it! gcc -O3 with code generated by dbf2c.b finishes in like 1.4s

// we'll do the discarding of comment chars here
fn compileLevel1(allocator: mem.Allocator, src: []const u8) !ArrayList(Level1) {
    var cur = Level1.nothing();
    var out = ArrayList(Level1).init(allocator);
    var openers = ArrayList(u24).init(allocator);
    errdefer {
        out.deinit();
        openers.deinit();
    }
    var skip: u8 = 0;

    for (src, 0..) |b, i| {
        if (skip > 0) {
            skip -= 1;
            continue;
        }

        switch (b) {
            '+', '-' => {
                const n: i8 = if (b == '+') 1 else -1;
                switch (cur) {
                    .inc => |*ns| ns.* +%= n,
                    else => {
                        if (!cur.pointless()) try out.append(cur);
                        cur = .{ .inc = n };
                    },
                }
            },
            '>', '<' => {
                const n: i16 = if (b == '>') 1 else -1;
                switch (cur) {
                    .shift => |*ns| ns.* +%= n,
                    else => {
                        if (!cur.pointless()) try out.append(cur);
                        cur = .{ .shift = n };
                    },
                }
            },
            '[' => {
                if (!cur.pointless()) try out.append(cur);
                // bad (what if space yk)
                if ((src[i + 1] == '-' or src[i + 1] == '+') and src[i + 2] == ']') {
                    try out.append(.zero);
                    skip = 2;
                    cur = Level1.nothing();
                } else {
                    cur = .{ .loop_start = 0 };
                    try openers.append(@as(u24, @truncate(out.items.len)));
                }
            },
            ']' => {
                if (!cur.pointless()) try out.append(cur);
                const opener_idx = openers.popOrNull() orelse return error.UnopenedCloser;
                out.items[opener_idx].loop_start = @as(u24, @truncate(out.items.len));
                cur = .{ .loop_end = opener_idx };
            },
            ',' => {
                switch (cur) {
                    .in => |*ns| ns.* += 1,
                    else => {
                        if (!cur.pointless()) try out.append(cur);
                        cur = .{ .in = 1 };
                    },
                }
            },
            '.' => {
                switch (cur) {
                    .out => |*ns| ns.* += 1,
                    else => {
                        if (!cur.pointless()) try out.append(cur);
                        cur = .{ .out = 1 };
                    },
                }
            },
            else => {},
        }
    }

    if (!cur.pointless()) try out.append(cur);

    const unclosed_openers = openers.items.len;
    if (unclosed_openers > 0) return error.UnclosedOpener;
    openers.deinit();

    return out;
}

fn testLevel1(src: []const u8, expected: []const Level1) !void {
    const level1 = try compileLevel1(testing.allocator, src);
    defer level1.deinit();

    errdefer {
        std.debug.print("{s}\n", .{src});
        std.debug.print("      OBTAINED  EXPECTED\n", .{});
        if (level1.items.len < expected.len) {
            for (level1.items, expected[0..level1.items.len], 0..) |a, b, i| {
                std.debug.print("{:0>4}: {}  {}\n", .{ i, a, b });
            }
            for (expected[level1.items.len..], level1.items.len..) |b, i| {
                std.debug.print("{:0>4}: --------  {}\n", .{ i, b });
            }
        } else {
            for (level1.items[0..expected.len], expected, 0..) |a, b, i| {
                std.debug.print("{:0>4}: {}  {}\n", .{ i, a, b });
            }
            for (level1.items[expected.len..], expected.len..) |a, i| {
                std.debug.print("{:0>4}: {}  --------\n", .{ i, a });
            }
        }
    }

    try testing.expectEqual(expected.len, level1.items.len);

    for (
        level1.items,
        expected,
    ) |a, b| {
        try testing.expectEqual(a, b);
    }
}

test "level 1" {
    try testLevel1(
        "+++-- [ >><<< +-  ] ...,,,",
        &[_]Level1{
            .{ .inc = 1 },
            .{ .loop_start = 3 },
            .{ .shift = -1 },
            .{ .loop_end = 1 },
            .{ .out = 3 },
            .{ .in = 3 },
        },
    );
    try testLevel1(
        \\ 0 [
        \\ 1   [
        \\ 2   ]
        \\ 3   [
        \\ 4     [
        \\ 5     ]
        \\ 6     [
        \\ 7     ]
        \\ 8   ]
        \\ 9 ]
    ,
        &[_]Level1{
            .{ .loop_start = 9 },
            .{ .loop_start = 2 },
            .{ .loop_end = 1 },
            .{ .loop_start = 8 },
            .{ .loop_start = 5 },
            .{ .loop_end = 4 },
            .{ .loop_start = 7 },
            .{ .loop_end = 6 },
            .{ .loop_end = 3 },
            .{ .loop_end = 0 },
        },
    );

    testLevel1(
        "bla bla - [ ++ ",
        &[_]Level1{},
    ) catch |err| try testing.expectEqual(err, error.UnclosedOpener);
    testLevel1(
        "wow wow += ]",
        &[_]Level1{},
    ) catch |err| try testing.expectEqual(err, error.UnopenedCloser);

    try testLevel1(
        "+++[>>[-]-+]",
        &[_]Level1{
            .{ .inc = 3 },
            .{ .loop_start = 4 },
            .{ .shift = 2 },
            .zero,
            .{ .loop_end = 1 },
        },
    );

    try testLevel1(
        "[>][<][-]",
        &[_]Level1{
            .find_zero_right,
            .find_zero_left,
            .zero,
        },
    );
}

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    var allocator = arena.allocator();
    defer arena.deinit();

    var out = io.getStdOut().writer();
    var in = io.getStdIn().reader();

    var args = process.args();
    defer args.deinit();
    _ = args.skip();

    const prog_path = args.next() orelse {
        try io.getStdErr().writer().writeAll("i need a program path :(\n");
        process.exit(1);
    };

    var prog_file = try fs.cwd().openFile(prog_path, .{});
    const src = try prog_file.readToEndAlloc(allocator, 1 << 24);
    const prog_arr = try compileLevel1(allocator, src);
    const prog = prog_arr.items;
    const prog_len = prog.len;

    var memory = try allocator.alloc(u8, 1 << 16);
    for (memory) |*m| {
        m.* = 0;
    }
    var mp: u16 = 0;
    var ip: u32 = 0;

    // should definitely add a command to print the cool asm
    // std.debug.print("{s}\n", .{src});
    // for (prog, 0..) |ins, i| {
    //     std.debug.print("{:0>4}: {}\n", .{ i, ins });
    // }

    // inc shift in out loop_start loop_end
    while (ip < prog_len) {
        switch (prog[ip]) {
            .inc => |n| memory[mp] +%= @bitCast(n),
            .shift => |n| mp +%= @bitCast(n),
            .in => |n| {
                for (0..(n - 1)) |_| {
                    _ = in.readByte() catch 0;
                }
                memory[mp] = in.readByte() catch 0;
            },
            .out => |n| {
                try out.writeByteNTimes(memory[mp], n);
            },
            .loop_start => |new_ip| {
                if (memory[mp] == 0) ip = new_ip;
            },
            .loop_end => |new_ip| {
                if (memory[mp] != 0) ip = new_ip;
            },
            .zero => memory[mp] = 0,
        }
        ip += 1;
    }
}