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,
        proc_call: ProcCall,
    };

    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 ProcCall = struct {
        dest: VReg,
        arg: VReg,
        proc: Proc,

        const Proc = enum {
            print,
            read_int,
        };

        pub fn sources(self: ProcCall) 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, .proc_call => |s| s.dest,
        };
    }

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

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

    fn init(allocator: Allocator, instrs: []Instr) !BasicBlock {
        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);
            if (instr.dest()) |dest|
                try vreg_last_use.put(dest, i);
        }
        return .{
            .instrs = instrs,
            .vreg_last_use = vreg_last_use,
        };
    }
};

pub fn compile(allocator: Allocator, source: []const u8, block: parse.Block) !BasicBlock {
    const instrs: std.ArrayListUnmanaged(Instr) = try .initCapacity(allocator, 0);
    var ctx: CompileContext = .{
        .allocator = allocator,
        .source = source,
        .register_counter = 0,
        .scope = .{ .locals = .empty, .parent = null },
        .instrs = instrs,
    };
    try ctx.compileBlock(block);
    return .init(allocator, ctx.instrs.items);
}

const CompileError = error{
    OutOfMemory,
    CanOnlyCallIdentifiers,
    UnknownProcedure,
    UnknownVariable,
};

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

    const Scope = struct {
        locals: std.StringHashMapUnmanaged(VReg),
        parent: ?*Scope,
    };

    const Self = @This();

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

    fn compileBlock(self: *Self, block: parse.Block) !void {
        const parent = try self.allocator.create(Scope);
        defer self.allocator.destroy(parent);
        parent.* = self.scope;
        self.scope = .{
            .locals = .empty,
            .parent = parent,
        };
        for (block.stmts) |stmt| {
            try self.compileStmt(stmt);
        }
        self.scope.locals.deinit(self.allocator);
        self.scope = parent.*;
    }

    fn compileStmt(self: *Self, stmt: parse.Stmt) CompileError!void {
        switch (stmt.type) {
            .expr => |expr| _ = try self.compileExpr(expr),
            .block => |block| try self.compileBlock(block),
            .declare_var => |declare_var| {
                const val = try self.compileExpr(declare_var.value);
                const name = declare_var.ident.getIdent(self.source);
                try self.scope.locals.put(self.allocator, name, val);
            },
        }
    }

    fn compileExpr(self: *Self, expr: *const parse.Expr) !VReg {
        // This is not used by all expression types, but creating an unused virtual register is not a problem.
        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) return error.CanOnlyCallIdentifiers;
                const proc = std.meta.stringToEnum(Instr.ProcCall.Proc, call.proc.loc.getIdent(self.source)) orelse return error.UnknownProcedure;

                const arg = try self.compileExpr(call.arg);
                try self.addInstr(.{
                    .loc = expr.loc,
                    .type = .{ .proc_call = .{ .dest = dest, .arg = arg, .proc = proc } },
                });
            },
            .identifier => {
                var scope: ?*Scope = &self.scope;
                while (scope) |s| : (scope = s.parent) {
                    if (s.locals.get(expr.loc.getIdent(self.source))) |reg| {
                        return reg;
                    }
                }
                return error.UnknownVariable;
            },
        }
        return dest;
    }

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