const std = @import("std");
const mem = std.mem;
const debug = std.debug;
const log = std.log.scoped(.vm);
const testAllocator = std.testing.allocator;

const bytecode = @import("bytecode.zig");
const Canvas = @import("canvas.zig");
const system = @import("system.zig");
const value = @import("value.zig");
const Value = value.Value;

const Vm = @This();

stack: []Value = &.{},
stack_top: u32 = 0,
call_stack: []CallFrame = &.{},
call_stack_top: u32 = 0,
defs: []Value = &.{},
fuel: u32 = 0, // NOTE: VM must be refueled via reset() before running code
exception: ?Exception = null,
// closure, ip are for exception reporting. These variables are read-only, system functions
// cannot affect control flow using them.
closure: ?*const value.Closure = null,
ip: ?[*]const u8 = null,

pub const Limits = struct {
    stack_capacity: usize = 256,
    call_stack_capacity: usize = 256,
};

pub const CallFrame = struct {
    closure: *const value.Closure,
    return_addr: ?[*]const u8,
    bottom: u32,
};

pub const Exception = struct {
    len: usize,
    format: *const fn (buf: []u8, args: *align(64) const anyopaque) []u8,
    args: [40]u8 align(64), // increase the size if we ever throw a larger exception
};

/// All errors coming from inside the VM get turned into a single Exception type, which signals
/// any kind of error that prevented the machine from running any further.
/// Details about the exception can be found in the VM's `exception` field.
pub const Error = error{Exception};

// NOTE: A VM is only ever initialized. There is no function for deallocating its resources.
// The intent is to use it with a freshly-reset Scratch, and deallocate its resources through that.
pub fn init(a: mem.Allocator, defs: *const bytecode.Defs, limits: *const Limits) !Vm {
    return .{
        .stack = try a.alloc(Value, limits.stack_capacity),
        .call_stack = try a.alloc(CallFrame, limits.call_stack_capacity),
        .defs = try a.alloc(Value, defs.num_defs),
    };
}

pub fn reset(vm: *Vm, fuel: u32) void {
    vm.stack_top = 0;
    vm.call_stack_top = 0;
    vm.fuel = fuel;
    vm.exception = null;
}

// NOTE: When the VM throws an exception, it must be reset. There's no resuming from an exception.
pub fn throw(vm: *Vm, comptime fmt: []const u8, args: anytype) Error {
    log.info("throw: fmt={s}", .{fmt});
    log.debug("implicit stack frame: closure={?any} ip={?*}", .{ vm.closure, vm.ip });

    const Args = @TypeOf(args);
    const max_args_size = @sizeOf(@TypeOf(vm.exception.?.args));
    if (@sizeOf(Args) > max_args_size) {
        @compileError(std.fmt.comptimePrint(
            "format arguments are too large; size={}, max={}",
            .{ @sizeOf(Args), max_args_size },
        ));
    }

    const Formatter = struct {
        fn format(buf: []u8, erased_args: *align(64) const anyopaque) []u8 {
            return std.fmt.bufPrint(buf, fmt, @as(*const Args, @ptrCast(erased_args)).*) catch unreachable;
        }
    };

    var exn = Exception{
        .len = @truncate(std.fmt.count(fmt, args)),
        .format = Formatter.format,
        .args = undefined,
    };
    @memcpy(exn.args[0..@sizeOf(Args)], mem.asBytes(&args));
    vm.exception = exn;

    return error.Exception;
}

pub fn outOfMemory(vm: *Vm) Error {
    return vm.throw("out of memory", .{});
}

// NOTE: Different from CallFrame, this is a helper type for obtaining stack traces.
pub const StackFrame = struct {
    closure: *const value.Closure,
    ip: ?[*]const u8,
};

pub fn stackFrameCount(vm: *const Vm) usize {
    if (vm.closure != null) {
        return vm.call_stack_top + 1;
    } else {
        return vm.call_stack_top;
    }
}

pub fn stackFrame(vm: *const Vm, index: usize) StackFrame {
    if (vm.closure) |closure| {
        if (index == 0) {
            return .{
                .closure = closure,
                .ip = vm.ip,
            };
        } else {
            // NOTE: Minus two, because remember index == 0 is occupied by the current stack frame.
            const call_frame_index = vm.call_stack_top - 1 - (index - 1);
            const call_frame = vm.call_stack[call_frame_index];
            return .{
                .closure = call_frame.closure,
                .ip = call_frame.return_addr,
            };
        }
    } else {
        const call_frame = vm.call_stack[vm.call_stack_top - 1 - index];
        return .{
            .closure = call_frame.closure,
            .ip = call_frame.return_addr,
        };
    }
}

/// Debug assertion for bytecode validity.
/// In future versions, this may become disabled in release builds.
fn validateBytecode(vm: *Vm, ok: bool, comptime fmt: []const u8, args: anytype) Error!void {
    if (!ok) {
        return vm.throw("corrupted bytecode: " ++ fmt, args);
    }
}

pub fn consumeFuel(vm: *Vm, fuel: *u32, amount: u32) Error!void {
    const new_fuel, const overflow = @subWithOverflow(fuel.*, amount);
    if (overflow > 0) {
        return vm.throw("code ran for too long (out of fuel!)", .{});
    }
    fuel.* = new_fuel;
}

pub fn push(vm: *Vm, val: Value) Error!void {
    if (vm.stack_top >= vm.stack.len) {
        return vm.throw("too many live temporary values (local variables and expression operands)", .{});
    }
    log.debug("PUSH {any} <- {}", .{ vm.stack[0..vm.stack_top], val });
    vm.stack[vm.stack_top] = val;
    vm.stack_top += 1;
}

pub fn pop(vm: *Vm) Error!Value {
    try vm.validateBytecode(vm.stack_top > 0, "value stack underflow", .{});
    vm.stack_top -= 1;
    const result = vm.stack[vm.stack_top];
    log.debug("POP  {any} -> {}", .{ vm.stack[0..vm.stack_top], result });
    return vm.stack[vm.stack_top];
}

pub fn top(vm: *const Vm) Value {
    if (vm.stack_top > 0) {
        return vm.stack[vm.stack_top - 1];
    } else {
        return .nil;
    }
}

pub fn pushCall(vm: *Vm, frame: CallFrame) Error!void {
    if (vm.call_stack_top >= vm.call_stack.len) {
        return vm.throw("too much recursion", .{});
    }
    log.debug("PUSH CALL {}", .{frame});
    vm.call_stack[vm.call_stack_top] = frame;
    vm.call_stack_top += 1;
}

pub fn popCall(vm: *Vm) Error!CallFrame {
    try vm.validateBytecode(vm.call_stack_top > 0, "call stack underflow", .{});
    log.debug("POP CALL", .{});
    vm.call_stack_top -= 1;
    return vm.call_stack[vm.call_stack_top];
}

pub fn local(vm: *Vm, bottom: u32, offset: u8) Error!*Value {
    const index = bottom + offset;
    try vm.validateBytecode(index < vm.stack_top, "stack index out of bounds", .{});
    return &vm.stack[bottom + offset];
}

pub fn def(vm: *Vm, index: u16) Error!*Value {
    try vm.validateBytecode(index < vm.defs.len, "def index out of bounds", .{});
    return &vm.defs[index];
}

// Utility struct for storing and restoring the VM's state variables across FFI boundaries.
const Context = struct {
    ip: ?[*]const u8,
    closure: ?*const value.Closure,
    fuel: u32,
};

fn restoreContext(vm: *Vm) Context {
    return .{
        .ip = vm.ip,
        .closure = vm.closure,
        .fuel = vm.fuel,
    };
}

fn storeContext(vm: *Vm, context: Context) void {
    vm.ip = context.ip;
    vm.closure = context.closure;
    vm.fuel = context.fuel;
}

inline fn read(comptime T: type, ip: *[*]const u8) T {
    const result = mem.readInt(T, ip.*[0..@sizeOf(T)], std.builtin.Endian.little);
    ip.* = ip.*[@sizeOf(T)..];
    return result;
}

inline fn readOpcode(vm: *Vm, ip: *[*]const u8) bytecode.Opcode {
    vm.ip = ip.*;
    const opcode: bytecode.Opcode = @enumFromInt(read(u8, ip));
    log.debug("OP {*} {}", .{ ip.*, opcode });
    return opcode;
}

/// Before calling this, vm.stack_top should be saved and the appropriate amount of parameters must
/// be pushed onto the stack to execute the provided closure.
/// The saved vm.stack_top must be passed to init_bottom.
///
/// allocator should be a scratch buffer; there is no way to free the memory allocated by a VM.
pub fn run(
    vm: *Vm,
    allocator: mem.Allocator,
    init_closure: *const value.Closure,
    init_bottom: u32,
) Error!void {
    log.debug("BEGIN RUN {}", .{init_closure});

    // NOTE: This will need swapping out once we implement first-class system functions, because
    // system functions should be able to call themselves (like: map (range 1 10) sin)
    debug.assert(init_closure.impl == .bytecode);

    var closure = init_closure;
    var ip: [*]const u8 = closure.bytecodePtr(0);
    var bottom = init_bottom;
    var fuel = vm.fuel;

    vm.closure = closure;
    vm.ip = ip;

    for (0..closure.impl.bytecode.local_count) |_| {
        try vm.push(.nil);
    }

    try vm.pushCall(.{
        .closure = closure,
        .return_addr = null, // nowhere to return to
        .bottom = bottom,
    });

    next: switch (vm.readOpcode(&ip)) {
        .nil => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.nil);
            continue :next vm.readOpcode(&ip);
        },

        .false => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.false);
            continue :next vm.readOpcode(&ip);
        },

        .true => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.true);
            continue :next vm.readOpcode(&ip);
        },

        .tag => {
            try vm.consumeFuel(&fuel, 1);
            const tag_id: value.TagId = @enumFromInt(read(u16, &ip));
            try vm.push(.{ .tag = tag_id });
            continue :next vm.readOpcode(&ip);
        },

        .number => {
            try vm.consumeFuel(&fuel, 1);
            const number: f32 = @bitCast(read(u32, &ip));
            try vm.push(.{ .number = number });
            continue :next vm.readOpcode(&ip);
        },

        .rgba => {
            try vm.consumeFuel(&fuel, 1);
            const r = read(u8, &ip);
            const g = read(u8, &ip);
            const b = read(u8, &ip);
            const a = read(u8, &ip);
            try vm.push(.{ .rgba = value.rgbaFrom8(value.Rgba8{ r, g, b, a }) });
            continue :next vm.readOpcode(&ip);
        },

        .local => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            const l = try vm.local(bottom, index);
            try vm.push(l.*);
            continue :next vm.readOpcode(&ip);
        },

        .set_local => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            const new = try vm.pop();
            const l = try vm.local(bottom, index);
            l.* = new;
            continue :next vm.readOpcode(&ip);
        },

        .capture => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            const captures = closure.impl.bytecode.captures;
            try vm.validateBytecode(index < captures.len, "capture index out of bounds", .{});
            const capture = captures[index];
            try vm.push(capture);
            continue :next vm.readOpcode(&ip);
        },

        .def => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u16, &ip);
            const d = try vm.def(index);
            try vm.push(d.*);
            continue :next vm.readOpcode(&ip);
        },

        .set_def => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u16, &ip);
            const new = try vm.pop();
            const d = try vm.def(index);
            d.* = new;
            continue :next vm.readOpcode(&ip);
        },

        .list => {
            const len = read(u16, &ip);
            try vm.consumeFuel(&fuel, len);
            const list_end = vm.stack_top;
            vm.stack_top -= len;
            const elements = vm.stack[vm.stack_top..list_end];
            const list = allocator.dupe(Value, elements) catch return vm.outOfMemory();
            const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
            ref.* = .{ .list = list };
            try vm.push(.{ .ref = ref });
            continue :next vm.readOpcode(&ip);
        },

        .function => {
            const param_count = read(u8, &ip);
            const then = read(u16, &ip);
            const name_len = read(u8, &ip);
            const name = ip[0..name_len];
            ip += name_len;
            const body = ip;
            ip = closure.bytecodePtr(then);

            const local_count = read(u8, &ip);
            const capture_count = read(u8, &ip);

            try vm.consumeFuel(&fuel, 1 + capture_count);
            const captures = allocator.alloc(Value, capture_count) catch return vm.outOfMemory();
            for (captures) |*capture| {
                const capture_kind = read(u8, &ip);
                const index = read(u8, &ip);
                vm.ip = ip;
                capture.* = switch (capture_kind) {
                    bytecode.capture_local => (try vm.local(bottom, index)).*,
                    bytecode.capture_capture => blk: {
                        const closureCaptures = closure.impl.bytecode.captures;
                        try vm.validateBytecode(index < closureCaptures.len, "capture index out of bounds", .{});
                        break :blk closureCaptures[index];
                    },
                    else => .nil,
                };
            }

            const new_closure = value.Closure{
                .name = name,
                .param_count = param_count,
                .impl = .{ .bytecode = .{
                    .chunk = closure.impl.bytecode.chunk,
                    .start = @truncate(body - closure.bytecodePtr(0)),
                    .local_count = local_count,
                    .captures = captures,
                } },
            };
            const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
            ref.* = .{ .closure = new_closure };
            try vm.push(.{ .ref = ref });

            continue :next vm.readOpcode(&ip);
        },

        .jump => {
            const offset = read(u16, &ip);
            try vm.consumeFuel(&fuel, 1);
            ip = closure.bytecodePtr(offset);
            continue :next vm.readOpcode(&ip);
        },

        .jump_if_not => {
            const offset = read(u16, &ip);
            try vm.consumeFuel(&fuel, 1);
            const condition = try vm.pop();
            if (!condition.isTruthy()) {
                ip = closure.bytecodePtr(offset);
            }
            continue :next vm.readOpcode(&ip);
        },

        .field => {
            try vm.consumeFuel(&fuel, 1);

            const tag_count = read(u8, &ip);
            const field_tag = try vm.pop();
            if (field_tag != .tag) {
                return vm.throw("name of data field to look up must be a tag (starting with an uppercase letter)", .{});
            }
            const field_tag_id = field_tag.tag;

            var found_index: ?usize = null;
            for (0..tag_count) |i| {
                const tag_id: value.TagId = @enumFromInt(read(u16, &ip));
                if (tag_id == field_tag_id) {
                    found_index = i;
                }
            }

            if (found_index) |index| {
                const fields = closure.impl.bytecode.captures;
                try vm.validateBytecode(index < fields.len, "field index out of bounds", .{});
                const field = fields[index];
                try vm.push(field);
            } else {
                return vm.throw("field with this name does not exist", .{});
            }

            continue :next vm.readOpcode(&ip);
        },

        .call => {
            const arg_count = read(u8, &ip);

            try vm.consumeFuel(&fuel, 1);

            const function_value = try vm.pop();
            if (function_value != .ref or function_value.ref.* != .closure) {
                return vm.throw("attempt to call a value that is not a function", .{});
            }
            const called_closure = &function_value.ref.closure;

            // NOTE: Will need replacing for first-class system functions.
            debug.assert(called_closure.impl == .bytecode);

            if (arg_count != called_closure.param_count) {
                return vm.throw(
                    "function expects {} arguments, but it received {}",
                    .{ called_closure.param_count, arg_count },
                );
            }

            const call_frame = CallFrame{
                .closure = closure,
                .return_addr = ip,
                .bottom = bottom,
            };

            const new_bottom, const overflow = @subWithOverflow(vm.stack_top, arg_count);
            try vm.validateBytecode(overflow == 0, "not enough values on the stack for arguments", .{});

            closure = called_closure;
            ip = closure.bytecodePtr(0);
            bottom = new_bottom;

            // pushCall before setting vm.closure and vm.ip, such that stack overflow errors are
            // reported at the call site rather than within the called function.
            try vm.pushCall(call_frame);

            // Remember to update the closure used for exception reporting, since readOpcode does
            // not do that as it does with ip.
            // ip also needs updating, because the topmost stack frame will assume that vm.ip is
            // based in the new closure.
            // We want errors for no stack space being left for locals to appear at the start of the
            // function in source code, not earlier or later. Hence we do this before local slots
            // are pushed onto the stack.
            vm.closure = closure;
            vm.ip = ip;

            for (0..closure.impl.bytecode.local_count) |_| {
                try vm.push(.nil);
            }

            continue :next vm.readOpcode(&ip);
        },

        .system => {
            const index = read(u8, &ip);
            const arg_count = read(u8, &ip);
            const system_fn = &system.fns[index];

            try vm.consumeFuel(&fuel, 1);

            log.debug("system index={} arg_count={} system_fn={s}", .{ index, arg_count, system_fn.closure.name });

            // Push the current call frame onto the stack, such that the stack trace displays it
            // in addition to the ongoing system call.
            try vm.pushCall(.{
                .closure = closure,
                .return_addr = ip,
                .bottom = bottom,
            });

            // Also now, push the system function onto the stack, because nulling out vm.closure
            // and vm.ip removes the extra "current function" stack entry.
            try vm.pushCall(.{
                .closure = &system_fn.closure,
                .return_addr = null,
                .bottom = vm.stack_top - arg_count,
            });

            vm.storeContext(.{
                .fuel = fuel,
                .closure = null,
                .ip = null,
            });
            const result = try system_fn.impl(.{
                .vm = vm,
                .allocator = allocator,
                .args = vm.stack[vm.stack_top - arg_count .. vm.stack_top],
            });
            const context = vm.restoreContext();
            fuel = context.fuel;

            vm.stack_top -= arg_count;
            try vm.push(result);

            // Restore the exception handling ip and closure to what it was.
            log.debug("system restore", .{});
            vm.ip = ip;
            vm.closure = closure;

            // Remove the temporary call frames.
            _ = try vm.popCall(); // systemClosure
            _ = try vm.popCall(); // closure

            continue :next vm.readOpcode(&ip);
        },

        .ret => {
            try vm.consumeFuel(&fuel, 1);

            const result = try vm.pop();
            const call_frame = try vm.popCall();

            try vm.validateBytecode(
                bottom <= vm.stack_top,
                "called function popped too many values. bottom={} stack_top={}",
                .{ bottom, vm.stack_top },
            );

            vm.stack_top = bottom;
            try vm.push(result);

            if (call_frame.return_addr) |addr| {
                closure = call_frame.closure;
                ip = addr;
                bottom = call_frame.bottom;

                vm.closure = closure;
                vm.ip = ip;
            } else {
                // If there's nowhere to return to, stop interpreting.
                vm.storeContext(Context{
                    .closure = closure,
                    .ip = ip,
                    .fuel = fuel,
                });
                break :next;
            }

            continue :next vm.readOpcode(&ip);
        },

        _ => |invalid_opcode| {
            try vm.consumeFuel(&fuel, 1);
            // NOTE: Not a validateBytecode call because this is zero-cost on the happy path,
            // so we don't need to disable it on release builds.
            return vm.throw("corrupted bytecode: invalid opcode {}", .{invalid_opcode});
        },
    }

    log.debug("END RUN", .{});
}

pub const Dotter = struct {
    from: value.Vec4,
    to: value.Vec4,
    num: f32,
};

/// NOTE: Assumes the value at the top is a dotter reticle.
pub fn runDotter(
    vm: *Vm,
    allocator: mem.Allocator,
    from: value.Vec4,
    to: value.Vec4,
    num: f32,
) Error!void {
    const reticle = try vm.pop();
    const draw = reticle.ref.reticle.dotter.draw; // parameter count checked on construction

    const data = allocator.dupe(Value, &[_]Value{
        .{ .vec4 = from },
        .{ .vec4 = to },
        .{ .number = num },
    }) catch return vm.outOfMemory();
    const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
    ref.* = value.Ref{ .closure = .{
        .name = "dotter",
        .param_count = 1,
        .impl = .{ .bytecode = .{
            .chunk = &system.record_dotter,
            .start = 0,
            .local_count = 0,
            .captures = data,
        } },
    } };

    const bottom = vm.stack_top;
    try vm.push(.{ .ref = ref });
    try vm.run(allocator, draw, bottom);
}