const std = @import("std");
const Allocator = std.mem.Allocator;
const root = @import("root");
const Token = root.Lexer.Token;
const parse = root.parse;
const Location = root.Lexer.Location;

pub const VReg = enum(u32) { _ };

pub const Instr = struct {
    loc: Location,
    type: Type,

    pub const Type = union(enum) {
        constant: Constant,
        bin_op: BinOp,
        print: Print,
    };

    pub const Constant = struct {
        dest: VReg,
        value: u64,

        pub fn sources(_: Constant) Sources {
            return Sources.init(0) catch unreachable;
        }
    };

    pub const BinOp = struct {
        dest: VReg,
        lhs: VReg,
        rhs: VReg,
        op: Op,

        const Op = enum {
            add,
            sub,
        };

        pub fn sources(self: BinOp) Sources {
            return Sources.fromSlice(&.{ self.lhs, self.rhs }) catch unreachable;
        }
    };

    pub const Print = struct {
        arg: VReg,

        pub fn sources(self: Print) Sources {
            return Sources.fromSlice(&.{self.arg}) catch unreachable;
        }
    };

    pub fn sources(self: Instr) Sources {
        return switch (self.type) {
            inline else => |instr| instr.sources(),
        };
    }

    pub fn dest(self: *const Instr) ?VReg {
        return switch (self.type) {
            inline .constant, .bin_op => |s| s.dest,
            inline .print => null,
        };
    }

    pub const Sources = std.BoundedArray(VReg, 2);
};

pub const Block = struct {
    // arguments: []Reg,
    instrs: []Instr,
    vreg_last_use: std.AutoHashMap(VReg, usize),

    fn init(allocator: Allocator, instrs: []Instr) !Block {
        var vreg_last_use: std.AutoHashMap(VReg, usize) = .init(allocator);
        for (0.., instrs) |i, instr| {
            for (instr.sources().slice()) |src|
                try vreg_last_use.put(src, i);
        }
        return .{
            .instrs = instrs,
            .vreg_last_use = vreg_last_use,
        };
    }
};

pub fn compile(allocator: Allocator, source: []const u8, stmts: []parse.Stmt) !Block {
    const instrs: std.ArrayListUnmanaged(Instr) = try .initCapacity(allocator, 0);
    var ctx: CompileContext = .{
        .allocator = allocator,
        .source = source,
        .register_counter = 0,
        .instrs = instrs,
    };
    for (stmts) |stmt| {
        try ctx.compileStmt(stmt);
    }
    return .init(allocator, ctx.instrs.items);
}

const CompileContext = struct {
    allocator: Allocator,
    source: []const u8,
    register_counter: u32,
    instrs: std.ArrayListUnmanaged(Instr),

    const Self = @This();

    fn addInstr(self: *Self, instr: Instr) !void {
        try self.instrs.append(self.allocator, instr);
    }

    fn compileStmt(self: *Self, stmt: parse.Stmt) !void {
        switch (stmt.type) {
            .expr => |expr| _ = try self.compileExpr(expr),
        }
    }

    fn compileExpr(self: *Self, expr: *const parse.Expr) !VReg {
        const dest = self.register();
        switch (expr.type) {
            .integer_literal => try addInstr(self, .{
                .loc = expr.loc,
                .type = .{ .constant = .{ .dest = dest, .value = expr.loc.getInt(self.source) } },
            }),
            .bin_op => |binop| {
                const lhs = try self.compileExpr(binop.lhs);
                const rhs = try self.compileExpr(binop.rhs);
                try addInstr(self, .{
                    .loc = expr.loc,
                    .type = .{
                        .bin_op = .{
                            .dest = dest,
                            .lhs = lhs,
                            .rhs = rhs,
                            .op = switch (binop.op) {
                                .plus => .add,
                                .minus => .sub,
                            },
                        },
                    },
                });
            },
            .call => |call| {
                if (call.proc.type == .identifier and
                    std.mem.eql(u8, call.proc.loc.getIdent(self.source), "print"))
                {
                    const arg = try self.compileExpr(call.arg);
                    try self.addInstr(.{
                        .loc = expr.loc,
                        .type = .{ .print = .{ .arg = arg } },
                    });
                    // BUG: we're returning a bogus virtual register that we didn't write to
                } else {
                    return error.CantCallAnythingButPrint;
                }
            },
            .identifier => return error.CantCompileIdentifierExpr,
        }
        return dest;
    }

    fn register(self: *Self) VReg {
        const reg: VReg = @enumFromInt(self.register_counter);
        self.register_counter += 1;
        return reg;
    }
};