const std = @import("std");

// We're using pipes as character devices
pub const CharDev = [2]i32;
pub const Word = i64;
pub const MemSize = 1024 * 2;

pub fn loadFromStream(alloc: *std.mem.Allocator, stream: *std.fs.File.InStream.Stream) anyerror![]Word {
    var program = try alloc.alloc(Word, 1024);

    var buf = [_]u8{0} ** 255;
    var i: u32 = 0;

    while (try stream.readUntilDelimiterOrEof(&buf, ',')) |num| {
        var trimmed = std.mem.trimRight(u8, num, "\r\n");
        program[i] = try std.fmt.parseInt(Word, trimmed, 10);
        i += 1;
    }

    std.debug.warn("\n");

    return program[0..program.len];
}

pub fn loadFromStdIn(alloc: *std.mem.Allocator) anyerror![]Word {
    const file = std.io.getStdIn();
    const stream = &file.inStream().stream;

    return loadFromStream(alloc, stream);
}

const Opcode = enum(u8) {
    ADD = 1,
    MULT = 2,
    IN = 3,
    OUT = 4,
    JNZ = 5,
    JZ = 6,
    LT = 7,
    EQ = 8,
    RBO = 9,
    END = 99,
};

const Mode = enum(u8) {
    Position = 0,
    Immediate = 1,
    Relative = 2,
};

const Instruction = struct {
    op: Opcode,

    // Modes for each parameter of the opcode. Parameter 0 is in position 0, etc
    modes: [3]Mode,

    pub fn decode(value: Word) !Instruction {
        assert(value > 0 and value < 99999);

        var buf: [6]u8 = undefined;
        const str = try std.fmt.bufPrint(&buf, "{d:0<6}", value);
        const op = ((str[4] - '0') * 10) + (str[5] - '0');

        return Instruction{
            .op = @intToEnum(Opcode, op),
            .modes = [3]Mode{
                @intToEnum(Mode, buf[3] - '0'),
                @intToEnum(Mode, buf[2] - '0'),
                @intToEnum(Mode, buf[1] - '0'),
            },
        };
    }

    /// Number of bytes taken up by a particular instruction
    pub fn size(self: Instruction) usize {
        return switch (self.op) {
            .ADD => 4,
            .MULT => 4,
            .IN => 2,
            .OUT => 2,
            .JNZ => 3,
            .JZ => 3,
            .LT => 4,
            .EQ => 4,
            .RBO => 2,
            .END => 1,
        };
    }

    pub fn halt(self: Instruction) bool {
        return self.op == .END;
    }
};

pub const Machine = struct {
    alloc: *std.mem.Allocator,
    idx: usize = 0, // handy for debugging

    memory: []Word,
    rb: Word = 0,

    ip: usize = 0,

    halted: bool = false,

    input: CharDev,
    output: CharDev,

    pub fn Run(alloc: *std.mem.Allocator, program: []Word) anyerror!Machine {
        var machine = try init(alloc, program);

        try machine.run();

        return machine;
    }

    pub fn init(alloc: *std.mem.Allocator, program: []Word) anyerror!Machine {
        var memory = try alloc.alloc(Word, MemSize);
        std.mem.set(Word, memory, 0);
        std.mem.copy(Word, memory, program);

        return Machine{
            .alloc = alloc,
            .memory = memory,
            .input = try std.os.pipe(),
            .output = try std.os.pipe(),
        };
    }

    pub fn run(self: *Machine) anyerror!void {
        while (try self.step()) {}
    }

    pub fn deinit(self: *Machine) void {
        self.alloc.free(self.memory);
        std.os.close(self.input[0]);
        std.os.close(self.input[1]);
        std.os.close(self.output[0]);
        std.os.close(self.output[1]);
    }

    pub fn readFromInput(self: *Machine) !Word {
        return self.__in(self.input);
    }

    pub fn writeToInput(self: *Machine, val: Word) !void {
        return self.__out(self.input, val);
    }

    pub fn readFromOutput(self: *Machine) !Word {
        return self.__in(self.output);
    }

    pub fn writeToOutput(self: *Machine, val: Word) !void {
        return self.__out(self.output, val);
    }

    fn __in(self: *Machine, dev: CharDev) !Word {
        const fd = dev[0]; // Read from the pipe
        var store: [8]u8 = undefined;

        // std.debug.warn("{}: __in({} <- {})\n", self.idx, dev[0], dev[1]);
        var n = try std.os.read(fd, &store);
        std.debug.assert(n == 8); // TODO: handle partial reads

        return std.mem.readIntNative(Word, &store);
    }

    fn __out(self: *Machine, dev: CharDev, val: Word) !void {
        const fd = dev[1]; // Write to the pipe
        const bytes = std.mem.asBytes(&val);

        // std.debug.warn("{}: __out({} -> {}, {})\n", self.idx, dev[1], dev[0], val);
        try std.os.write(fd, bytes);
    }

    /// Read an immediate or position value from memory. Parameter determines
    /// which field following the current instruction to read from
    inline fn __read(self: *Machine, parameter: usize, mode: Mode) Word {
        const immediate = self.memory[self.ip + parameter + 1];

        return switch (mode) {
            .Immediate => immediate,
            .Position => self.memory[@intCast(usize, immediate)],
            .Relative => self.memory[@intCast(usize, immediate + self.rb)],
        };
    }

    inline fn __write(self: *Machine, parameter: usize, mode: Mode, value: Word) void {
        const dest = self.__read(parameter, .Immediate);

        switch (mode) {
            .Immediate => unreachable,
            .Position => self.memory[@intCast(usize, dest)] = value,
            .Relative => self.memory[@intCast(usize, dest + self.rb)] = value,
        }
    }

    pub fn step(self: *Machine) anyerror!bool {
        const insn = try Instruction.decode(self.memory[self.ip]);
        const start_ip = self.ip;

        const opcode = switch (insn.op) {
            .ADD => {
                const a = self.__read(0, insn.modes[0]);
                const b = self.__read(1, insn.modes[1]);

                self.__write(2, insn.modes[2], a + b);
            },

            .MULT => {
                const a = self.__read(0, insn.modes[0]);
                const b = self.__read(1, insn.modes[1]);

                self.__write(2, insn.modes[2], a * b);
            },

            .IN => {
                self.__write(0, insn.modes[0], try self.readFromInput()); // blocking read
            },

            .OUT => {
                try self.writeToOutput(self.__read(0, insn.modes[0]));
            },

            .JNZ => {
                if (self.__read(0, insn.modes[0]) != 0)
                    self.ip = @intCast(usize, self.__read(1, insn.modes[1]));
            },

            .JZ => {
                if (self.__read(0, insn.modes[0]) == 0)
                    self.ip = @intCast(usize, self.__read(1, insn.modes[1]));
            },

            .LT => {
                const a = self.__read(0, insn.modes[0]);
                const b = self.__read(1, insn.modes[1]);

                if (a < b) {
                    self.__write(2, insn.modes[2], 1);
                } else {
                    self.__write(2, insn.modes[2], 0);
                }
            },

            .EQ => {
                const a = self.__read(0, insn.modes[0]);
                const b = self.__read(1, insn.modes[1]);

                if (a == b) {
                    self.__write(2, insn.modes[2], 1);
                } else {
                    self.__write(2, insn.modes[2], 0);
                }
            },

            .RBO => {
                self.rb += self.__read(0, insn.modes[0]);
            },

            .END => self.halted = true,
        };

        // Only modify IP if the instruction itself has not
        if (self.ip == start_ip) self.ip += insn.size();

        return !insn.halt();
    }
};

const assert = std.debug.assert;
const test_allocator = std.heap.page_allocator;

test "day 2 example 1" {
    var before: [12]Word = .{ 1, 9, 10, 3, 2, 3, 11, 0, 99, 30, 40, 50 };
    var after: [12]Word = .{ 3500, 9, 10, 70, 2, 3, 11, 0, 99, 30, 40, 50 };
    var machine = try Machine.Run(test_allocator, before[0..before.len]);
    defer machine.deinit();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
}

test "day 2 example 2" {
    var before: [5]Word = .{ 1, 0, 0, 0, 99 };
    var after: [5]Word = .{ 2, 0, 0, 0, 99 };
    var machine = try Machine.Run(test_allocator, before[0..before.len]);
    defer machine.deinit();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
}

test "day 2 example 3" {
    var before: [5]Word = .{ 2, 3, 0, 3, 99 };
    var after: [5]Word = .{ 2, 3, 0, 6, 99 };
    var machine = try Machine.Run(test_allocator, before[0..before.len]);
    defer machine.deinit();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
}

test "day 2 example 4" {
    var before: [6]Word = .{ 2, 4, 4, 5, 99, 0 };
    var after: [6]Word = .{ 2, 4, 4, 5, 99, 9801 };
    var machine = try Machine.Run(test_allocator, before[0..before.len]);
    defer machine.deinit();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
}

test "day 2 example 5" {
    var before: [9]Word = .{ 1, 1, 1, 4, 99, 5, 6, 0, 99 };
    var after: [9]Word = .{ 30, 1, 1, 4, 2, 5, 6, 0, 99 };
    var machine = try Machine.Run(test_allocator, before[0..before.len]);
    defer machine.deinit();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
}

test "day 5 example 1" {
    var before: [5]Word = .{ 3, 0, 4, 0, 99 };
    var after: [5]Word = .{ 666, 0, 4, 0, 99 };
    var machine = try Machine.init(test_allocator, before[0..before.len]);
    defer machine.deinit();

    try machine.writeToInput(666);
    try machine.run();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
    assert(666 == try machine.readFromOutput());
}

test "day 9 example 1" {
    var before: [16]Word = .{ 109, 1, 204, -1, 1001, 100, 1, 100, 1008, 100, 16, 101, 1006, 101, 0, 99 };
    var after: [16]Word = before;
    var output: [16]Word = undefined;

    var machine = try Machine.init(test_allocator, before[0..before.len]);
    defer machine.deinit();
    try machine.run();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));

    for (output) |*b| {
        b.* = try machine.readFromOutput();
    }

    // This program produces a copy of itself as output
    assert(std.mem.eql(Word, machine.memory[0..before.len], output[0..after.len]));
}

test "day 9 example 2" {
    var before: [8]Word = .{ 1102, 34915192, 34915192, 7, 4, 7, 99, 0 };
    var after: [8]Word = .{ 1102, 34915192, 34915192, 7, 4, 7, 99, 1219070632396864 };

    var machine = try Machine.init(test_allocator, before[0..before.len]);
    defer machine.deinit();
    try machine.run();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
    std.debug.assert(1219070632396864 == try machine.readFromOutput());
}

test "day 9 example 3" {
    var before: [3]Word = .{ 104, 1125899906842624, 99 };
    var after: [3]Word = before;

    var machine = try Machine.init(test_allocator, before[0..before.len]);
    defer machine.deinit();
    try machine.run();

    assert(std.mem.eql(Word, machine.memory[0..before.len], after[0..after.len]));
    std.debug.assert(1125899906842624 == try machine.readFromOutput());
}

test "Instruction#decode ADD" {
    const insn = try Instruction.decode(1);

    assert(insn.op == Opcode.ADD);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Position);
}

test "Instruction#decode MULT" {
    const insn = try Instruction.decode(2);

    assert(insn.op == Opcode.MULT);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Position);
}

test "Instruction#decode IN" {
    const insn = try Instruction.decode(3);

    assert(insn.op == Opcode.IN);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Position);
}

test "Instruction#decode OUT" {
    const insn = try Instruction.decode(4);

    assert(insn.op == Opcode.OUT);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Position);
}

test "Instruction#decode END Position" {
    const insn = try Instruction.decode(99);

    assert(insn.op == Opcode.END);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Position);
}

test "Instruction#decode END Immediate" {
    const insn = try Instruction.decode(11199);

    assert(insn.op == Opcode.END);
    assert(insn.modes[0] == .Immediate);
    assert(insn.modes[1] == .Immediate);
    assert(insn.modes[2] == .Immediate);
}

test "Instruction#decode END mixed" {
    const insn = try Instruction.decode(10099);

    assert(insn.op == Opcode.END);
    assert(insn.modes[0] == .Position);
    assert(insn.modes[1] == .Position);
    assert(insn.modes[2] == .Immediate);
}