#![cfg(feature = "tailcall")]

use gc_arena::Mutation;

use crate::{Break, BytecodeReader, Step, Vm, VmRuntimeCtx};

pub(crate) enum TailResult {
    YieldFuel(u32),
    Done,
}

pub(crate) type OpFn<'gc, C> = extern "rust-preserve-none" fn(
    &mut Vm<'gc>,
    &Mutation<'gc>,
    &mut C,
    &[u8],
    &DispatchTable<'gc, C>,
    u32,
    u32,
) -> TailResult;

pub(crate) struct DispatchTable<'gc, C: VmRuntimeCtx>(pub(crate) [OpFn<'gc, C>; 256]);

extern "rust-preserve-none" fn op_illegal<'gc, C: VmRuntimeCtx>(
    _vm: &mut Vm<'gc>,
    _mc: &Mutation<'gc>,
    _ctx: &mut C,
    _bc: &[u8],
    _table: &DispatchTable<'gc, C>,
    pc: u32,
    _fuel: u32,
) -> TailResult {
    panic!("illegal opcode at pc = {pc}");
}

macro_rules! tail_dispatch_after {
    ($result:expr, $new_pc:expr, $vm:ident, $mc:ident, $ctx:ident, $bc:ident, $table:ident, $fuel:ident) => {{
        match $result {
            Step::Continue(()) | Step::Break(Break::Force) => {}
            Step::Break(Break::Done) => return TailResult::Done,
        }
        let new_pc: u32 = $new_pc;
        if $fuel == 0 {
            return TailResult::YieldFuel(new_pc);
        }
        let next_op = $bc[new_pc as usize] as usize;
        become $table.0[next_op]($vm, $mc, $ctx, $bc, $table, new_pc, $fuel - 1)
    }};
}

macro_rules! tail_fn {
    ($name:ident, ()) => {
        extern "rust-preserve-none" fn $name<'gc, C: VmRuntimeCtx>(
            vm: &mut Vm<'gc>,
            mc: &Mutation<'gc>,
            ctx: &mut C,
            bc: &[u8],
            table: &DispatchTable<'gc, C>,
            pc: u32,
            fuel: u32,
        ) -> TailResult {
            let result = vm.$name();
            tail_dispatch_after!(result, pc + 1, vm, mc, ctx, bc, table, fuel)
        }
    };
    ($name:ident, (reader)) => {
        extern "rust-preserve-none" fn $name<'gc, C: VmRuntimeCtx>(
            vm: &mut Vm<'gc>,
            mc: &Mutation<'gc>,
            ctx: &mut C,
            bc: &[u8],
            table: &DispatchTable<'gc, C>,
            pc: u32,
            fuel: u32,
        ) -> TailResult {
            let mut reader = BytecodeReader::from_after_op(bc, pc as usize);
            let result = vm.$name(&mut reader);
            tail_dispatch_after!(result, reader.pc() as u32, vm, mc, ctx, bc, table, fuel)
        }
    };
    ($name:ident, (reader, mc)) => {
        extern "rust-preserve-none" fn $name<'gc, C: VmRuntimeCtx>(
            vm: &mut Vm<'gc>,
            mc: &Mutation<'gc>,
            ctx: &mut C,
            bc: &[u8],
            table: &DispatchTable<'gc, C>,
            pc: u32,
            fuel: u32,
        ) -> TailResult {
            let mut reader = BytecodeReader::from_after_op(bc, pc as usize);
            let result = vm.$name(&mut reader, mc);
            tail_dispatch_after!(result, reader.pc() as u32, vm, mc, ctx, bc, table, fuel)
        }
    };
    ($name:ident, (ctx, reader, mc)) => {
        extern "rust-preserve-none" fn $name<'gc, C: VmRuntimeCtx>(
            vm: &mut Vm<'gc>,
            mc: &Mutation<'gc>,
            ctx: &mut C,
            bc: &[u8],
            table: &DispatchTable<'gc, C>,
            pc: u32,
            fuel: u32,
        ) -> TailResult {
            let mut reader = BytecodeReader::from_after_op(bc, pc as usize);
            let result = vm.$name(ctx, &mut reader, mc);
            tail_dispatch_after!(result, reader.pc() as u32, vm, mc, ctx, bc, table, fuel)
        }
    };
    ($name:ident, (ctx)) => {
        extern "rust-preserve-none" fn $name<'gc, C: VmRuntimeCtx>(
            vm: &mut Vm<'gc>,
            mc: &Mutation<'gc>,
            ctx: &mut C,
            bc: &[u8],
            table: &DispatchTable<'gc, C>,
            pc: u32,
            fuel: u32,
        ) -> TailResult {
            let result = vm.$name(ctx);
            tail_dispatch_after!(result, pc + 1, vm, mc, ctx, bc, table, fuel)
        }
    };
}

tail_fn!(op_push_smi, (reader));
tail_fn!(op_push_bigint, (reader, mc));
tail_fn!(op_push_float, (reader));
tail_fn!(op_push_string, (reader));
tail_fn!(op_push_null, ());
tail_fn!(op_push_true, ());
tail_fn!(op_push_false, ());

tail_fn!(op_load_local, (reader));
tail_fn!(op_load_outer, (reader));
tail_fn!(op_store_local, (reader, mc));
tail_fn!(op_alloc_locals, (reader, mc));

tail_fn!(op_make_thunk, (reader, mc));
tail_fn!(op_make_closure, (reader, mc));
tail_fn!(op_make_pattern_closure, (reader, mc));

tail_fn!(op_call, (ctx, reader, mc));
tail_fn!(op_dispatch_primop, (ctx, reader, mc));
tail_fn!(op_return, (ctx, reader, mc));

tail_fn!(op_make_attrs, (ctx, reader, mc));
tail_fn!(op_make_empty_attrs, ());
tail_fn!(op_select_static, (ctx, reader, mc));
tail_fn!(op_select_dynamic, (ctx, reader, mc));
tail_fn!(op_has_attr_path_static, (ctx, reader, mc));
tail_fn!(op_has_attr_path_dynamic, (ctx, reader, mc));
tail_fn!(op_jump_if_select_failed, (reader));
tail_fn!(op_jump_if_select_succeeded, (reader));
tail_fn!(op_has_attr_static, (reader, mc));
tail_fn!(op_has_attr_dynamic, (ctx, reader, mc));
tail_fn!(op_has_attr_resolve, ());

tail_fn!(op_make_list, (ctx, reader, mc));
tail_fn!(op_make_empty_list, ());

tail_fn!(op_add, (ctx, reader, mc));
tail_fn!(op_sub, (reader, mc));
tail_fn!(op_mul, (reader, mc));
tail_fn!(op_div, (reader, mc));
tail_fn!(op_eq, (ctx, reader, mc));
tail_fn!(op_neq, (ctx, reader, mc));
tail_fn!(op_lt, (ctx, reader, mc));
tail_fn!(op_gt, (ctx, reader, mc));
tail_fn!(op_leq, (ctx, reader, mc));
tail_fn!(op_geq, (ctx, reader, mc));
tail_fn!(op_concat, (reader, mc));
tail_fn!(op_update, (reader, mc));

tail_fn!(op_neg, ());
tail_fn!(op_not, ());

tail_fn!(op_jump_if_false, (reader, mc));
tail_fn!(op_jump_if_true, (reader, mc));
tail_fn!(op_jump, (reader));

tail_fn!(op_concat_strings, (ctx, reader, mc));
tail_fn!(op_resolve_path, (ctx));

tail_fn!(op_assert, (reader));

tail_fn!(op_push_with, (ctx, reader, mc));
tail_fn!(op_pop_with, ());
tail_fn!(op_lookup_with, (ctx, reader, mc));
tail_fn!(op_prepare_with, ());

tail_fn!(op_load_builtins, ());
tail_fn!(op_load_builtin, (reader));

tail_fn!(op_mk_pos, (reader));
tail_fn!(op_load_repl_binding, (reader));
tail_fn!(op_load_scoped_binding, (reader));

macro_rules! table {
    ($($variant:ident => $fn:ident),* $(,)?) => {
        impl<'gc, C: VmRuntimeCtx> DispatchTable<'gc, C> {
            pub(crate) const NEW: Self = {
                let mut arr: [OpFn<'gc, C>; 256] = [op_illegal; 256];
                $( arr[fix_codegen::Op::$variant as usize] = $fn; )*
                DispatchTable(arr)
            };
        }

        // Exhaustiveness check: fails to compile if `fix_codegen::Op` gains,
        // loses, or renames a variant that isn't wired up above.
        #[allow(dead_code)]
        const _: fn(fix_codegen::Op) = |op| match op {
            $( fix_codegen::Op::$variant => (), )*
        };
    };
}

table! {
    PushSmi               => op_push_smi,
    PushBigInt            => op_push_bigint,
    PushFloat             => op_push_float,
    PushString            => op_push_string,
    PushNull              => op_push_null,
    PushTrue              => op_push_true,
    PushFalse             => op_push_false,

    LoadLocal             => op_load_local,
    LoadOuter             => op_load_outer,
    StoreLocal            => op_store_local,
    AllocLocals           => op_alloc_locals,

    MakeThunk             => op_make_thunk,
    MakeClosure           => op_make_closure,
    MakePatternClosure    => op_make_pattern_closure,

    Call                  => op_call,
    DispatchPrimOp        => op_dispatch_primop,
    Return                => op_return,

    MakeAttrs             => op_make_attrs,
    MakeEmptyAttrs        => op_make_empty_attrs,
    SelectStatic          => op_select_static,
    SelectDynamic         => op_select_dynamic,
    HasAttrPathStatic     => op_has_attr_path_static,
    HasAttrPathDynamic    => op_has_attr_path_dynamic,
    HasAttrStatic         => op_has_attr_static,
    HasAttrDynamic        => op_has_attr_dynamic,
    HasAttrResolve        => op_has_attr_resolve,
    JumpIfSelectSucceeded => op_jump_if_select_succeeded,
    JumpIfSelectFailed    => op_jump_if_select_failed,

    MakeList              => op_make_list,
    MakeEmptyList         => op_make_empty_list,

    OpAdd                 => op_add,
    OpSub                 => op_sub,
    OpMul                 => op_mul,
    OpDiv                 => op_div,
    OpEq                  => op_eq,
    OpNeq                 => op_neq,
    OpLt                  => op_lt,
    OpGt                  => op_gt,
    OpLeq                 => op_leq,
    OpGeq                 => op_geq,
    OpConcat              => op_concat,
    OpUpdate              => op_update,

    OpNeg                 => op_neg,
    OpNot                 => op_not,

    JumpIfFalse           => op_jump_if_false,
    JumpIfTrue            => op_jump_if_true,
    Jump                  => op_jump,

    ConcatStrings         => op_concat_strings,
    ResolvePath           => op_resolve_path,

    Assert                => op_assert,

    PushWith              => op_push_with,
    PopWith               => op_pop_with,
    LookupWith            => op_lookup_with,
    PrepareWith           => op_prepare_with,

    LoadBuiltins          => op_load_builtins,
    LoadBuiltin           => op_load_builtin,

    MkPos                 => op_mk_pos,
    LoadReplBinding       => op_load_repl_binding,
    LoadScopedBinding     => op_load_scoped_binding,

    Illegal               => op_illegal,
}

pub(crate) fn run_tailcall<'gc, C: VmRuntimeCtx>(
    vm: &mut Vm<'gc>,
    mc: &Mutation<'gc>,
    ctx: &mut C,
    bc: &[u8],
    pc: u32,
) -> TailResult {
    let table = &DispatchTable::<'gc, C>::NEW;
    let op = bc[pc as usize] as usize;
    table.0[op](vm, mc, ctx, bc, table, pc, Vm::DEFAULT_FUEL_AMOUNT)
}