//! Virtual machine state for story execution.
const std = @import("std");
const Compilation = @import("compile.zig").Compilation;
const tokenizer = @import("tokenizer.zig");
const Ast = @import("Ast.zig");
const AstGen = @import("AstGen.zig");
const Sema = @import("Sema.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,
constants_pool: std.ArrayListUnmanaged(*Object) = .empty,
globals: std.StringHashMapUnmanaged(?*Object) = .empty,
stack: std.ArrayListUnmanaged(?*Object) = .empty,
call_stack: std.ArrayListUnmanaged(CallFrame) = .empty,
stack_max: usize = 128,
gc_objects: std.SinglyLinkedList = .{},
// FIXME: This was a hack to keep string bytes alive.
string_bytes: []const u8 = &.{},

pub const default_knot_name: [:0]const u8 = "$__main__$";

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.constants_pool.deinit(gpa);
    story.globals.deinit(gpa);
    story.stack.deinit(gpa);
    story.call_stack.deinit(gpa);
    gpa.free(story.string_bytes);
}

pub fn dump(story: *Story, writer: *std.Io.Writer) !void {
    const story_dumper: Dumper = .{ .story = story, .writer = writer };

    try writer.writeAll("=== Constants ===\n");
    for (story.constants_pool.items) |global_constant| {
        try story_dumper.dumpObject(global_constant);
        try writer.writeAll("\n");
    }

    try writer.writeAll("\n");
    try writer.writeAll("=== Globals ===\n");

    var global_iter = story.globals.iterator();
    while (global_iter.next()) |entry| {
        try writer.print("{s} => ...\n", .{entry.key_ptr.*});
    }

    try writer.writeAll("\n");
    try writer.writeAll("=== Knots ===\n");

    var knots_iter = story.globals.iterator();
    while (knots_iter.next()) |entry| {
        if (entry.value_ptr.*) |global| {
            switch (global.tag) {
                .content_path => try story_dumper.dump(@ptrCast(global)),
                else => {},
            }
        }
    }
}

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) {
        // FIXME: There has to be a better way to do this.
        if (stack_top > 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 (stack.items[stack.items.len - 1]) |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;
    if (stack_top <= offset) return null;

    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: *Story, frame: *CallFrame, offset: u8) !*Object {
    if (offset >= frame.callee.const_pool.len) return error.InvalidArgument;

    const constant_index = frame.callee.const_pool[offset];
    return story.constants_pool.items[constant_index];
}

fn getLocal(vm: *Story, frame: *CallFrame, offset: u8) ?*Object {
    const stack_top = vm.stack.items.len;
    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.items.len;
    const stack_offset = frame.sp + offset;
    assert(stack_top > stack_offset);

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

// TODO: This should probably check the constants table first.
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;
}

// TODO: This should probably check the constants table first.
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;

    var stream_writer = std.Io.Writer.Allocating.init(gpa);
    defer stream_writer.deinit();

    while (true) {
        const frame = vm.currentFrame();
        const code = std.mem.bytesAsSlice(Opcode, frame.callee.bytes);
        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 => {
                // FIXME: This should be more strict.
                // Its not because theres a bug in when these instructions are
                // emitted.
                if (vm.peekStack(0)) |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);
                    _ = vm.popStack();
                } //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;
            },
            .divert => {
                const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                frame.ip += 2;

                if (peekStack(vm, arg_offset)) |knot| {
                    try divertToKnot(vm, @ptrCast(knot));
                } else {
                    return error.InvalidArgument;
                }
            },
            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);
}

pub fn getKnot(vm: *Story, name: []const u8) ?*Object.ContentPath {
    const knot: ?*Object.ContentPath = blk: {
        if (vm.globals.get(name)) |object| {
            break :blk @ptrCast(object);
        }
        break :blk null;
    };
    return knot;
}

// TODO(Brett): Add arguments?
fn divertToKnot(vm: *Story, knot: *Object.ContentPath) !void {
    const gpa = vm.allocator;
    const stack_ptr = vm.stack.items.len - knot.arity;
    const stack_needed = knot.locals_count;

    try vm.stack.ensureUnusedCapacity(gpa, stack_needed);
    try vm.call_stack.ensureUnusedCapacity(gpa, 1);

    vm.call_stack.appendAssumeCapacity(.{
        .callee = knot,
        .ip = 0,
        .sp = stack_ptr,
    });
    vm.stack.appendNTimesAssumeCapacity(null, stack_needed);
    vm.can_advance = true;
}

fn divert(vm: *Story, knot_name: []const u8) !void {
    return if (getKnot(vm, knot_name)) |knot| {
        return divertToKnot(vm, knot);
    } else return error.InvalidPath;
}

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

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 comp = try Compilation.compile(gpa, .{
        .source_bytes = source_bytes,
        .filename = "<STDIN>",
        .dump_writer = options.dump_writer,
        .dump_use_color = options.use_color,
        .dump_ast = options.dump_ast,
        .dump_ir = options.dump_ir,
    });
    defer comp.deinit();

    if (comp.errors.len > 0) {
        for (comp.errors) |err| {
            try comp.renderError(options.error_writer, err);
        }
        return error.Fail;
    }

    var story: Story = .{
        .allocator = gpa,
        .can_advance = false,
        .dump_writer = options.dump_writer,
    };
    errdefer story.deinit();
    try comp.setupStoryRuntime(gpa, &story);
    if (story.getKnot(Story.default_knot_name)) |knot| {
        try story.divertToKnot(knot);
        story.can_advance = true;
    }
    return story;
}