//! A comptime assembler.
//!
//! Not at all ready for real-world use; this is only designed
//! to make the test cases easier to write.
const std = @import("std");

pub fn assemble(
    in: []const u8,
    writer: *std.Io.Writer,
) !void {
    var it = std.mem.tokenizeScalar(u8, in, '\n');

    while (it.next()) |l| {
        var line = l;
        if (std.mem.indexOfScalar(u8, line, ';')) |semicolon|
            line = line[0..semicolon];
        line = std.mem.trim(u8, line, whitespace);

        var tokens = std.mem.tokenizeScalar(u8, line, ' ');

        const opcode_s = tokens.next() orelse continue;
        const opcode: Opcode = for (std.meta.tags(Opcode)) |op| {
            if (std.ascii.eqlIgnoreCase(@tagName(op), opcode_s)) break op;
        } else return error.InvalidOpcode;

        const rest = tokens.rest();
        const operand: Operand = if (rest.len > 0) try .parse(rest) else .implied;
        const op_byte = opcode.encode(operand) orelse return error.InvalidPairing;

        try writer.writeByte(op_byte);
        switch (operand) {
            .implied => {},

            .immediate,
            .zeropage,
            .zeropage_x,
            .zeropage_y,
            .indexed_indirect,
            .indirect_indexed,
            => |v| try writer.writeByte(v),

            .absolute,
            .absolute_x,
            .absolute_y,
            .indirect,
            => |v| try writer.writeInt(u16, v, .little),
        }
    }
}

pub fn assembleComptime(comptime in: []const u8) []const u8 {
    // Heuristically defined
    @setEvalBranchQuota(in.len * 500);

    return comptime v: {
        var count_w: std.Io.Writer.Discarding = .init(&.{});
        assemble(in, &count_w.writer) catch |err| @compileError(std.fmt.comptimePrint("error: {}", .{err}));
        var buf: [count_w.fullCount()]u8 = undefined;
        var w: std.Io.Writer = .fixed(&buf);
        assemble(in, &w) catch unreachable;

        const result = buf;
        break :v &result;
    };
}

pub const Opcode = enum {
    brk,
    lda,
    ldx,
    ldy,
    sta,
    stx,
    sty,

    php,
    plp,
    pha,
    pla,

    ora,
    @"and",
    eor,
    adc,
    sbc,
    asl,
    lsr,
    rol,
    ror,
    inc,
    dec,
    inx,
    dex,
    iny,
    dey,
    cmp,
    cpx,
    cpy,
    bit,

    clc,
    sec,
    cli,
    sei,
    clv,
    cld,
    sed,

    bpl,
    bmi,
    bvc,
    bvs,
    bcc,
    bcs,
    bne,
    beq,
    jmp,
    jsr,
    rts,

    nop,

    fn encode(op: Opcode, opr: Operand) ?u8 {
        return switch (op) {
            .brk => switch (opr) {
                // A common extension is to allow BRK to take an operand,
                // though it is usually completely ignored.
                .implied, .immediate => 0x00,
                else => null,
            },
            .lda => switch (opr) {
                .zeropage => 0xa5,
                .zeropage_x => 0xb5,
                .absolute => 0xad,
                .absolute_x => 0xbd,
                .absolute_y => 0xb9,
                .immediate => 0xa9,
                .indexed_indirect => 0xa1,
                .indirect_indexed => 0xb1,
                else => null,
            },
            .ldx => switch (opr) {
                .zeropage => 0xa6,
                .zeropage_y => 0xb6,
                .absolute => 0xae,
                .absolute_y => 0xbe,
                .immediate => 0xa2,
                else => null,
            },
            .ldy => switch (opr) {
                .zeropage => 0xa4,
                .zeropage_x => 0xb4,
                .absolute => 0xac,
                .absolute_x => 0xbc,
                .immediate => 0xa0,
                else => null,
            },
            .sta => switch (opr) {
                .zeropage => 0x85,
                .zeropage_x => 0x95,
                .absolute => 0x8d,
                .absolute_x => 0x9d,
                .absolute_y => 0x99,
                .indexed_indirect => 0x81,
                .indirect_indexed => 0x91,
                else => null,
            },
            .stx => switch (opr) {
                .zeropage => 0x86,
                .zeropage_y => 0x96,
                .absolute => 0x8e,
                else => null,
            },
            .sty => switch (opr) {
                .zeropage => 0x84,
                .zeropage_y => 0x94,
                .absolute => 0x8c,
                else => null,
            },
            .ora => switch (opr) {
                .zeropage => 0x05,
                .zeropage_x => 0x15,
                .absolute => 0x0d,
                .absolute_x => 0x1d,
                .absolute_y => 0x19,
                .immediate => 0x09,
                .indexed_indirect => 0x01,
                .indirect_indexed => 0x11,
                else => null,
            },
            .@"and" => switch (opr) {
                .zeropage => 0x25,
                .zeropage_x => 0x35,
                .absolute => 0x2d,
                .absolute_x => 0x3d,
                .absolute_y => 0x39,
                .immediate => 0x29,
                .indexed_indirect => 0x21,
                .indirect_indexed => 0x31,
                else => null,
            },
            .eor => switch (opr) {
                .zeropage => 0x45,
                .zeropage_x => 0x55,
                .absolute => 0x4d,
                .absolute_x => 0x5d,
                .absolute_y => 0x59,
                .immediate => 0x49,
                .indexed_indirect => 0x41,
                .indirect_indexed => 0x51,
                else => null,
            },
            .adc => switch (opr) {
                .zeropage => 0x65,
                .zeropage_x => 0x75,
                .absolute => 0x6d,
                .absolute_x => 0x7d,
                .absolute_y => 0x79,
                .immediate => 0x69,
                .indexed_indirect => 0x61,
                .indirect_indexed => 0x71,
                else => null,
            },
            .sbc => switch (opr) {
                .zeropage => 0xe5,
                .zeropage_x => 0xf5,
                .absolute => 0xed,
                .absolute_x => 0xfd,
                .absolute_y => 0xf9,
                .immediate => 0xe9,
                .indexed_indirect => 0xe1,
                .indirect_indexed => 0xf1,
                else => null,
            },
            .asl => switch (opr) {
                .zeropage => 0x06,
                .zeropage_x => 0x16,
                .absolute => 0x0e,
                .absolute_x => 0x1e,
                .implied => 0x0a,
                else => null,
            },
            .lsr => switch (opr) {
                .zeropage => 0x46,
                .zeropage_x => 0x56,
                .absolute => 0x4e,
                .absolute_x => 0x5e,
                .implied => 0x4a,
                else => null,
            },
            .rol => switch (opr) {
                .zeropage => 0x26,
                .zeropage_x => 0x36,
                .absolute => 0x2e,
                .absolute_x => 0x3e,
                .implied => 0x2a,
                else => null,
            },
            .ror => switch (opr) {
                .zeropage => 0x66,
                .zeropage_x => 0x76,
                .absolute => 0x6e,
                .absolute_x => 0x7e,
                .implied => 0x6a,
                else => null,
            },
            .inc => switch (opr) {
                .zeropage => 0xe6,
                .zeropage_x => 0xf6,
                .absolute => 0xee,
                .absolute_x => 0xfe,
                else => null,
            },
            .dec => switch (opr) {
                .zeropage => 0xc6,
                .zeropage_x => 0xd6,
                .absolute => 0xce,
                .absolute_x => 0xde,
                else => null,
            },
            .inx => switch (opr) {
                .implied => 0xe8,
                else => null,
            },
            .dex => switch (opr) {
                .implied => 0xca,
                else => null,
            },
            .iny => switch (opr) {
                .implied => 0xc8,
                else => null,
            },
            .dey => switch (opr) {
                .implied => 0x88,
                else => null,
            },
            .cmp => switch (opr) {
                .zeropage => 0xc5,
                .zeropage_x => 0xd5,
                .absolute => 0xcd,
                .absolute_x => 0xdd,
                .absolute_y => 0xd9,
                .immediate => 0xc9,
                .indexed_indirect => 0xc1,
                .indirect_indexed => 0xd1,
                else => null,
            },
            .cpx => switch (opr) {
                .immediate => 0xe0,
                .zeropage => 0xe4,
                .absolute => 0xec,
                else => null,
            },
            .cpy => switch (opr) {
                .immediate => 0xc0,
                .zeropage => 0xc4,
                .absolute => 0xcc,
                else => null,
            },
            .bit => switch (opr) {
                .zeropage => 0x24,
                .absolute => 0x2c,
                else => null,
            },
            .clc => switch (opr) {
                .implied => 0x18,
                else => null,
            },
            .sec => switch (opr) {
                .implied => 0x38,
                else => null,
            },
            .cli => switch (opr) {
                .implied => 0x58,
                else => null,
            },
            .sei => switch (opr) {
                .implied => 0x78,
                else => null,
            },
            .clv => switch (opr) {
                .implied => 0xb8,
                else => null,
            },
            .cld => switch (opr) {
                .implied => 0xd8,
                else => null,
            },
            .sed => switch (opr) {
                .implied => 0xf8,
                else => null,
            },

            .bpl => switch (opr) {
                .immediate => 0x10,
                else => null,
            },
            .bmi => switch (opr) {
                .immediate => 0x30,
                else => null,
            },
            .bvc => switch (opr) {
                .immediate => 0x50,
                else => null,
            },
            .bvs => switch (opr) {
                .immediate => 0x70,
                else => null,
            },
            .bcc => switch (opr) {
                .immediate => 0x90,
                else => null,
            },
            .bcs => switch (opr) {
                .immediate => 0xb0,
                else => null,
            },
            .bne => switch (opr) {
                .immediate => 0xd0,
                else => null,
            },
            .beq => switch (opr) {
                .immediate => 0xf0,
                else => null,
            },
            .nop => switch (opr) {
                .implied => 0xea,
                else => null,
            },
            .pha => switch (opr) {
                .implied => 0x48,
                else => null,
            },
            .pla => switch (opr) {
                .implied => 0x68,
                else => null,
            },
            .php => switch (opr) {
                .implied => 0x08,
                else => null,
            },
            .plp => switch (opr) {
                .implied => 0x28,
                else => null,
            },
            .jmp => switch (opr) {
                .absolute => 0x4c,
                .indirect => 0x6c,
                else => null,
            },
            .jsr => switch (opr) {
                .absolute => 0x20,
                else => null,
            },
            .rts => switch (opr) {
                .implied => 0x60,
                else => null,
            },
        };
    }
};

const whitespace = " \t";

pub const Operand = union(enum) {
    implied,
    zeropage: u8,
    zeropage_x: u8,
    zeropage_y: u8,
    absolute: u16,
    absolute_x: u16,
    absolute_y: u16,
    immediate: u8,
    indirect: u16,
    indexed_indirect: u8,
    indirect_indexed: u8,

    fn parse(s: []const u8) !Operand {
        const maybe_comma = std.mem.indexOfScalar(u8, s, ',');

        switch (s[0]) {
            '#' => return .{ .immediate = try parseInt(u8, s[1..]) },
            '<' => {
                if (maybe_comma) |comma| {
                    if (std.mem.indexOfAny(
                        u8,
                        s[comma + 1 ..],
                        "Xx",
                    )) |_| return .{
                        .zeropage_x = try parseInt(u8, std.mem.trim(
                            u8,
                            s[1..comma],
                            whitespace,
                        )),
                    };
                    if (std.mem.indexOfAny(
                        u8,
                        s[comma + 1 ..],
                        "Yy",
                    )) |_| return .{
                        .zeropage_y = try parseInt(u8, std.mem.trim(
                            u8,
                            s[1..comma],
                            whitespace,
                        )),
                    };
                }

                return .{ .zeropage = try parseInt(u8, s[1..]) };
            },
            '(' => {
                const r_paren = std.mem.lastIndexOfScalar(
                    u8,
                    s,
                    ')',
                ) orelse return error.InvalidSyntax;

                if (maybe_comma) |comma| {
                    // (zp,X)
                    if (comma < r_paren) {
                        if (std.mem.indexOfAny(
                            u8,
                            s[comma + 1 .. r_paren],
                            "Xx",
                        ) == null) return error.InvalidSyntax;

                        return .{
                            .indexed_indirect = try parseInt(u8, std.mem.trim(
                                u8,
                                s[1..comma],
                                whitespace,
                            )),
                        };
                    }

                    // (zp),Y
                    if (std.mem.indexOfAny(
                        u8,
                        s[comma + 1 ..],
                        "Yy",
                    ) == null) return error.InvalidSyntax;

                    return .{
                        .indirect_indexed = try parseInt(u8, std.mem.trim(
                            u8,
                            s[1..comma],
                            whitespace,
                        )),
                    };
                }

                // (ind)
                return .{
                    .indirect = try parseInt(u16, std.mem.trim(
                        u8,
                        s[1..r_paren],
                        whitespace,
                    )),
                };
            },
            else => {
                if (maybe_comma) |comma| {
                    if (std.mem.indexOfAny(
                        u8,
                        s[comma + 1 ..],
                        "Xx",
                    )) |_| return .{
                        .absolute_x = try parseInt(u16, std.mem.trim(
                            u8,
                            s[0..comma],
                            whitespace,
                        )),
                    };
                    if (std.mem.indexOfAny(
                        u8,
                        s[comma + 1 ..],
                        "Yy",
                    )) |_| return .{
                        .absolute_y = try parseInt(u16, std.mem.trim(
                            u8,
                            s[0..comma],
                            whitespace,
                        )),
                    };
                    return error.InvalidSyntax;
                }

                return .{ .absolute = try parseInt(u16, s) };
            },
        }
    }
};

fn parseInt(comptime T: type, s: []const u8) !T {
    const base: u8, const slice = switch (s[0]) {
        '$' => .{ 16, s[1..] },
        '%' => .{ 2, s[1..] },
        else => .{ 10, s },
    };

    var v: T = 0;
    for (slice) |d| {
        const digit = try std.fmt.charToDigit(d, base);
        v = try std.math.mul(T, v, base);
        v = try std.math.add(T, v, digit);
    }
    return v;
}

test "basic ops" {
    try std.testing.expectEqualSlices(
        u8,
        &.{ 0xa9, 0x42, 0xae, 0x00, 0x69 },
        assembleComptime(
            \\LDA #$42
            \\LDX $6900
        ),
    );
}