use gc_arena::{Gc, Mutation};
use inkwell::AddressSpace;
use inkwell::context::Context;
use inkwell::execution_engine::ExecutionEngine;
use inkwell::module::Module;
use inkwell::types::{FloatType, FunctionType, IntType, PointerType, StructType};
use inkwell::values::{BasicValueEnum, FunctionValue};

use crate::bytecode::OpCodes;
use crate::env::VmEnv;
use crate::jit::JITValueData;
use crate::ty::internal::{Thunk, Value};
use crate::vm::VM;

use super::{JITContext, JITValue, ValueTag};

pub struct Helpers<'ctx> {
    pub int_type: IntType<'ctx>,
    pub float_type: FloatType<'ctx>,
    pub bool_type: IntType<'ctx>,
    pub ptr_int_type: IntType<'ctx>,
    pub ptr_type: PointerType<'ctx>,
    pub value_type: StructType<'ctx>,
    pub func_type: FunctionType<'ctx>,

    pub new_thunk: FunctionValue<'ctx>,

    pub debug: FunctionValue<'ctx>,
    pub capture_env: FunctionValue<'ctx>,
    pub neg: FunctionValue<'ctx>,
    pub not: FunctionValue<'ctx>,
    pub add: FunctionValue<'ctx>,
    pub sub: FunctionValue<'ctx>,
    pub eq: FunctionValue<'ctx>,
    pub or: FunctionValue<'ctx>,
    pub call: FunctionValue<'ctx>,
    pub lookup: FunctionValue<'ctx>,
    pub force: FunctionValue<'ctx>,
}

impl<'ctx> Helpers<'ctx> {
    pub fn new(
        context: &'ctx Context,
        module: &Module<'ctx>,
        execution_engine: &ExecutionEngine<'ctx>,
    ) -> Self {
        let int_type = context.i64_type();
        let float_type = context.f64_type();
        let bool_type = context.bool_type();
        let ptr_int_type = context.ptr_sized_int_type(execution_engine.get_target_data(), None);
        let ptr_type = context.ptr_type(AddressSpace::default());
        let value_type = context.struct_type(&[int_type.into(), int_type.into()], false);
        let func_type = value_type.fn_type(&[ptr_type.into(), ptr_type.into()], false);

        let new_thunk = module.add_function(
            "new_thunk",
            value_type.fn_type(&[ptr_type.into(), ptr_type.into()], false),
            None,
        );
        let debug = module.add_function(
            "debug",
            context.void_type().fn_type(&[value_type.into()], false),
            None,
        );
        let capture_env = module.add_function(
            "capture_env",
            context.void_type().fn_type(
                &[value_type.into(), ptr_type.into(), ptr_type.into()],
                false,
            ),
            None,
        );
        let neg = module.add_function(
            "neg",
            value_type.fn_type(&[value_type.into(), ptr_type.into()], false),
            None,
        );
        let not = module.add_function(
            "not",
            value_type.fn_type(&[value_type.into(), ptr_type.into()], false),
            None,
        );
        let add = module.add_function(
            "add",
            value_type.fn_type(&[value_type.into(), value_type.into()], false),
            None,
        );
        let sub = module.add_function(
            "sub",
            value_type.fn_type(&[value_type.into(), value_type.into()], false),
            None,
        );
        let eq = module.add_function(
            "eq",
            value_type.fn_type(&[value_type.into(), value_type.into()], false),
            None,
        );
        let or = module.add_function(
            "or",
            value_type.fn_type(&[value_type.into(), value_type.into()], false),
            None,
        );
        let call = module.add_function(
            "call",
            value_type.fn_type(
                &[
                    value_type.into(),
                    value_type.into(),
                    ptr_type.into(),
                    ptr_type.into(),
                ],
                false,
            ),
            None,
        );
        let lookup = module.add_function(
            "lookup",
            value_type.fn_type(&[ptr_int_type.into(), ptr_type.into()], false),
            None,
        );
        let force = module.add_function(
            "force",
            value_type.fn_type(
                &[
                    value_type.into(),
                    ptr_type.into(),
                    ptr_type.into(),
                    ptr_type.into(),
                ],
                false,
            ),
            None,
        );

        execution_engine.add_global_mapping(&new_thunk, helper_new_thunk as _);
        execution_engine.add_global_mapping(&debug, helper_debug as _);
        execution_engine.add_global_mapping(&capture_env, helper_capture_env as _);
        execution_engine.add_global_mapping(&neg, helper_neg as _);
        execution_engine.add_global_mapping(&not, helper_not as _);
        execution_engine.add_global_mapping(&add, helper_add as _);
        execution_engine.add_global_mapping(&sub, helper_sub as _);
        execution_engine.add_global_mapping(&eq, helper_eq as _);
        execution_engine.add_global_mapping(&or, helper_or as _);
        execution_engine.add_global_mapping(&call, helper_call as _);
        execution_engine.add_global_mapping(&lookup, helper_lookup as _);
        execution_engine.add_global_mapping(&force, helper_force as _);

        Helpers {
            int_type,
            float_type,
            bool_type,
            ptr_int_type,
            ptr_type,
            value_type,
            func_type,

            new_thunk,

            debug,
            capture_env,
            neg,
            not,
            add,
            sub,
            eq,
            or,
            call,
            lookup,
            force,
        }
    }

    pub fn new_int(&self, int: i64) -> BasicValueEnum<'ctx> {
        self.value_type
            .const_named_struct(&[
                self.int_type.const_int(ValueTag::Int as _, false).into(),
                self.int_type.const_int(int as _, false).into(),
            ])
            .into()
    }

    pub fn new_float(&self, float: f64) -> BasicValueEnum<'ctx> {
        self.value_type
            .const_named_struct(&[
                self.int_type.const_int(ValueTag::Float as _, false).into(),
                self.float_type.const_float(float).into(),
            ])
            .into()
    }

    pub fn new_bool(&self, bool: bool) -> BasicValueEnum<'ctx> {
        self.value_type
            .const_named_struct(&[
                self.int_type.const_int(ValueTag::Bool as _, false).into(),
                self.bool_type.const_int(bool as _, false).into(),
            ])
            .into()
    }

    pub fn new_null(&self) -> BasicValueEnum<'ctx> {
        self.value_type
            .const_named_struct(&[
                self.int_type.const_int(ValueTag::Null as _, false).into(),
                self.int_type.const_zero().into(),
            ])
            .into()
    }

    pub fn const_string(&self, string: *const u8) -> BasicValueEnum<'ctx> {
        self.value_type
            .const_named_struct(&[
                self.int_type.const_int(ValueTag::String as _, false).into(),
                self.ptr_int_type.const_int(string as _, false).into(),
            ])
            .into()
    }
}

extern "C" fn helper_debug(value: JITValue) {
    dbg!(value.tag);
}

extern "C" fn helper_capture_env<'gc>(
    thunk: JITValue,
    env: *const VmEnv<'gc>,
    mc: *const Mutation<'gc>,
) {
    let thunk = unsafe { (thunk.data.ptr as *const Thunk).as_ref().unwrap() };
    let env = unsafe { Gc::from_ptr(env) };
    thunk.capture_env(env, unsafe { mc.as_ref() }.unwrap());
}

extern "C" fn helper_neg(rhs: JITValue, _env: *const VmEnv) -> JITValue {
    use ValueTag::*;
    match rhs.tag {
        Int => JITValue {
            tag: Int,
            data: JITValueData {
                int: -unsafe { rhs.data.int },
            },
        },
        Float => JITValue {
            tag: Float,
            data: JITValueData {
                float: -unsafe { rhs.data.float },
            },
        },
        _ => todo!(),
    }
}

extern "C" fn helper_not(rhs: JITValue, _env: *const VmEnv) -> JITValue {
    use ValueTag::*;
    match rhs.tag {
        Bool => JITValue {
            tag: Bool,
            data: JITValueData {
                bool: !unsafe { rhs.data.bool },
            },
        },
        _ => todo!(),
    }
}

extern "C" fn helper_add(lhs: JITValue, rhs: JITValue) -> JITValue {
    use ValueTag::*;
    match (lhs.tag, rhs.tag) {
        (Int, Int) => JITValue {
            tag: Int,
            data: JITValueData {
                int: unsafe { lhs.data.int + rhs.data.int },
            },
        },
        _ => todo!(
            "Addition not implemented for {:?} and {:?}",
            lhs.tag,
            rhs.tag
        ),
    }
}

extern "C" fn helper_sub(lhs: JITValue, rhs: JITValue) -> JITValue {
    use ValueTag::*;
    match (lhs.tag, rhs.tag) {
        (Int, Int) => JITValue {
            tag: Int,
            data: JITValueData {
                int: unsafe { lhs.data.int - rhs.data.int },
            },
        },
        _ => todo!(
            "Substruction not implemented for {:?} and {:?}",
            lhs.tag,
            rhs.tag
        ),
    }
}

extern "C" fn helper_eq(lhs: JITValue, rhs: JITValue) -> JITValue {
    use ValueTag::*;
    match (lhs.tag, rhs.tag) {
        (Int, Int) => JITValue {
            tag: Bool,
            data: JITValueData {
                bool: unsafe { lhs.data.int == rhs.data.int },
            },
        },
        _ => todo!(
            "Equation not implemented for {:?} and {:?}",
            lhs.tag,
            rhs.tag
        ),
    }
}

extern "C" fn helper_or(lhs: JITValue, rhs: JITValue) -> JITValue {
    use ValueTag::*;
    match (lhs.tag, rhs.tag) {
        (Bool, Bool) => JITValue {
            tag: Bool,
            data: JITValueData {
                bool: unsafe { lhs.data.bool || rhs.data.bool },
            },
        },
        _ => todo!(
            "Substraction not implemented for {:?} and {:?}",
            lhs.tag,
            rhs.tag
        ),
    }
}

extern "C" fn helper_call<'gc>(
    func: JITValue,
    arg: JITValue,
    vm: *const VM<'gc>,
    mc: *const Mutation<'gc>,
) -> JITValue {
    let vm = unsafe { vm.as_ref() }.unwrap();
    let mc = unsafe { mc.as_ref() }.unwrap();
    let arg = Value::from(arg);
    match func.tag {
        ValueTag::Function => {
            let func = Value::from(func).unwrap_func();
            func.call_compile(arg, vm, mc).unwrap().into()
        }
        _ => todo!(),
    }
}

extern "C" fn helper_lookup(sym: usize, env: *const VmEnv) -> JITValue {
    let env = unsafe { env.as_ref() }.unwrap();
    let val: JITValue = env.lookup_slow(&sym).unwrap().into();
    val
}

extern "C" fn helper_force<'gc>(
    thunk: JITValue,
    vm: *const VM<'gc>,
    mc: *const Mutation<'gc>,
    jit: *const JITContext<'gc>,
) -> JITValue {
    if !matches!(thunk.tag, ValueTag::Thunk) {
        return thunk;
    }

    let vm = unsafe { vm.as_ref() }.unwrap();
    let mc = unsafe { mc.as_ref() }.unwrap();
    let thunk = Value::from(thunk).unwrap_thunk();
    if let Some(val) = thunk.get_value() {
        return val.into();
    }
    let (opcodes, env) = thunk.suspend(mc).unwrap();
    let func = unsafe { jit.as_ref() }
        .unwrap()
        .compile_seq(opcodes.iter().copied(), vm)
        .unwrap();
    let val = unsafe { func.call(env.as_ref() as *const _, mc as *const _) };
    thunk.insert_value(val.into(), mc);
    val
}

extern "C" fn helper_new_thunk(opcodes: *const OpCodes, mc: *const Mutation) -> JITValue {
    Value::Thunk(Thunk::new(
        unsafe { opcodes.as_ref() }.unwrap(),
        unsafe { mc.as_ref() }.unwrap(),
    ))
    .into()
}