const std = @import("std");
const parser = @import("parser.zig");
const ast = @import("ast.zig");
const errors = @import("error.zig");
const lir = @import("ir.zig");
const Allocator = std.mem.Allocator;
const testing = std.testing;
const Type = @import("Value.zig").Type;

//! Contains the Compiler internals.
//! This compiles all AST nodes generated by the parser and turned into `ByteCode` instructions
//! Those instructions are than used by the vm to be executed.
//! Potentially, we could provide a different back-end to support
//! different use cases.

/// Compiles the source code into Luf's bytecode
pub fn compile(
    allocator: *Allocator,
    source: []const u8,
    _errors: *errors.Errors,
) (parser.Parser.Error || Compiler.Error)!lir.CompileUnit {
    const tree = try parser.parse(allocator, source, _errors);
    defer tree.deinit();

    var root_scope = Compiler.Scope{
        .symbols = Compiler.SymbolTable{},
        .id = .global,
        .allocator = allocator,
    };
    defer {
        for (root_scope.symbols.items()) |entry| {
            allocator.destroy(entry.value);
        }
        root_scope.symbols.deinit(allocator);
    }
    var arena = std.heap.ArenaAllocator.init(allocator);
    errdefer arena.deinit();

    var module = lir.Module{ .gpa = &arena.allocator };

    var compiler = Compiler{
        .instructions = lir.Instructions{},
        .allocator = allocator,
        .scope = &root_scope,
        .errors = _errors,
        .ir = &module,
        .modules = std.StringHashMapUnmanaged(*Compiler.Module){},
        .module_states = std.ArrayListUnmanaged(std.heap.ArenaAllocator.State){},
    };
    defer compiler.deinit();

    // declare all functions first
    try compiler.forwardDeclareNodes(tree.nodes);

    for (tree.nodes) |node| {
        if (node == .declaration and canBeForwardDeclared(node.declaration.value)) continue; // already compiled above
        try compiler.instructions.append(allocator, try compiler.resolveInst(node));
    }

    return lir.CompileUnit{
        .gpa = allocator,
        .state = arena.state,
        .instructions = compiler.instructions.toOwnedSlice(allocator),
    };
}

/// Returns true if the given `Node` can be forward declared
fn canBeForwardDeclared(node: ast.Node) bool {
    return switch (node) {
        .int_lit,
        .string_lit,
        .func_lit,
        .data_structure,
        .nil,
        .slice,
        => true,
        else => false,
    };
}

/// Contains the memory and generated instructions
/// after calling Compile()
pub const Compiler = struct {
    /// The instruction set the compiler is building
    instructions: lir.Instructions,
    /// Allocator used to allocate to the heap, call deinit() on compiler to free all memory
    allocator: *Allocator,
    /// Any errors generated by the compiler will be appended to `errors`
    errors: *errors.Errors,
    /// Imported modules and their global symboltable
    modules: std.StringHashMapUnmanaged(*Module),
    /// Arena states which will be freed once we finished compilation
    /// This ensures all Nodes that are imported from other files have a long enough
    /// lifetime to be used for typechecking imported functions/variables
    module_states: std.ArrayListUnmanaged(std.heap.ArenaAllocator.State),
    /// Globals counter, this is increased by 1 each time a global is defined
    /// This is increased by both global and module scope
    gc: u16 = 0,
    /// Current scope of the compilation, i.e. .loop
    scope: *Scope,
    /// Luf's IR, used to generate instructions
    ir: *lir.Module,

    /// Compiler errorset
    pub const Error = error{ CompilerError, OutOfMemory };

    /// Hashmap of `Symbol` where the key is the Symbol's name
    const SymbolTable = std.StringArrayHashMapUnmanaged(*Symbol);

    /// Symbols known by the compiler
    const Symbol = struct {
        name: []const u8,
        /// Not mutable by default
        mutable: bool = false,
        /// Not public by default
        is_pub: bool = false,
        /// Index of symbol table
        index: u16 = 0,
        /// The scope the symbol belongs to
        scope: Scope.Id,
        /// The `Ast.Node` that declared the symbol, used to retrieve the symbol's type
        node: ast.Node,
        /// Set to `true` if the symbol was created during forward declaration
        forward_declared: bool,
        /// IR declaration
        ident: *lir.Inst,
    };

    /// Module with its global symbols
    const Module = struct {
        /// symbols existing inside the module
        symbols: SymbolTable,
        /// used to determine if module still requires codegen
        compiled: bool,
    };

    /// Scope of the current state (function, global, etc)
    const Scope = struct {
        /// Symbols that exist within the current scope
        symbols: SymbolTable,
        /// Type of the Scope. i.e. Root, function, etc
        id: Tag,
        /// If not a global scope, the scope will own a parent
        parent: ?*Scope = null,
        /// used to create child scopes
        allocator: *Allocator,
        /// The offset used to calculate symbol indices
        offset: u16 = 0,

        /// The type of the scope
        const Id = enum {
            global,
            function,
            loop,
            module,
        };

        const Tag = union(Id) {
            global,
            loop,
            module: void,
            function: *ast.Node.FunctionLiteral,
        };

        /// Creates a new `Scope` from the current Scope.
        /// The new Scope will have its parent set to the current Scope
        fn fork(self: *Scope, id: Scope.Id) !*Scope {
            const new_scope = try self.allocator.create(Scope);
            new_scope.* = .{
                .symbols = SymbolTable{},
                .id = switch (id) {
                    .global => .{ .global = {} },
                    .module => .{ .module = {} },
                    .function => .{ .function = undefined },
                    .loop => .{ .loop = {} },
                },
                .offset = if (self.id == .loop) @truncate(u16, self.symbols.items().len) else 0,
                .parent = self,
                .allocator = self.allocator,
            };
            return new_scope;
        }

        /// Defines a new symbol and saves it in the symbol table
        /// Returns an error if Symbol already exists
        fn define(
            self: *Scope,
            name: []const u8,
            mutable: bool,
            node: ast.Node,
            forward_declared: bool,
            is_pub: bool,
            index: ?u16,
        ) !?*Symbol {
            if (self.symbols.contains(name)) return null;

            const new_index = index orelse @intCast(u16, self.symbols.items().len);

            const symbol = try self.allocator.create(Symbol);
            symbol.* = .{
                .name = name,
                .mutable = mutable,
                .index = new_index + self.offset,
                .scope = self.id,
                .node = node,
                .forward_declared = forward_declared,
                .ident = undefined,
                .is_pub = is_pub,
            };

            try self.symbols.putNoClobber(self.allocator, name, symbol);
            return symbol;
        }

        /// Retrieves a `Symbol` from the Scopes symbol table, returns null if not found
        fn resolve(self: *const Scope, name: []const u8) ?*Symbol {
            return self.symbols.get(name);
        }

        /// Cleans up the memory of the `Scope`
        fn deinit(self: *Scope) void {
            for (self.symbols.items()) |entry| {
                self.allocator.destroy(entry.value);
            }
            self.symbols.deinit(self.allocator);
            self.allocator.destroy(self);
        }
    };

    /// Frees up any memory allocated by the compiler
    fn deinit(self: *Compiler) void {
        var it = self.modules.iterator();
        while (it.next()) |mod| {
            for (mod.value.symbols.items()) |entry| self.allocator.destroy(entry.value);
            mod.value.symbols.deinit(self.allocator);
            self.allocator.destroy(mod.value);
        }
        self.modules.deinit(self.allocator);
        self.exitScope();
        self.instructions.deinit(self.allocator);
        for (self.module_states.items) |state| {
            state.promote(self.allocator).deinit();
        }
        self.module_states.deinit(self.allocator);
        self.* = undefined;
    }

    /// Returns `Error.CompilerError` and appends an error message to the `errors` list.
    fn fail(self: *Compiler, comptime msg: []const u8, index: usize, args: anytype) Error {
        try self.errors.add(msg, index, .err, args);
        return Error.CompilerError;
    }

    /// Sets the current `Scope` to its parent's scope and cleans up the closing scope's memory
    fn exitScope(self: *Compiler) void {
        if (self.scope.id == .global) return; // can't escape the global scope
        if (self.scope.parent) |parent| {
            const old = self.scope;
            self.scope = parent;
            old.deinit();
        }
    }

    /// Creates a new Scope with the given Id, then sets the new scope as the current
    fn createScope(self: *Compiler, id: Scope.Id) !void {
        self.scope = try self.scope.fork(id);
    }

    /// Recursively looks for the scope with the given Tag
    fn findScope(self: *Compiler, scope: *const Scope, tag: Scope.Id) ?*const Scope {
        if (scope.id == tag) return scope;
        return if (scope.parent) |parent| self.findScope(parent, tag) else null;
    }

    /// Attempts to resolve a symbol from the symbol table
    /// If not found, will attempt to resolve it from a parent scope
    fn resolveSymbol(self: *Compiler, scope: *const Scope, name: []const u8) ?*Symbol {
        return scope.resolve(name) orelse if (scope.id != .global and scope.parent != null)
            self.resolveSymbol(scope.parent.?, name)
        else
            null;
    }

    /// Defines a new symbol in the current scope
    fn defineSymbol(
        self: *Compiler,
        name: []const u8,
        mutable: bool,
        node: ast.Node,
        forward_declared: bool,
        is_pub: bool,
    ) !?*Symbol {
        const index: ?u16 = if (self.scope.id == .global or self.scope.id == .module) self.gc else null;
        if (try self.scope.define(name, mutable, node, forward_declared, is_pub, index)) |symbol| {
            if (index != null) self.gc += 1;
            return symbol;
        } else return null;
    }

    /// Returns a Luf `Type` based on the given node
    fn resolveType(self: *Compiler, node: ast.Node) !Type {
        if (node.getType()) |t| return t;

        return switch (node) {
            .identifier => |id| {
                const symbol = self.resolveSymbol(self.scope, id.value) orelse
                    return self.fail("'{s}' is undefined", id.token.start, .{id.value});

                return self.resolveType(symbol.node);
            },
            .infix => |inf| return self.resolveType(inf.left),
            .declaration => |decl| {
                if (decl.type_def) |td| return td.getType().?;

                return self.resolveType(decl.value);
            },
            .index => |index| return self.resolveType(index.left),
            .call_expression => |cal| return self.resolveType(cal.function),
            .slice => |slice| return self.resolveType(slice.left),
            else => std.debug.panic("TODO, implement type resolving for type: {s}\n", .{node}),
        };
    }

    /// Resolves the inner type of a node. i.e. this will return `Type.integer` for a list []int
    fn resolveScalarType(self: *Compiler, node: ast.Node) Error!Type {
        if (node.getInnerType()) |i| return i;
        return switch (node) {
            .identifier => |id| {
                const symbol = self.resolveSymbol(
                    self.scope,
                    id.value,
                ) orelse return self.fail("'{s}' is undefined", id.token.start, .{id.value});

                return self.resolveScalarType(symbol.node);
            },
            .infix => |inf| return self.resolveScalarType(inf.left),
            .declaration => |decl| {
                if (decl.type_def) |td| return td.getInnerType().?;

                return self.resolveScalarType(decl.value);
            },
            .index => |index| {
                const temp_type = try self.resolveScalarType(index.left);
                if (temp_type == .module) {
                    const function_name = index.index.string_lit.value;
                    const symbol = self.resolveSymbol(self.scope, index.left.identifier.value) orelse
                        return self.fail(
                        "'{s}' is undefined",
                        index.left.tokenPos(),
                        .{index.left.identifier.value},
                    );

                    const mod: *Module = self.modules.get(symbol.node.declaration.value.import.value.string_lit.value) orelse
                        return Type.function; // Module imported by VM, therefore just return it as a function

                    const function_symbol: *Symbol = mod.symbols.get(function_name) orelse
                        return self.fail("Module does not contain function '{s}'", index.index.tokenPos(), .{function_name});

                    return self.resolveScalarType(function_symbol.node);
                }

                return self.resolveScalarType(index.left);
            },
            .call_expression => |cal| return self.resolveScalarType(cal.function),
            .slice => |slice| return self.resolveScalarType(slice.left),
            else => std.debug.panic("TODO, implement scalar type resolving for type: {s}\n", .{node.getType()}),
        };
    }

    /// Retrieves all symbols for Nodes of type declaration
    fn forwardDeclareNodes(self: *Compiler, nodes: []const ast.Node) !void {
        for (nodes) |node| {
            if (node != .declaration) continue;
            const decl = node.declaration;

            if (decl.value == .import) {
                try self.forwardDeclareImport(decl.value.import);
                continue;
            } else if (!canBeForwardDeclared(decl.value)) continue;

            const symbol = (try self.defineSymbol(
                decl.name.identifier.value,
                decl.mutable,
                node,
                true,
                decl.is_pub,
            )) orelse return self.fail(
                "'{s}' has already been declared",
                node.tokenPos(),
                .{decl.name.identifier.value},
            );

            symbol.ident = try self.ir.emitIdent(
                try self.resolveType(node),
                node.tokenPos(),
                symbol.name,
                .global,
                symbol.index,
            );
        }
        for (self.scope.symbols.items()) |entry| {
            const symbol: *Symbol = entry.value;

            const value = try self.resolveInst(symbol.node.declaration.value);
            const decl = try self.ir.emitDecl(
                symbol.node.tokenPos(),
                symbol.name,
                symbol.index,
                .global,
                symbol.is_pub,
                symbol.mutable,
                value,
            );

            try self.instructions.append(self.allocator, decl);
        }
    }

    /// Forward declares all functions from a module, so we can do type checking for imported
    /// functions.
    fn forwardDeclareImport(self: *Compiler, node: *const ast.Node.Import) !void {
        const file_name = node.value.string_lit.value;
        if (self.modules.contains(file_name)) return;
        if (!std.mem.endsWith(u8, file_name, ".luf")) return;

        const file = std.fs.cwd().openFile(file_name, .{}) catch
            return self.fail("Cannot open file with path '{s}'", node.value.tokenPos(), .{file_name});
        defer file.close();

        const size = file.getEndPos() catch return self.fail("Unable to find file size", node.value.tokenPos(), .{});
        const source = file.reader().readAllAlloc(self.allocator, size) catch return self.fail(
            "Failed to read content of file with path '{s}'",
            node.value.tokenPos(),
            .{file_name},
        );
        defer self.allocator.free(source);

        const tree = try parser.parse(self.allocator, source, self.errors);
        errdefer tree.deinit();

        const module = try self.allocator.create(Module);
        errdefer self.allocator.destroy(module);

        module.* = .{ .symbols = SymbolTable{}, .compiled = false };
        errdefer module.symbols.deinit(self.allocator);

        for (tree.nodes) |n| {
            if (n != .declaration or n.declaration.value != .func_lit) continue;
            if (!n.declaration.is_pub) continue; // only public declarations

            const decl = n.declaration;
            const name = decl.name.identifier.value;

            if (module.symbols.contains(name))
                return self.fail("Identifier '{s}' already exists", n.declaration.name.tokenPos(), .{name});

            const symbol = try self.allocator.create(Symbol);
            symbol.* = .{
                .name = name,
                .mutable = decl.mutable,
                .scope = .global,
                .node = n,
                .forward_declared = true,
                .index = self.gc,
                .ident = undefined,
                .is_pub = decl.is_pub,
            };

            try module.symbols.putNoClobber(self.allocator, name, symbol);

            self.gc += 1;
        }

        try self.modules.putNoClobber(self.allocator, file_name, module);
        try self.module_states.append(self.allocator, tree.arena);
    }

    /// Depending on an `ast.Node` type, generates a `lir.Inst` instruction recursively
    fn resolveInst(self: *Compiler, node: ast.Node) (Error || lir.Module.Error || parser.Parser.Error)!*lir.Inst {
        return try switch (node) {
            .expression => |exp| self.compileExpr(exp),
            .block_statement => |block| self.compileBlock(block),
            .infix => |infix| self.compileInfix(infix),
            .prefix => |prefix| self.compilePrefix(prefix),
            .boolean => |boolean| self.compileBool(boolean),
            .int_lit => |int| self.ir.emitInt(node.tokenPos(), int.value),
            .if_expression => |if_exp| self.compileCond(if_exp),
            .declaration => |decl| self.compileDecl(decl, node),
            .identifier => |ident| self.compileIdent(ident),
            .string_lit => |string| self.ir.emitString(.string, node.tokenPos(), string.value),
            .data_structure => |ds| self.compileData(ds),
            .map_pair => |pair| self.compilePair(pair),
            .index => |index| self.compileIndex(index),
            .func_lit => |func| self.compileFunc(func),
            .call_expression => |call| self.compileCall(call),
            .@"return" => |ret| self.compileRet(ret),
            .while_loop => |loop| self.compileWhile(loop),
            .for_loop => |loop| self.compileFor(loop),
            .assignment => |asgn| self.compileAssign(asgn),
            .nil => self.ir.emitPrimitive(node.tokenPos(), .nil),
            .import => |imp| self.compileImport(imp),
            .@"break" => self.ir.emitNoOp(node.tokenPos(), .@"break"),
            .@"continue" => self.ir.emitNoOp(node.tokenPos(), .@"continue"),
            .range => |range| self.compileRange(range),
            .@"enum" => |enm| self.compileEnum(enm),
            .switch_statement => |sw| self.compileSwitch(sw),
            .switch_prong => |prong| self.comileProng(prong),
            .comment => |cmt| self.ir.emitString(.comment, node.tokenPos(), cmt.value),
            .type_def => |td| self.compileTypeDef(td),
            .func_arg => |arg| self.compileFuncArg(arg),
            .slice => |slice| self.compileSlice(slice),
        };
    }

    /// Compiles an `ast.Node.Expression` into a `lir.Inst.Single` this can be used
    /// by backends to know when an expression ends. For example, to pop a value.
    fn compileExpr(self: *Compiler, exp: *ast.Node.Expression) !*lir.Inst {
        const rhs = try self.resolveInst(exp.value);
        return self.ir.emitSingle(exp.token.start, .expr, rhs);
    }

    /// Compiles a `ast.Node.BlockStatement` node into a list of instructions
    fn compileBlock(self: *Compiler, block: *ast.Node.BlockStatement) !*lir.Inst {
        const list = try self.ir.gpa.alloc(*lir.Inst, block.nodes.len);

        for (block.nodes) |node, i| list[i] = try self.resolveInst(node);

        return self.ir.emitBlock(block.token.start, list);
    }

    /// Compiles an `ast.Node.Infix` node into a `lir.Inst.Double`
    fn compileInfix(self: *Compiler, infix: *ast.Node.Infix) !*lir.Inst {
        var lhs = try self.resolveType(infix.left);
        var rhs = try self.resolveType(infix.right);

        if (lhs == .list or lhs == .range or lhs == .function) lhs = try self.resolveScalarType(infix.left);
        if (rhs == .list or rhs == .range or rhs == .function) rhs = try self.resolveScalarType(infix.right);

        if (lhs != rhs)
            return self.fail("Expected type {s} but found type {s}", infix.right.tokenPos(), .{ lhs, rhs });

        const lhs_inst = try self.resolveInst(infix.left);
        const rhs_inst = try self.resolveInst(infix.right);

        return self.ir.emitDouble(infix.token.start, switch (infix.operator) {
            .add => .add,
            .multiply => .mul,
            .sub => .sub,
            .divide => .div,
            .less_than => .lt,
            .greater_than => .gt,
            .equal => .eql,
            .not_equal => .nql,
            .mod => .mod,
            .@"and" => .@"and",
            .@"or" => .@"or",
            .bitwise_xor => .bitwise_xor,
            .bitwise_or => .bitwise_or,
            .bitwise_and => .bitwise_and,
            .shift_left => .shift_left,
            .shift_right => .shift_right,
            .assign_add => .assign_add,
            .assign_sub => .assign_sub,
            .assign_mul => .assign_mul,
            .assign_div => .assign_div,
            .assign => .assign,
        }, lhs_inst, rhs_inst);
    }

    /// Compiles an `ast.Node.Prefix` node into a `lir.Inst.Single`
    fn compilePrefix(self: *Compiler, prefix: *ast.Node.Prefix) !*lir.Inst {
        return self.ir.emitSingle(
            prefix.token.start,
            switch (prefix.operator) {
                .minus => .negate,
                .bang => .not,
                .bitwise_not => .bitwise_not,
            },
            try self.resolveInst(prefix.right),
        );
    }

    /// Compiles an `ast.Node.Prefix` node into a `lir.Inst.Primitive`
    fn compileBool(self: *Compiler, boolean: *ast.Node.Bool) !*lir.Inst {
        return self.ir.emitPrimitive(boolean.token.start, if (boolean.value) .@"true" else .@"false");
    }

    /// Compiles an `ast.Node.IfExpression` node into a `lir.Inst.If`
    fn compileCond(self: *Compiler, node: *ast.Node.IfExpression) !*lir.Inst {
        const cond = try self.resolveInst(node.condition);
        const then_block = try self.resolveInst(node.true_pong);
        const else_block = if (node.false_pong) |f| try self.resolveInst(f) else null;

        return self.ir.emitCond(node.token.start, cond, then_block, else_block);
    }

    /// Compiles an `ast.Node.Declaration` node into a `lir.Inst.Decl`
    /// This will create a new symbol in the symbol table for the current scope
    fn compileDecl(self: *Compiler, decl: *ast.Node.Declaration, node: ast.Node) !*lir.Inst {
        if (decl.type_def) |td| {
            const explicit_type = try self.resolveType(td);

            const rhs_type = blk: {
                const tmp = try self.resolveType(decl.value);

                // if function, get the return type of the function
                if (tmp == .function or tmp == .module) break :blk try self.resolveScalarType(decl.value);

                break :blk tmp;
            };

            if (rhs_type == .nil and explicit_type != .optional)
                return self.fail("Cannot assign `nil` value to non-optional", decl.value.tokenPos(), .{})
            else if (explicit_type != .optional and explicit_type != rhs_type) {
                return self.fail("Expected type {s} but found type {s}", decl.value.tokenPos(), .{ explicit_type, rhs_type });
            }
        }

        const symbol = if (self.resolveSymbol(self.scope, decl.name.identifier.value)) |s| blk: {
            if (s.forward_declared and (self.scope.id == .global or self.scope.id == .module)) {
                s.forward_declared = false;
                break :blk s;
            } else {
                return self.fail("Identifier '{s}' has already been declared", decl.token.start, .{decl.name.identifier.value});
            }
        } else blk: {
            const s = (try self.defineSymbol(decl.name.identifier.value, decl.mutable, node, false, decl.is_pub)).?;
            s.ident = try self.ir.emitIdent(
                try self.resolveScalarType(decl.value),
                node.tokenPos(),
                s.name,
                switch (s.scope) {
                    .global, .module => .global,
                    else => .local,
                },
                s.index,
            );
            break :blk s;
        };

        const decl_value = try self.resolveInst(decl.value);

        return try self.ir.emitDecl(
            decl.token.start,
            symbol.name,
            symbol.index,
            switch (symbol.scope) {
                .global, .module => .global,
                else => .local,
            },
            symbol.is_pub,
            symbol.mutable,
            decl_value,
        );
    }

    /// Compiles an `ast.Node.Identifier` node into a `lir.Inst.Ident`
    fn compileIdent(self: *Compiler, id: *ast.Node.Identifier) !*lir.Inst {
        const symbol = self.resolveSymbol(self.scope, id.value) orelse
            return self.fail("Identifier '{s}' is undefined", id.token.start, .{id.value});

        return symbol.ident;
    }

    /// Compiles an `ast.Node.DataStructure` node into a `lir.Inst.DataStructure`
    /// Also compiles all elements
    fn compileData(self: *Compiler, ds: *ast.Node.DataStructure) !*lir.Inst {
        const inner_type = try self.resolveScalarType(ds.type_def_key);

        const elements = try self.ir.gpa.alloc(*lir.Inst, ds.value.?.len);

        // ds.value will only be null if it's specified as a type rather than an expression
        for (ds.value.?) |element, i| {
            const el_type = try self.resolveScalarType(element);
            if (inner_type != el_type)
                return self.fail(
                    "Expected type {s} but found type {s} for element {d}",
                    element.tokenPos(),
                    .{ inner_type, el_type, i },
                );

            elements[i] = try self.resolveInst(element);
        }

        return self.ir.emitList(
            if (ds.d_type == .array) .list else .map,
            ds.token.start,
            elements,
        );
    }

    /// Compiles an `ast.Node.MapPair` node into a `lir.Inst.Double`
    fn compilePair(self: *Compiler, node: *ast.Node.MapPair) !*lir.Inst {
        const key = try self.resolveInst(node.key);
        const value = try self.resolveInst(node.value);
        return self.ir.emitDouble(node.token.start, .pair, key, value);
    }

    /// Compiles an `ast.Node.Index` node into a `lir.Inst.Double`
    fn compileIndex(self: *Compiler, index: *ast.Node.IndexExpression) !*lir.Inst {
        if ((try self.resolveType(index.left)) == .module and index.left == .identifier) {
            const module_symol = self.resolveSymbol(self.scope, index.left.identifier.value).?;
            const module_name = module_symol.node.declaration.value.import.value.string_lit.value;
            if (self.modules.get(module_name)) |module| {
                const func_name = index.index.string_lit.value;
                var func: *Symbol = module.symbols.get(func_name).?;

                return try self.ir.emitIdent(
                    try self.resolveType(func.node),
                    index.token.start,
                    func.name,
                    .global,
                    func.index,
                );
            }
        }

        const lhs = try self.resolveInst(index.left);
        const rhs = try self.resolveInst(index.index);
        return self.ir.emitDouble(index.token.start, .load, lhs, rhs);
    }

    /// Compiles an `ast.Node.FuntionLiteral` node into a `lir.Inst.Func`
    fn compileFunc(self: *Compiler, function: *ast.Node.FunctionLiteral) !*lir.Inst {
        try self.createScope(.function);
        self.scope.id.function = function;

        const args = try self.ir.gpa.alloc(*lir.Inst, function.params.len);

        for (function.params) |param, i| {
            args[i] = try self.resolveInst(param);
        }

        const body = try self.resolveInst(function.body orelse return self.fail(
            "Function body missing",
            function.token.start,
            .{},
        ));

        // get locals
        const locals = try self.ir.gpa.alloc(*lir.Inst, self.scope.symbols.items().len);
        for (self.scope.symbols.items()) |entry, i|
            locals[i] = entry.value.ident;

        self.exitScope();

        const ret_type = try self.resolveInst(function.ret_type);

        return self.ir.emitFunc(
            function.token.start,
            body,
            locals,
            args,
            ret_type,
        );
    }

    /// Compiles an `ast.Node.FunctionArgument` into a `lir.Inst.FuncArg`
    fn compileFuncArg(self: *Compiler, arg: *ast.Node.FunctionArgument) !*lir.Inst {
        const ty = try self.resolveType(arg.arg_type);
        const name = arg.value;

        // function arguments are not mutable by default
        const symbol = (try self.defineSymbol(name, false, arg.arg_type, false, false)) orelse
            return self.fail(
            "Identifier '{s}' has already been declared",
            arg.token.start,
            .{name},
        );

        symbol.ident = try self.ir.emitIdent(
            ty,
            arg.token.start,
            symbol.name,
            .local,
            symbol.index,
        );

        return self.ir.emitFuncArg(arg.token.start, ty, name);
    }

    /// Compiles a slice expression `ast.Node.SliceExpression` into a `lir.Inst.Triple`
    fn compileSlice(self: *Compiler, slice: *ast.Node.SliceExpression) !*lir.Inst {
        const ty = try self.resolveType(slice.left);
        if (ty != .list) return self.fail("Expected a list but instead found type '{s}'", slice.left.tokenPos(), .{ty});

        const left = try self.resolveInst(slice.left);
        const start = blk: {
            if (slice.start) |start| break :blk try self.resolveInst(start);

            break :blk try self.ir.emitInt(slice.token.start, 0);
        };

        const end = blk: {
            if (slice.end) |end| break :blk try self.resolveInst(end);

            break :blk try self.ir.emitPrimitive(slice.token.start, .nil);
        };

        return self.ir.emitSlice(slice.token.start, left, start, end);
    }

    /// Compiles an `ast.Node.CallExpression` node into a `lir.Inst.Call`
    fn compileCall(self: *Compiler, call: *ast.Node.CallExpression) !*lir.Inst {
        const initial_function_type = try self.resolveType(call.function);

        const args = try self.ir.gpa.alloc(*lir.Inst, call.arguments.len);

        // function is either builtin or defined on a type
        // for now this can only be checked at runtime as the compiler is unaware of builtins
        if (call.function == .index and initial_function_type != .module) {
            for (call.arguments) |arg, i| {
                args[i] = try self.resolveInst(arg);
            }
            const func = try self.resolveInst(call.function);
            return self.ir.emitCall(
                try self.resolveScalarType(call.function),
                call.token.start,
                func,
                args,
            );
        }
        //if it's an identifier, we first do a lookup to retrieve it's declaration node
        //then return the function declaration. else we return the function itself
        //in the case of an anonymous function
        const function_node = if (call.function == .identifier and initial_function_type != .module) blk: {
            const function = self.resolveSymbol(self.scope, call.function.identifier.value) orelse
                return self.fail(
                "Function '{s}' is undefined",
                call.function.tokenPos(),
                .{call.function.identifier.value},
            );
            break :blk function.node.declaration.value.func_lit;
        } else if (initial_function_type != .module)
            call.function.func_lit
        else blk: {
            if (call.function == .index and call.function.index.left == .identifier) {
                // we handle module functions here
                const function_name = call.function.index.index.string_lit.value;
                const symbol = self.resolveSymbol(self.scope, call.function.index.left.identifier.value) orelse
                    return self.fail(
                    "Identifier '{s}' does not exist",
                    call.function.index.left.tokenPos(),
                    .{call.function.index.left.identifier.value},
                );

                const module_name = symbol.node.declaration.value.import.value.string_lit.value;
                if (self.modules.get(module_name)) |mod| {
                    const decl_symbol: *Symbol = mod.symbols.get(function_name) orelse
                        return self.fail(
                        "Module does not contain function '{s}'",
                        call.function.tokenPos(),
                        .{function_name},
                    );

                    break :blk decl_symbol.node.declaration.value.func_lit;
                } else {
                    // check if it's a luf source file or possible library
                    if (std.mem.endsWith(u8, module_name, ".luf"))
                        return self.fail("Module '{s}' does not exist", call.token.start, .{module_name});
                }
            }

            // not a module, so expect a library
            for (call.arguments) |arg, i| {
                args[i] = try self.resolveInst(arg);
            }

            const func = try self.resolveInst(call.function);

            return self.ir.emitCall(
                try self.resolveScalarType(call.function),
                call.token.start,
                func,
                args,
            );
        };

        if (function_node.params.len != call.arguments.len)
            return self.fail("Expected {d} arguments, but found {d}", call.token.start, .{
                function_node.params.len,
                call.arguments.len,
            });

        for (call.arguments) |arg, i| {
            const arg_type = try self.resolveType(arg);
            const func_arg_type = try self.resolveType(function_node.params[i]);

            if (arg_type != func_arg_type)
                return self.fail(
                    "Expected type '{s}' but found type '{s}'",
                    arg.tokenPos(),
                    .{ func_arg_type, arg_type },
                );
            args[i] = try self.resolveInst(arg);
        }

        const func = try self.resolveInst(call.function);
        return self.ir.emitCall(
            try self.resolveScalarType(call.function),
            call.token.start,
            func,
            args,
        );
    }

    /// Compiles an `ast.Node.ReturnStatement` into a `lir.Inst.Single`
    fn compileRet(self: *Compiler, ret: *ast.Node.Return) !*lir.Inst {
        const scope = self.findScope(self.scope, .function) orelse
            return self.fail(
            "Unexpected return statement. Return can only be done while inside a function body",
            ret.token.start,
            .{},
        );

        const func_ret_type = scope.id.function.ret_type.getType();
        const ret_type = blk: {
            const tmp = try self.resolveType(ret.value);

            if (tmp == .function or tmp == .list) break :blk try self.resolveScalarType(ret.value);
            break :blk tmp;
        };

        if (func_ret_type != ret_type)
            return self.fail(
                "Expected return type '{s}' but found type '{s}'",
                ret.value.tokenPos(),
                .{ func_ret_type, ret_type },
            );

        const value = try self.resolveInst(ret.value);
        return self.ir.emitSingle(ret.token.start, .@"return", value);
    }

    /// Compiles an `ast.Node.WhileLoop` into an `lir.Inst.Double`
    fn compileWhile(self: *Compiler, loop: *ast.Node.WhileLoop) !*lir.Inst {
        try self.createScope(.loop);

        const cond = try self.resolveInst(loop.condition);
        const body = try self.resolveInst(loop.block);

        self.exitScope();
        return self.ir.emitDouble(loop.token.start, .@"while", cond, body);
    }

    /// Compiles an `ast.Node.ForLoop` into an `lir.Inst.Loop`
    fn compileFor(self: *Compiler, loop: *ast.Node.ForLoop) !*lir.Inst {
        try self.createScope(.loop);
        const iterator_type = try self.resolveType(loop.iter);

        if (iterator_type != .list and iterator_type != .range and iterator_type != .string) {
            return self.fail(
                "Expected a list, range or string, but found type '{s}'",
                loop.iter.tokenPos(),
                .{iterator_type},
            );
        }
        const it = try self.resolveInst(loop.iter);

        // as above, parser ensures it's an identifier
        const index = if (loop.index) |i| blk: {
            // loop is just a boolean, so create a new Node that represents the integer for the counter
            const index_node = &ast.Node.ForLoop.index_node;
            index_node.token = i.identifier.token;

            const symbol = (try self.defineSymbol(
                i.identifier.value,
                false, // not mutable
                ast.Node{ .int_lit = index_node },
                false, // not forward declared
                false, // not public
            )) orelse return self.fail(
                "Identifier '{s}' has already been declared",
                loop.token.start,
                .{i.identifier.value},
            );

            // generate an identifier to attach to the symbol
            symbol.ident = try self.ir.emitIdent(.integer, i.tokenPos(), symbol.name, .local, symbol.index);
            break :blk symbol.ident;
        } else null;

        // parser already parses it as an identifier, no need to check here again
        const capture = (try self.defineSymbol(
            loop.capture.identifier.value,
            false, // not mutable
            loop.iter,
            false, // not forward declared
            false, // not public
        )) orelse return self.fail(
            "Capture identifier '{s}' has already been declared",
            loop.token.start,
            .{loop.capture.identifier.value},
        );

        capture.ident = try self.ir.emitIdent(
            try self.resolveScalarType(capture.node),
            loop.capture.tokenPos(),
            capture.name,
            .local,
            capture.index,
        );

        const body = try self.resolveInst(loop.block);

        const loop_declaration = self.ir.emitFor(loop.token.start, it, body, capture.ident, index);

        self.exitScope();

        return loop_declaration;
    }

    /// Compiles an `ast.Node.Assignment` node into a `lir.Inst.Assign`
    fn compileAssign(self: *Compiler, asg: *ast.Node.Assignment) !*lir.Inst {
        if (asg.left == .identifier) {
            const symbol = self.resolveSymbol(
                self.scope,
                asg.left.identifier.value,
            ) orelse return self.fail(
                "Identifier '{s}' is undefined",
                asg.token.start,
                .{asg.left.identifier.value},
            );

            if (!symbol.mutable) return self.fail("Identifier '{s}' is constant", asg.token.start, .{symbol.name});

            const right_type = try self.resolveType(asg.right);
            const left_type = try self.resolveType(symbol.node);

            if (right_type == .nil and left_type != .optional) {
                return self.fail("Cannot assign `nil` to non-optional", asg.token.start, .{});
            } else if (left_type == .optional) {
                const child_type = try self.resolveScalarType(symbol.node);
                if (child_type != right_type and right_type != .nil) {
                    return self.fail("Assignment expected type '{s}' but found type '{s}'", asg.token.start, .{
                        child_type,
                        right_type,
                    });
                }
            } else if (left_type != right_type)
                return self.fail("Assignment expected type '{s}' but found type '{s}'", asg.token.start, .{
                    left_type,
                    right_type,
                });

            const rhs = try self.resolveInst(asg.right);

            return self.ir.emitDouble(asg.token.start, .assign, symbol.ident, rhs);
        } else if (asg.left == .index) {
            const index = asg.left.index;

            const lhs_type = try self.resolveScalarType(index.left);
            const rhs_type = try self.resolveScalarType(asg.right);

            if (lhs_type != rhs_type)
                return self.fail("Expected type '{s}' but found type '{s}'", asg.right.tokenPos(), .{
                    lhs_type,
                    rhs_type,
                });

            const lhs = try self.resolveInst(index.left);
            const mid = try self.resolveInst(index.index);
            const rhs = try self.resolveInst(asg.right);

            return self.ir.emitTriple(asg.token.start, .store, lhs, mid, rhs);
        } else {
            return self.fail("Expected an identifier on left hand side", asg.left.tokenPos(), .{});
        }
    }

    /// Compile an `ast.Node.Import` node into a `lir.Inst.String` while also
    /// compiling the modules
    fn compileImport(self: *Compiler, imp: *ast.Node.Import) !*lir.Inst {
        // no need for resolveType, this is faster
        if (imp.value != .string_lit) return self.fail("Expected a string literal", imp.value.tokenPos(), .{});
        const file_name = imp.value.string_lit.value;
        if (std.mem.endsWith(u8, file_name, ".luf")) {
            // forward declaration would have defined it, so we can be sure it exists
            var mod: *Module = self.modules.get(file_name).?;
            if (mod.compiled) {
                return self.ir.emitString(.import, imp.token.start, file_name);
            }

            try self.createScope(.module);

            const file = std.fs.cwd().openFile(file_name, .{}) catch
                return self.fail("Could not open file at path '{s}'", imp.token.start, .{file_name});
            defer file.close();
            const size = file.getEndPos() catch return self.fail("Could not determine file size", imp.token.start, .{});
            const source = file.reader().readAllAlloc(self.allocator, size) catch
                return self.fail("Could not read file", imp.token.start, .{});
            defer self.allocator.free(source);
            const tree = try parser.parse(self.allocator, source, self.errors);
            defer tree.deinit();

            for (mod.symbols.items()) |entry| {
                const symbol: *Symbol = entry.value;
                const new_symbol = try self.scope.define(
                    symbol.name,
                    symbol.mutable,
                    symbol.node,
                    true,
                    symbol.is_pub,
                    symbol.index,
                );
                new_symbol.?.ident = try self.ir.emitIdent(
                    try self.resolveType(symbol.node),
                    symbol.node.tokenPos(),
                    symbol.name,
                    .global,
                    symbol.index,
                );
            }

            // we can't do it in the loop above because all symbols need to be defined first so we
            // can call functions from other functions
            for (self.scope.symbols.items()) |*entry| {
                // var symbol: *Symbol = entry.value;
                // const value = try self.resolveInst(symbol.node.declaration.value);
                // symbol.decl = try self.ir.emitDecl(symbol.node.tokenPos(), symbol.name, symbol.index, switch (symbol.scope) {
                //     .global, .module => .global,
                //     else => .local,
                // }, false, symbol.mutable, value);
            }

            for (tree.nodes) |n| {
                try self.instructions.append(self.allocator, try self.resolveInst(n));
            }
            mod.compiled = true;
            self.exitScope();
        }

        return self.ir.emitString(.import, imp.token.start, file_name);
    }

    /// Compiles an `ast.Node.Range` node into a `lir.Inst.Double`
    fn compileRange(self: *Compiler, range: *ast.Node.Range) !*lir.Inst {
        const lhs_type = try self.resolveScalarType(range.left);
        const rhs_type = try self.resolveScalarType(range.right);

        if (lhs_type != .integer) return self.fail("Expected an integer but found type '{s}'", range.left.tokenPos(), .{lhs_type});
        if (rhs_type != .integer) return self.fail("Expected an integer but found type '{s}'", range.right.tokenPos(), .{rhs_type});

        const lhs = try self.resolveInst(range.left);
        const rhs = try self.resolveInst(range.right);
        return self.ir.emitDouble(range.token.start, .range, lhs, rhs);
    }

    /// Compiles an `ast.Node.Enum` node into a `lir.Inst.Enum`
    fn compileEnum(self: *Compiler, enm: *ast.Node.EnumLiteral) !*lir.Inst {
        const list = try self.ir.gpa.alloc(*lir.Inst, enm.nodes.len);

        for (enm.nodes) |n, i| {
            if (n != .identifier)
                return self.fail(
                    "Expected an identifier but found type '{s}' inside the Enum declaration",
                    enm.token.start,
                    .{try self.resolveType(n)},
                );

            list[i] = try self.ir.emitString(.string, n.tokenPos(), n.identifier.value);
        }

        return self.ir.emitEnum(enm.token.start, list);
    }

    /// Compiles an `ast.Node.SwitchLiteral` into a `lir.Inst.Switch`
    fn compileSwitch(self: *Compiler, sw: *ast.Node.SwitchLiteral) !*lir.Inst {
        const prongs = try self.ir.gpa.alloc(*lir.Inst, sw.prongs.len);

        const capture = try self.resolveInst(sw.capture);

        const capture_type = try self.resolveType(sw.capture);

        if (capture_type != .integer and capture_type != .string)
            return self.fail("Switches are only allowed for integers, enums and strings. Found type '{s}'", sw.capture.tokenPos(), .{capture_type});

        for (sw.prongs) |p, i| {
            prongs[i] = try self.resolveInst(p);
        }

        return self.ir.emitSwitch(sw.token.start, capture, prongs);
    }

    /// Compiles an `ast.Node.SwitchProng` into a `lir.Inst.Double`
    /// ensures each branch is of type integer, string or range.
    fn comileProng(self: *Compiler, prong: *ast.Node.SwitchProng) !*lir.Inst {
        const prong_type = try self.resolveType(prong.left);
        switch (prong_type) {
            .integer, .string, .range => {},
            else => return self.fail("Unpermitted type '{s}' in switch case", prong.left.tokenPos(), .{prong_type}),
        }

        const lhs = try self.resolveInst(prong.left);
        const rhs = try self.resolveInst(prong.right);

        return self.ir.emitDouble(prong.token.start, .branch, lhs, rhs);
    }

    /// Compiles an `ast.Node.TypeDef` into a `lir.Inst.NoOp`
    /// This is only used for the compiler to ensure type safety
    /// and generates a no op instruction that codegen can ignore
    fn compileTypeDef(self: *Compiler, node: *ast.Node.TypeDef) !*lir.Inst {
        return self.ir.emitType(node.getType().?, node.token.start);
    }
};

/// Tests if given input results in the expected instructions
fn testInput(input: []const u8, expected: []const lir.Inst.Tag) !void {
    const alloc = testing.allocator;

    var err = errors.Errors.init(alloc);
    defer err.deinit();

    var result = compile(alloc, input, &err) catch |_err| {
        try err.write(input, std.io.getStdOut().writer());
        return _err;
    };
    defer result.deinit();

    try testing.expectEqual(expected.len, result.instructions.len);
    for (result.instructions) |inst, i| {
        try testing.expectEqual(expected[i], inst.tag);
    }
}

test "Arithmetic" {
    const test_cases = .{
        .{
            .input = "1 + 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "3 - 1",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "1 * 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "2 / 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "1 + 2 * 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "5 * 2 / 5 - 8",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "1 < 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "2 > 1",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "1 == 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "1 != 2",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "true == false",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "-1",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
    };

    inline for (test_cases) |case| {
        try testInput(case.input, case.tags);
    }
}

test "Conditional" {
    const test_cases = .{
        .{
            .input = "if (true) { 5 } 10",
            .tags = &[_]lir.Inst.Tag{ .expr, .expr },
        },
        .{
            .input = "if (true) { 5 } else { 7 } 10",
            .tags = &[_]lir.Inst.Tag{ .expr, .expr },
        },
    };

    inline for (test_cases) |case| {
        try testInput(case.input, case.tags);
    }
}

test "Declaration" {
    const input = "const x = \"foo\"";

    const alloc = testing.allocator;
    var err = errors.Errors.init(alloc);
    defer err.deinit();

    var result = try compile(alloc, input, &err);
    defer result.deinit();

    var inst = result.instructions[0];

    try testing.expectEqual(lir.Inst.Tag.decl, inst.tag);

    const decl = inst.as(lir.Inst.Decl);
    try testing.expectEqualStrings("x", decl.name);
    try testing.expectEqual(lir.Inst.Tag.string, decl.value.tag);
    try testing.expectEqualStrings("foo", decl.value.as(lir.Inst.String).value);
}

test "Lists" {
    const test_cases = .{
        .{
            .input = "const x = []int{1, 2, 3}",
            .tags = &[_]lir.Inst.Tag{.decl},
        },
        .{
            .input = "const x = []int{1, 2, 3}[0]",
            .tags = &[_]lir.Inst.Tag{.decl},
        },
        .{
            .input = "const x = []int:int{1: 2, 2: 1, 5: 6}",
            .tags = &[_]lir.Inst.Tag{.decl},
        },
        .{
            .input = "const list = []int{1, 2, 3} list[1] = 10 list[1]",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr, .expr },
        },
    };

    inline for (test_cases) |case| {
        try testInput(case.input, case.tags);
    }
}

test "Functions" {
    const test_cases = .{
        .{
            .input = "fn() void { 1 + 2 }",
            .tags = &[_]lir.Inst.Tag{.expr},
        },
        .{
            .input = "const x = fn() void { 1 } x()",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr },
        },
        .{
            .input = "const func = fn(x: int) int { return x } func(5)",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr },
        },
    };

    inline for (test_cases) |case| {
        try testInput(case.input, case.tags);
    }
}

test "Loop" {
    const input = "while true { 10 }";

    const alloc = testing.allocator;
    var err = errors.Errors.init(alloc);
    defer err.deinit();

    var result = try compile(alloc, input, &err);
    defer result.deinit();

    var inst = result.instructions[0];

    try testing.expectEqual(lir.Inst.Tag.@"while", inst.tag);

    const loop = inst.as(lir.Inst.Double);
    try testing.expectEqual(lir.Inst.Tag.primitive, loop.lhs.tag);
    try testing.expectEqual(lir.Inst.Primitive.PrimType.@"true", loop.lhs.as(lir.Inst.Primitive).prim_type);
    try testing.expectEqual(
        @as(u64, 10),
        loop.rhs.as(lir.Inst.Block).instructions[0].as(lir.Inst.Single).rhs.as(lir.Inst.Int).value,
    );
}

test "Assign" {
    const test_cases = .{
        .{
            .input = "mut x = 5 x = 6",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr },
        },
        .{
            .input = "mut x = 5 x += 6 x",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr, .expr },
        },
    };

    inline for (test_cases) |case| {
        try testInput(case.input, case.tags);
    }
}

test "Optionals" {
    const test_cases = .{
        .{
            .input = "mut x: ?int = nil",
            .tags = &[_]lir.Inst.Tag{.decl},
        },
        .{
            .input = "mut x: ?int = 5",
            .tags = &[_]lir.Inst.Tag{.decl},
        },
        .{
            .input = "mut x: ?int = nil x = 5",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr },
        },
        .{
            .input = "mut x: ?int = 5 x = nil",
            .tags = &[_]lir.Inst.Tag{ .decl, .expr },
        },
    };

    inline for (test_cases) |case| try testInput(case.input, case.tags);
}

test "Slices" {
    const test_case =
        \\ const list = []int{1, 2, 3}
        \\ const slice1 = list[1:2]
        \\ const slice2 = list[:1]
        \\ const slice3 = list[0:]
    ;

    const tags = &[_]lir.Inst.Tag{
        .decl,
        .decl,
        .decl,
        .decl,
    };

    try testInput(test_case, tags);
}
