const std = @import("std");
const chunk = @import("chunk.zig");
const value = @import("value.zig");
const values = value;
const compiler = @import("compiler.zig");
const objects = @import("object.zig");

const Chunk = chunk.Chunk;
const Value = value.Value;
const Compiler = compiler.Compiler;

pub const StdOut = *std.io.OutStream(std.fs.File.WriteError);

pub const InterpretResult = error{
    Ok,
    CompileError,
    RuntimeError,
};

fn isFalsey(val: value.Value) bool {
    return val.vtype == .Nil or (val.vtype == .Bool and !val.as.Bool);
}

fn valuesEqual(a: value.Value, b: value.Value) bool {
    if (a.vtype != b.vtype) return false;

    switch (a.vtype) {
        .Nil => return true,
        .Bool => return a.as.Bool == b.as.Bool,
        .Number => return a.as.Number == b.as.Number,
        .Object => blk: {
            var aStr = a.as.Object.value.String;
            var bStr = b.as.Object.value.String;
            return std.mem.compare(u8, aStr, bStr) == .Equal;
        },
    }
}

pub const ValueMap = std.StringHashMap(values.Value);

pub const VM = struct {
    chk: *Chunk = undefined,
    ip: usize = 0,

    stack: []Value,
    stackTop: usize = 0,

    stdout: StdOut,
    debug_flag: bool,
    allocator: *std.mem.Allocator,

    objs: ?*objects.Object = null,
    globals: ValueMap,

    fn resetStack(self: *VM) void {
        self.stackTop = 0;
    }

    pub fn init(
        allocator: *std.mem.Allocator,
        stdout: StdOut,
        debug_flag: bool,
    ) !VM {
        var self = VM{
            .stack = try allocator.alloc(Value, 256),
            .stdout = stdout,
            .debug_flag = debug_flag,
            .allocator = allocator,

            .globals = ValueMap.init(allocator),
        };

        self.resetStack();

        return self;
    }

    fn deinitObject(self: *VM, obj: *objects.Object) void {
        switch (obj.otype) {
            .String => blk: {
                self.allocator.free(obj.value.String);
                self.allocator.destroy(obj);
                break :blk;
            },
            else => unreachable,
        }
    }

    fn deinitObjects(self: *VM) void {
        var obj_opt: ?*objects.Object = self.objs;

        // doing a while(obj != null) but with optionals
        while (true) {
            if (obj_opt) |obj| {
                var next = obj.next;
                self.deinitObject(obj);
                obj_opt = next;
            } else {
                break;
            }
        }
    }

    pub fn deinit(self: *VM) void {
        self.globals.deinit();
        self.deinitObjects();
    }

    pub fn debug(self: *VM, comptime fmt: []const u8, args: ...) void {
        if (self.debug_flag) {
            std.debug.warn(fmt, args);
        }
    }

    fn readByte(self: *VM) u8 {
        var byte: u8 = self.chk.code[self.ip];
        self.ip += 1;
        return byte;
    }

    fn readConst(self: *VM) Value {
        return self.chk.constants.values[self.readByte()];
    }

    fn readConstLong(self: *VM) Value {
        const v3 = self.readByte();
        const v2 = self.readByte();
        const v1 = self.readByte();
        const const_idx = (@intCast(u24, v3) << 16) |
            (@intCast(u24, v2) << 8) |
            v1;

        return self.chk.constants.values[const_idx];
    }

    fn debugStack(self: *VM) !void {
        try self.stdout.print("          ");
        for (self.stack) |val, idx| {
            if (idx >= self.stackTop) break;

            try self.stdout.print("[ ");
            try value.printValue(self.stdout, val);
            try self.stdout.print(" ]");
        }
        try self.stdout.print("\n");
    }

    /// gets a f64 out of a value on the top of the stack.
    fn popNum(self: *VM) !f64 {
        var val: Value = self.pop();

        switch (val.vtype) {
            .Number => return val.as.Number,

            else => |vtype| blk: {
                self.runtimeError("Expected number, got {x}", vtype);
                return InterpretResult.RuntimeError;
            },
        }
    }

    fn concatenateStrings(self: *VM) !void {
        var b = self.pop().as.Object.value.String;
        var a = self.pop().as.Object.value.String;

        var res_str = try std.mem.join(
            self.allocator,
            "",
            [_][]u8{ a, b },
        );

        var val = values.ObjVal(try objects.takeString(self, res_str));
        try self.push(val);
    }

    fn doAdd(self: *VM) !void {
        if (values.isObjType(self.peek(0), .String) and
            values.isObjType(self.peek(1), .String))
        {
            return try self.concatenateStrings();
        }

        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.NumberVal(a + b));
    }

    fn doSub(self: *VM) !void {
        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.NumberVal(a - b));
    }

    fn doMul(self: *VM) !void {
        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.NumberVal(a * b));
    }

    fn doDiv(self: *VM) !void {
        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.NumberVal(a / b));
    }

    fn doGreater(self: *VM) !void {
        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.BoolVal(a > b));
    }

    fn doLess(self: *VM) !void {
        var b = try self.popNum();
        var a = try self.popNum();
        try self.push(values.BoolVal(a < b));
    }

    fn runtimeError(self: *VM, comptime fmt: []const u8, args: ...) void {
        std.debug.warn(fmt, args);
        std.debug.warn("\n[line {}] in script\n", self.chk.lines[self.ip]);
        self.resetStack();
    }

    fn defGlobal(self: *VM, name: []const u8) !void {
        _ = try self.globals.put(name, self.peek(0));
        _ = self.pop();
    }

    fn readString(self: *VM) []u8 {
        return self.readConst().as.Object.value.String;
    }

    fn readStringLong(self: *VM) []u8 {
        return self.readConstLong().as.Object.value.String;
    }

    fn doGetGlobal(self: *VM, name: []u8) !void {
        var kv_opt = self.globals.get(name);

        if (kv_opt) |kv| {
            try self.push(kv.value);
        } else {
            self.runtimeError("Undefined variable '{}'.", name);
            return InterpretResult.RuntimeError;
        }
    }

    fn doSetGlobal(self: *VM, name: []u8) !void {
        var res = try self.globals.getOrPut(name);

        if (res.found_existing) {
            res.kv.value = self.peek(0);
        } else {
            self.runtimeError("Undefined variable '{}'.", name);
            return InterpretResult.RuntimeError;
        }
    }

    fn run(self: *VM) !void {
        while (true) {
            if (self.debug_flag) {
                try self.debugStack();
                _ = try self.chk.disassembleInstruction(self.stdout, self.ip);
            }

            var instruction = self.readByte();

            switch (instruction) {
                chunk.OpCode.Constant => blk: {
                    var constant = self.readConst();
                    try self.push(constant);
                    break :blk;
                },
                chunk.OpCode.ConstantLong => blk: {
                    var constant = self.readConstLong();
                    try self.push(constant);
                    break :blk;
                },

                chunk.OpCode.Print => blk: {
                    try value.printValue(self.stdout, self.pop());
                    try self.stdout.print("\n");
                    break :blk;
                },

                chunk.OpCode.Return => blk: {
                    // Exit VM
                    return InterpretResult.Ok;
                },

                chunk.OpCode.Nil => try self.push(values.NilVal()),
                chunk.OpCode.True => try self.push(values.BoolVal(true)),
                chunk.OpCode.False => try self.push(values.BoolVal(false)),

                chunk.OpCode.Pop => blk: {
                    _ = self.pop();
                },

                chunk.OpCode.GetLocal => blk: {
                    var slot = self.readByte();
                    try self.push(self.stack[slot]);
                },
                chunk.OpCode.SetLocal => blk: {
                    var slot = self.readByte();
                    self.stack[slot] = self.peek(0);
                },

                chunk.OpCode.GetGlobal => blk: {
                    try self.doGetGlobal(self.readString());
                },
                chunk.OpCode.SetGlobal => blk: {
                    try self.doSetGlobal(self.readString());
                    break :blk;
                },

                // extracting the name is different depending of the
                // op code since one just uses a single byte, the other
                // uses three bytes since its a u24.
                chunk.OpCode.DefineGlobal => blk: {
                    try self.defGlobal(self.readString());
                    break :blk;
                },
                chunk.OpCode.DefineGlobalLong => blk: {
                    try self.defGlobal(self.readStringLong());
                    break :blk;
                },

                chunk.OpCode.Equal => blk: {
                    var a = self.pop();
                    var b = self.pop();
                    try self.push(values.BoolVal(valuesEqual(a, b)));
                },

                chunk.OpCode.Greater => try self.doGreater(),
                chunk.OpCode.Less => try self.doLess(),

                chunk.OpCode.Add => try self.doAdd(),
                chunk.OpCode.Subtract => try self.doSub(),
                chunk.OpCode.Multiply => try self.doMul(),
                chunk.OpCode.Divide => try self.doDiv(),
                chunk.OpCode.Not => blk: {
                    try self.push(values.BoolVal(isFalsey(self.pop())));
                },

                chunk.OpCode.Negate => blk: {
                    var val = self.peek(0);
                    if (val.vtype != .Bool) {
                        self.runtimeError("Operand must be a number.");
                        return InterpretResult.RuntimeError;
                    }

                    val = self.pop();
                    switch (val.as) {
                        .Number => |num| {
                            try self.push(values.NumberVal(-num));
                        },
                        else => unreachable,
                    }
                },

                else => blk: {
                    std.debug.warn("Unknown instruction: {x}\n", instruction);
                    return InterpretResult.RuntimeError;
                },
            }
        }
    }

    pub fn interpret(self: *VM, src: []const u8) !void {
        //self.ip = 0;
        //self.debug("VM start\n");
        //var res = try self.run();
        //self.debug("VM end\n");
        //return res;
        var chk = try Chunk.init(self.allocator);

        var cmpr = Compiler.init(
            self.allocator,
            &chk,
            self.stdout,
            src,
            self.debug_flag,
            self,
        );
        if (!try cmpr.compile(&chk)) {
            return InterpretResult.CompileError;
        }

        self.chk = &chk;
        self.ip = 0;
        return try self.run();
    }

    pub fn push(self: *VM, val: Value) !void {
        if (self.stackTop > 0 and self.stackTop - 1 > self.stack.len) {
            self.stack = try self.allocator.realloc(self.stack, self.stack.len + 1);
        }

        self.stack[self.stackTop] = val;
        self.stackTop += 1;
    }

    pub fn pop(self: *VM) Value {
        self.stackTop -= 1;
        return self.stack[self.stackTop];
    }

    pub fn peek(self: *VM, distance: usize) Value {
        return self.stack[self.stackTop - 1 - distance];
    }
};