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 BlockRef = enum(u32) { _ };

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

    pub const Type = union(enum) {
        constant: Constant,
        bin_op: BinOp,
        proc_call: ProcCall,
        branch: Branch,
        jump: Jump,
        exit: Exit,
    };

    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 const Branch = struct {
        cond: VReg,
        true: BlockRef,
        false: BlockRef,

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

    pub const Jump = struct {
        to: BlockRef,

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

    pub const Exit = struct {
        pub fn sources(_: Exit) Sources {
            return Sources.init(0) 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,
            .branch, .jump, .exit => null,
        };
    }

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

    pub fn format(self: Instr, comptime fmt: []const u8, options: std.fmt.FormatOptions, writer: anytype) !void {
        _ = .{ fmt, options };

        switch (self.type) {
            .constant => |constant| try writer.print(
                "%{} = {}",
                .{ @intFromEnum(constant.dest), constant.value },
            ),
            .bin_op => |bin_op| try writer.print(
                "%{} = %{} {s} %{}",
                .{
                    @intFromEnum(bin_op.dest),
                    @intFromEnum(bin_op.lhs),
                    @tagName(bin_op.op),
                    @intFromEnum(bin_op.rhs),
                },
            ),
            .proc_call => |proc_call| try writer.print(
                "%{} = {s} %{}",
                .{
                    @intFromEnum(proc_call.dest),
                    @tagName(proc_call.proc),
                    @intFromEnum(proc_call.arg),
                },
            ),
            .branch => |branch| try writer.print(
                "branch %{} ? ${} : ${}",
                .{
                    @intFromEnum(branch.cond),
                    @intFromEnum(branch.true),
                    @intFromEnum(branch.false),
                },
            ),
            .jump => |jump| try writer.print("jump ${}", .{@intFromEnum(jump.to)}),
            .exit => |_| try writer.print("exit", .{}),
        }
    }
};

pub const Procedure = struct {
    blocks: []BasicBlock,

    fn init(allocator: Allocator, blocks: []BasicBlock) !Procedure {
        for (blocks) |*block| {
            try block.finalize(allocator);
        }
        return .{ .blocks = blocks };
    }
};

pub const BasicBlock = struct {
    // arguments: []Reg,
    instrs: std.ArrayListUnmanaged(Instr) = .empty,

    vreg_last_use: std.AutoHashMapUnmanaged(VReg, usize) = .empty,

    fn finalize(self: *BasicBlock, allocator: Allocator) !void {
        self.vreg_last_use = .empty;
        for (0.., self.instrs.items) |i, instr| {
            for (instr.sources().slice()) |src|
                try self.vreg_last_use.put(allocator, src, i);
            if (instr.dest()) |dest|
                try self.vreg_last_use.put(allocator, dest, i);
        }
    }
};

pub fn compile(allocator: Allocator, source: []const u8, block: parse.Block) !Procedure {
    var ctx: CompileContext = .{
        .allocator = allocator,
        .source = source,
        .register_counter = 0,
        .scope = .{ .locals = .empty, .parent = null },
        .blocks = .empty,
        .current_block = @enumFromInt(0),
    };
    try ctx.compileBlock(block);
    try ctx.addInstr(.{
        .loc = .{ .start = 0, .end = 0 },
        .type = .{ .exit = .{} },
    });
    return try .init(allocator, try ctx.blocks.toOwnedSlice(allocator));
}

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

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

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

    fn compileBlock(self: *CompileContext, block: parse.Block) !void {
        _ = try self.switchToNewBlock();

        var 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: *CompileContext, stmt: parse.Stmt) CompileError!void {
        switch (stmt.type) {
            .expr => |expr| _ = try self.compileExpr(expr),
            .block => |block| {
                const curr = self.current_block;
                const after = try self.switchToNewBlock();
                try self.compileBlock(block);
                const b = self.current_block;
                try self.addInstr(.{
                    .loc = stmt.loc,
                    .type = .{ .jump = .{ .to = after } },
                });
                self.current_block = curr;
                try self.addInstr(.{
                    .loc = stmt.loc,
                    .type = .{ .jump = .{ .to = b } },
                });
                self.current_block = after;
            },
            .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: *CompileContext, 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 self.addInstr(.{
                .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 self.addInstr(.{
                    .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;
            },
            .@"if" => |@"if"| {
                const cond = try self.compileExpr(@"if".cond);
                const curr = self.current_block;

                const after = try self.switchToNewBlock();

                try self.compileBlock(@"if".then);
                try self.addInstr(.{
                    .loc = expr.loc,
                    .type = .{ .jump = .{ .to = after } },
                });
                const t = self.current_block;

                const f = if (@"if".@"else") |@"else"| blk: {
                    try self.compileBlock(@"else");
                    try self.addInstr(.{
                        .loc = expr.loc,
                        .type = .{ .jump = .{ .to = after } },
                    });
                    break :blk self.current_block;
                } else after;

                self.current_block = curr;
                try self.addInstr(.{
                    .loc = expr.loc,
                    .type = .{ .branch = .{
                        .cond = cond,
                        .true = t,
                        .false = f,
                    } },
                });

                self.current_block = after;
            },
        }
        return dest;
    }

    fn switchToNewBlock(self: *CompileContext) !BlockRef {
        const ref: BlockRef = @enumFromInt(self.blocks.items.len);
        try self.blocks.append(self.allocator, .{});
        self.current_block = ref;
        return ref;
    }

    fn addInstr(self: *CompileContext, instr: Instr) !void {
        try self.blocks.items[@intFromEnum(self.current_block)]
            .instrs.append(self.allocator, instr);
    }

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