//! Virtual machine state for story execution.
const std = @import("std");
const tokenizer = @import("tokenizer.zig");
const Ast = @import("Ast.zig");
const AstGen = @import("AstGen.zig");
pub const Object = @import("Story/object.zig").Object;
const Dumper = @import("Story/Dumper.zig");
const assert = std.debug.assert;
const Story = @This();

allocator: std.mem.Allocator,
dump_writer: ?*std.Io.Writer = null,
is_exited: bool = false,
can_advance: bool = false,
choice_index: usize = 0,
current_choices: std.ArrayListUnmanaged(Choice) = .empty,
globals: std.StringHashMapUnmanaged(?*Object) = .empty,
paths: std.ArrayListUnmanaged(*Object) = .empty,
stack: std.ArrayListUnmanaged(?*Object) = .empty,
call_stack: std.ArrayListUnmanaged(CallFrame) = .empty,
stack_max: usize = 128,
gc_objects: std.SinglyLinkedList = .{},

pub const CallFrame = struct {
    ip: usize,
    sp: usize,
    callee: *Object.ContentPath,
};

pub const Choice = struct {
    text: std.ArrayListUnmanaged(u8),
    dest_offset: u16,
};

pub const Opcode = enum(u8) {
    /// Exit the VM normally.
    exit,
    ret,
    /// Pop a value off the stack, discarding it.
    pop,
    /// Push an object representing the boolean value of "true" to the stack.
    true,
    /// Push an object representing the boolean value of "false" to the stack.
    false,
    /// Pop two values off the stack and calculate their sum.
    /// The result will be pushed to the stack.
    add,
    /// Pop two values off the stack and calculate their difference.
    /// The result will be pushed to the stack.
    sub,
    /// Pop two values off the stack and calculate their product.
    /// The result will be pushed to the stack.
    mul,
    /// Pop two values off the stack and calculate their quotient.
    /// The result will be pushed to the stack.
    div,
    mod,
    neg,
    not,
    cmp_eq,
    cmp_lt,
    cmp_gt,
    cmp_lte,
    cmp_gte,
    /// Jump unconditionally to the target address.
    jmp,
    /// Jump conditionally to the target address if the boolean value at the
    /// top of the stack is true.
    jmp_t,
    /// Jump conditionally to the target address if the boolean value at the
    /// top of the stack is false.
    jmp_f,
    call,
    divert,
    load_const,
    load,
    store,
    load_global,
    store_global,
    /// Pop a value off the stack and write it to the content stream.
    stream_push,
    stream_line,
    stream_glue,
    /// Flush the content stream to the story consumer.
    stream_flush,
    br_push,
    br_table,
    br_dispatch,
    br_select_index,
    _,
};

pub fn deinit(story: *Story) void {
    const gpa = story.allocator;
    var next = story.gc_objects.first;
    while (next) |node| {
        const object: *Object = @alignCast(@fieldParentPtr("node", node));
        next = node.next;
        object.destroy(story);
    }

    story.current_choices.deinit(gpa);
    story.globals.deinit(gpa);
    story.paths.deinit(gpa);
    story.stack.deinit(gpa);
    story.call_stack.deinit(gpa);
}

pub fn dump(story: *Story, writer: *std.Io.Writer) !void {
    const story_dumper: Dumper = .{ .story = story, .writer = writer };
    for (story.paths.items) |path_object| {
        try story_dumper.dump(@ptrCast(path_object));
    }
}

pub fn trace(story: *Story, writer: *std.Io.Writer, frame: *CallFrame) !void {
    try writer.print("\tStack => stack_pointer={d}, objects=[", .{frame.sp});

    const story_dumper: Dumper = .{ .story = story, .writer = writer };
    const stack = &story.stack;
    const stack_top = story.stack.items.len;
    if (stack_top > 0) {
        const last_slot = stack.items[stack.items.len - 1];
        for (stack.items[frame.sp .. stack.items.len - 1]) |slot| {
            if (slot) |object| {
                try story_dumper.dumpObject(object);
            } else {
                try writer.writeAll("null");
            }
            try writer.writeAll(", ");
        }
        if (last_slot) |object| {
            try story_dumper.dumpObject(object);
        } else {
            try writer.writeAll("null");
        }
    }

    try writer.writeAll("]\n");
    _ = try story_dumper.dumpInst(frame.callee, frame.ip, true);
    return writer.flush();
}

fn isCallStackEmpty(vm: *const Story) bool {
    return vm.call_stack.items.len == 0;
}

fn currentFrame(vm: *Story) *CallFrame {
    return &vm.call_stack.items[vm.call_stack.items.len - 1];
}

fn peekStack(vm: *Story, offset: usize) ?*Object {
    const stack_top = vm.stack.items.len;
    assert(stack_top > offset);
    assert(stack_top != 0);

    return vm.stack.items[stack_top - offset - 1];
}

fn pushStack(vm: *Story, object: *Object) !void {
    const gpa = vm.allocator;
    const stack_top = vm.stack.items.len;
    const max_stack_top = vm.stack_max;
    if (stack_top >= max_stack_top) return error.StackOverflow;

    return vm.stack.append(gpa, object);
}

fn popStack(vm: *Story) ?*Object {
    return vm.stack.pop() orelse unreachable;
}

fn getConstant(_: *Story, frame: *CallFrame, offset: u8) !*Object {
    const constant_pool = frame.callee.const_pool;
    if (offset >= constant_pool.len) return error.InvalidArgument;
    return constant_pool[offset];
}

fn getLocal(vm: *Story, frame: *CallFrame, offset: u8) ?*Object {
    const stack_top = vm.stack.capacity;
    const stack_offset = frame.sp + offset;
    assert(stack_top > stack_offset);

    return vm.stack.items[stack_offset];
}

fn setLocal(vm: *Story, frame: *CallFrame, offset: u8, value: *Object) void {
    const stack_top = vm.stack.capacity;
    const stack_offset = frame.sp + offset;
    assert(stack_top > stack_offset);

    vm.stack.items[stack_offset] = value;
}

fn getGlobal(vm: *Story, key: *const Object.String) !*Object {
    const key_bytes = key.bytes[0..key.length];
    const val = vm.globals.get(key_bytes) orelse return error.InvalidVariable;
    return val orelse unreachable;
}

fn setGlobal(vm: *Story, key: *const Object.String, value: *Object) !void {
    const key_bytes = key.bytes[0..key.length];
    return vm.globals.putAssumeCapacity(key_bytes, value);
}

fn execute(vm: *Story) !std.ArrayListUnmanaged(u8) {
    const gpa = vm.allocator;
    errdefer {
        vm.can_advance = false;
    }
    if (vm.isCallStackEmpty()) return .empty;

    const frame = vm.currentFrame();
    const code = std.mem.bytesAsSlice(Opcode, frame.callee.bytes);
    var stream_writer = std.Io.Writer.Allocating.init(gpa);
    defer stream_writer.deinit();

    while (true) {
        if (vm.dump_writer) |w| {
            vm.trace(w, frame) catch {};
        }
        switch (code[frame.ip]) {
            .exit => {
                vm.is_exited = true;
                vm.can_advance = false;
                return .empty;
            },
            .true => {
                const true_object = try Object.Number.create(vm, .{
                    .boolean = true,
                });
                try vm.pushStack(@ptrCast(true_object));
                frame.ip += 1;
            },
            .false => {
                const false_object = try Object.Number.create(vm, .{
                    .boolean = false,
                });
                try vm.pushStack(@ptrCast(false_object));
                frame.ip += 1;
            },
            .pop => {
                if (vm.popStack()) |_| {} else {
                    return error.InvalidArgument;
                }
                frame.ip += 1;
            },
            .add => {
                const lhs = vm.peekStack(1) orelse return error.Bugged;
                const rhs = vm.peekStack(0) orelse return error.Bugged;
                const value = try Object.add(vm, lhs, rhs);

                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(value);
                frame.ip += 1;
            },
            .sub, .mul, .div, .mod => |op| {
                const lhs = vm.peekStack(1) orelse return error.Bugged;
                const rhs = vm.peekStack(0) orelse return error.Bugged;
                const value = try Object.Number.performArithmetic(vm, op, @ptrCast(lhs), @ptrCast(rhs));

                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(@ptrCast(value));
                frame.ip += 1;
            },
            .neg => {
                if (vm.peekStack(0)) |arg| {
                    _ = Object.Number.negate(@ptrCast(arg));
                } else {
                    return error.InvalidArgument;
                }
                frame.ip += 1;
            },
            .not => {
                if (vm.peekStack(0)) |arg| {
                    const value = try Object.Number.create(vm, .{
                        .boolean = arg.isFalsey(),
                    });

                    _ = vm.popStack();
                    try vm.pushStack(@ptrCast(value));
                } else {
                    return error.StackOverflow;
                }
                frame.ip += 1;
            },
            .cmp_eq => {
                const lhs = vm.peekStack(1) orelse return error.Bugged;
                const rhs = vm.peekStack(0) orelse return error.Bugged;
                const value = try Object.cmpEql(vm, @ptrCast(lhs), @ptrCast(rhs));
                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(@ptrCast(value));

                frame.ip += 1;
            },
            .cmp_lt, .cmp_gt, .cmp_lte, .cmp_gte => |op| {
                const lhs = vm.peekStack(1) orelse return error.Bugged;
                const rhs = vm.peekStack(0) orelse return error.Bugged;
                const value = try Object.Number.performLogic(vm, op, @ptrCast(lhs), @ptrCast(rhs));
                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(@ptrCast(value));

                frame.ip += 1;
            },
            .jmp => {
                const arg_offset = readAddress(code, frame.ip);
                frame.ip += 3 + arg_offset;
            },
            .jmp_t => {
                const arg_offset = readAddress(code, frame.ip);
                frame.ip += 3;

                if (vm.peekStack(0)) |condition| {
                    if (!condition.isFalsey()) {
                        frame.ip += arg_offset;
                    }
                } else {
                    return error.InvalidArgument;
                }
            },
            .jmp_f => {
                const arg_offset = readAddress(code, frame.ip);
                frame.ip += 3;

                if (vm.peekStack(0)) |condition| {
                    if (condition.isFalsey()) {
                        frame.ip += arg_offset;
                    }
                } else {
                    return error.InvalidArgument;
                }
            },
            .load_const => {
                const index: u8 = @intFromEnum(code[frame.ip + 1]);
                const value = try vm.getConstant(frame, index);
                try vm.pushStack(value);
                frame.ip += 2;
            },
            .load => {
                const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                const value = vm.getLocal(frame, arg_offset) orelse return error.Bugged;
                try vm.pushStack(value);
                frame.ip += 2;
            },
            .store => {
                if (vm.peekStack(0)) |arg| {
                    const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                    vm.setLocal(frame, arg_offset, arg);
                } else {
                    return error.InvalidArgument;
                }
                frame.ip += 2;
            },
            .load_global => {
                const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                const global_name = try vm.getConstant(frame, arg_offset);
                assert(global_name.tag == .string);

                const global_value = try vm.getGlobal(@ptrCast(global_name));
                try vm.pushStack(global_value);
                frame.ip += 2;
            },
            .store_global => {
                const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                const global_name = try vm.getConstant(frame, arg_offset);
                assert(global_name.tag == .string);

                if (vm.peekStack(0)) |arg| {
                    try vm.setGlobal(@ptrCast(global_name), arg);
                    _ = vm.popStack();
                    try vm.pushStack(arg);
                } else {
                    return error.InvalidArgument;
                }
                frame.ip += 2;
            },
            .stream_push => {
                const arg_object = vm.popStack();
                if (arg_object) |object| {
                    const string_object = try Object.String.fromObject(vm, object);
                    const string_bytes = string_object.bytes[0..string_object.length];
                    try stream_writer.writer.writeAll(string_bytes);
                } else {
                    return error.InvalidArgument;
                }

                frame.ip += 1;
            },
            .stream_flush => {
                frame.ip += 1;

                // FIXME: This is a bit of a hack, but we have to deal with this right now.
                const buffered = stream_writer.writer.buffered();
                if (buffered.len == 0) continue;
                return stream_writer.toArrayList();
            },
            .br_push => {
                const arg_offset = readAddress(code, frame.ip);

                try vm.current_choices.append(gpa, .{
                    .text = stream_writer.toArrayList(),
                    .dest_offset = arg_offset,
                });
                frame.ip += 3;
            },
            .br_table => {
                // No-op currently.
                frame.ip += 1;
            },
            .br_select_index => {
                vm.can_advance = false;
                frame.ip += 1;
                return .empty;
            },
            .br_dispatch => {
                const index = vm.choice_index;
                const branch_dispatch = vm.current_choices.items[index];
                defer vm.current_choices.clearRetainingCapacity();

                frame.ip = branch_dispatch.dest_offset;
            },
            else => return error.InvalidInstruction,
        }
    }
}

fn readAddress(code: []const Story.Opcode, offset: usize) u16 {
    const arg_offset = std.mem.bytesToValue(u16, code[offset + 1 ..][0..2]);
    return std.mem.bigToNative(u16, arg_offset);
}

pub fn advance(story: *Story, gpa: std.mem.Allocator) ![]const u8 {
    var content = try story.execute();
    return content.toOwnedSlice(gpa);
}

fn divert(vm: *Story, path_name: []const u8) !void {
    const gpa = vm.allocator;
    const path_object: ?*Object.ContentPath = blk: {
        for (vm.paths.items) |object| {
            const current_path: *Object.ContentPath = @ptrCast(object);
            const current_name = current_path.name;
            // TODO(Brett): We probably should create a method for doing this.
            const name_bytes = current_name.bytes[0..current_name.length];
            if (std.mem.eql(u8, name_bytes, path_name)) break :blk current_path;
        }
        break :blk null;
    };
    if (path_object) |path| {
        // TODO(Brett): Add arguments?
        const stack_needed = path.arity + path.locals_count;
        const stack_ptr = vm.stack.items.len;
        try vm.stack.ensureUnusedCapacity(gpa, stack_needed);
        try vm.call_stack.ensureUnusedCapacity(gpa, 1);

        vm.stack.appendNTimesAssumeCapacity(null, stack_needed);
        vm.call_stack.appendAssumeCapacity(.{ .ip = 0, .sp = stack_ptr, .callee = path });
    } else return error.InvalidPath;
}

pub const LoadOptions = struct {
    dump_writer: ?*std.Io.Writer = null,
    stderr_writer: *std.Io.Writer,
    use_color: bool = true,
};

pub fn selectChoiceIndex(story: *Story, index: usize) !void {
    if (index >= story.current_choices.items.len) return error.InvalidChoice;
    story.choice_index = index;
    story.can_advance = true;
}

pub fn loadFromString(
    gpa: std.mem.Allocator,
    source_bytes: [:0]const u8,
    options: LoadOptions,
) !Story {
    var arena_allocator = std.heap.ArenaAllocator.init(gpa);
    defer arena_allocator.deinit();

    const arena = arena_allocator.allocator();
    const ast = try Ast.parse(gpa, arena, source_bytes, "<STDIN>", 0);

    if (options.dump_writer) |w| {
        try ast.render(gpa, w, .{
            .use_color = options.use_color,
        });
    }
    if (ast.errors.len > 0) {
        try ast.renderErrors(gpa, options.stderr_writer, .{
            .use_color = options.use_color,
        });
        return error.Invalid;
    }

    var story = try AstGen.generate(gpa, &ast);
    errdefer story.deinit();

    try story.divert("@main@");
    story.dump_writer = options.dump_writer;
    story.can_advance = true;
    return story;
}