//! 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");
const Module = @import("compile.zig").Module;
pub const Object = @import("Story/Object.zig");
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,
code_chunks: std.ArrayListUnmanaged(*Object.Code) = .empty,
constants_pool: std.ArrayListUnmanaged(Value) = .empty,
globals: std.StringHashMapUnmanaged(?Value) = .empty,
stack: std.ArrayListUnmanaged(?Value) = .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 Value = union(enum) {
    bool: bool,
    int: i64,
    float: f64,
    object: *Object,

    pub fn tagBytes(v: Value) []const u8 {
        return switch (v) {
            .bool => "Bool",
            .int => "Int",
            .float => "Float",
            .object => "Object",
        };
    }

    pub fn isNumeric(v: Value) bool {
        return v == .int or v == .float;
    }

    pub fn isTruthy(v: Value) bool {
        return switch (v) {
            //.nil => false,
            .bool => |b| b,
            .int => |i| i != 0,
            .float => |f| f != 0.0,
            .object => true,
        };
    }

    pub fn toFloat(v: Value) f64 {
        return switch (v) {
            .int => |i| @floatFromInt(i),
            .float => |f| f,
            else => unreachable,
        };
    }

    pub fn add(lhs: Value, rhs: Value, story: *Story) !Value {
        return arith(lhs, rhs, .add, story);
    }

    pub fn arith(lhs: Value, rhs: Value, op: Opcode, story: *Story) !Value {
        if (lhs.isNumeric() and rhs.isNumeric())
            return numericArith(lhs, rhs, op);
        if (op == .add) {
            if (lhs == .object and lhs.object.tag == .string)
                return concat(lhs, rhs, story);
            if (rhs == .object and rhs.object.tag == .string)
                return concat(lhs, rhs, story);
        }
        return error.TypeError;
    }

    pub fn numericArith(lhs: Value, rhs: Value, op: Story.Opcode) !Value {
        if (!lhs.isNumeric() or !rhs.isNumeric()) return error.TypeError;

        if (lhs == .int and rhs == .int) {
            switch (op) {
                .add => return .{ .int = lhs.int +% rhs.int },
                .sub => return .{ .int = lhs.int -% rhs.int },
                .mul => return .{ .int = lhs.int *% rhs.int },
                .div => {
                    if (rhs.int == 0)
                        return error.DivisionByZero;
                    return .{ .int = @divTrunc(lhs.int, rhs.int) };
                },
                .mod => if (rhs.int == 0)
                    return error.DivisionByZero
                else
                    return .{ .int = @mod(lhs.int, rhs.int) },
                else => unreachable,
            }
        }

        const l = lhs.toFloat();
        const r = rhs.toFloat();
        switch (op) {
            .add => return .{ .float = l + r },
            .sub => return .{ .float = l - r },
            .mul => return .{ .float = l * r },
            .div => if (r == 0.0)
                return error.DivisionByZero
            else
                return .{ .float = l / r },
            .mod => if (r == 0.0)
                return error.DivisionByZero
            else
                return .{ .float = @mod(l, r) },
            else => unreachable,
        }
    }

    pub fn eql(lhs: Value, rhs: Value) bool {
        return switch (lhs) {
            .bool => |l| rhs == .bool and l == rhs.bool,
            .int => |l| switch (rhs) {
                .int => |r| l == r,
                .float => |r| @as(f64, @floatFromInt(l)) == r,
                else => false,
            },
            .float => |l| switch (rhs) {
                .int => |r| l == @as(f64, @floatFromInt(r)),
                .float => |r| l == r,
                else => false,
            },
            .object => |l| rhs == .object and l == rhs.object,
        };
    }

    pub fn compare(lhs: Value, rhs: Value, op: Story.Opcode) !Value {
        return switch (op) {
            .cmp_eq => .{ .bool = lhs.eql(rhs) },
            .cmp_neq => .{ .bool = !lhs.eql(rhs) },
            .cmp_lt, .cmp_gt, .cmp_lte, .cmp_gte => blk: {
                if (!lhs.isNumeric() or !rhs.isNumeric()) return error.TypeError;
                const l = lhs.toFloat();
                const r = rhs.toFloat();
                break :blk .{
                    .bool = switch (op) {
                        .cmp_lt => l < r,
                        .cmp_gt => l > r,
                        .cmp_lte => l <= r,
                        .cmp_gte => l >= r,
                        else => unreachable,
                    },
                };
            },
            else => unreachable,
        };
    }

    pub fn logicalNot(lhs: Value) Value {
        return .{ .bool = !lhs.isTruthy() };
    }

    pub fn concat(lhs: Value, rhs: Value, story: *Story) !Value {
        const lhs_object = try Object.String.fromValue(story, lhs);
        const rhs_object = try Object.String.fromValue(story, rhs);
        const str_object = try Object.String.concat(story, lhs_object, rhs_object);
        return .{ .object = &str_object.base };
    }

    pub fn negate(lhs: Value) !Value {
        switch (lhs) {
            .bool => return error.TypeError,
            .int => |int| return .{ .int = -int },
            .float => |float| return .{ .float = -float },
            .object => return error.TypeError,
        }
    }
};

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

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_neq,
    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,
    load_attr,
    store_attr,
    /// 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);
}

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) ?Value {
    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, value: Value) !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, value);
}

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

fn getConstant(story: *Story, frame: *CallFrame, offset: u8) !Value {
    if (offset >= frame.callee.code.constants.len) return error.InvalidArgument;

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

fn getLocal(vm: *Story, frame: *CallFrame, offset: u8) ?Value {
    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: Value) 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: Value) !Value {
    switch (key) {
        .object => |object| switch (object.tag) {
            .string => {
                const str_object: *Object.String = @ptrCast(object);
                if (vm.globals.get(str_object.toSlice())) |value| {
                    // FIXME: Support for nil
                    return value.?;
                }
                return error.InvalidVariable;
            },
            else => return error.TypeError,
        },
        else => return error.TypeError,
    }
}

// TODO: This should probably check the constants table first.
fn setGlobal(vm: *Story, key: Value, value: Value) !void {
    switch (key) {
        .object => |object| switch (object.tag) {
            .string => {
                const str_object: *Object.String = @ptrCast(object);
                return vm.globals.putAssumeCapacity(str_object.toSlice(), value);
            },
            else => return error.TypeError,
        },
        else => return error.TypeError,
    }
}

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.code.bytecode);
        if (vm.dump_writer) |w| {
            Dumper.trace(vm, w, frame) catch {};
        }
        switch (code[frame.ip]) {
            .exit => {
                vm.is_exited = true;
                vm.can_advance = false;
                return .empty;
            },
            .true => {
                const value: Value = .{ .bool = true };
                try vm.pushStack(value);
                frame.ip += 1;
            },
            .false => {
                const value: Value = .{ .bool = false };
                try vm.pushStack(value);
                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 Value.add(lhs, rhs, vm);

                _ = 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 Value.arith(lhs, rhs, op, vm);

                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(value);
                frame.ip += 1;
            },
            .neg => {
                if (vm.popStack()) |arg| {
                    const value = try Value.negate(arg);
                    try vm.pushStack(value);
                } else {
                    return error.InvalidArgument;
                }
                frame.ip += 1;
            },
            .not => {
                if (vm.popStack()) |arg| {
                    const value = Value.logicalNot(arg);
                    try vm.pushStack(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 Value.compare(lhs, rhs, .cmp_eq);
                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(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 Value.compare(lhs, rhs, op);
                _ = vm.popStack();
                _ = vm.popStack();
                try vm.pushStack(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.isTruthy()) {
                        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.isTruthy()) {
                        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);
                const global_value = try vm.getGlobal(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);

                if (vm.peekStack(0)) |arg| {
                    try vm.setGlobal(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)) |arg| {
                    const str_object = try Object.String.fromValue(vm, arg);
                    try stream_writer.writer.writeAll(str_object.toSlice());
                    _ = 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)) |value| {
                    try divertToValue(vm, value);
                } else {
                    return error.InvalidArgument;
                }
            },
            .load_attr => {
                const arg_offset: u8 = @intFromEnum(code[frame.ip + 1]);
                frame.ip += 2;

                if (peekStack(vm, 0)) |value| {
                    const knot_object: *Object.Knot = @ptrCast(value.object);
                    const arg_value = try vm.getConstant(frame, arg_offset);
                    const knot_attr: *Object.String = @ptrCast(arg_value.object);

                    _ = popStack(vm);

                    if (knot_object.members.get(knot_attr.toSlice())) |attr_object| {
                        try vm.pushStack(.{ .object = attr_object });
                    } else {
                        return error.InvalidArgument;
                    }
                } 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.Knot {
    const knot: ?*Object.Knot = blk: {
        if (vm.globals.get(name)) |value| {
            // TODO: Do a check here.
            break :blk @ptrCast(value.?.object);
        }
        break :blk null;
    };
    return knot;
}

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

    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 divertToValue(vm: *Story, value: Value) !void {
    switch (value) {
        .object => |object| switch (object.tag) {
            .knot => try divertToKnot(vm, @ptrCast(object)),
            else => return error.TypeError,
        },
        else => return error.TypeError,
    }
}

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,
    dump_trace: 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 dump(story: *Story, writer: *std.Io.Writer) !void {
    return Dumper.dump(story, writer);
}

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();
    var comp = try Module.compile(gpa, arena, .{
        .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.items.len > 0) {
        for (comp.errors.items) |err| {
            try comp.renderError(options.error_writer, err);
        }
        return error.LoadFailed;
    }

    var story: Story = .{
        .allocator = gpa,
        .can_advance = false,
        .dump_writer = if (options.dump_trace) options.dump_writer else null,
    };
    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;
}