use crate::relocation::{TrapData, TrapSink};
use crate::resolver::FuncResolver;
use crate::trampoline::Trampolines;
use hashbrown::HashSet;
use libc::c_void;
use std::{any::Any, cell::Cell, ptr::NonNull, sync::Arc};
use wasmer_runtime_core::{
    backend::{RunnableModule, UserTrapper},
    export::Context,
    module::{ExportIndex, ModuleInfo, ModuleInner},
    typed_func::{Wasm, WasmTrapInfo},
    types::{FuncIndex, FuncSig, LocalFuncIndex, LocalOrImport, SigIndex, Type, Value},
    vm::{self, ImportBacking},
};

#[cfg(unix)]
mod unix;

#[cfg(windows)]
mod windows;

#[cfg(unix)]
pub use self::unix::*;

#[cfg(windows)]
pub use self::windows::*;

thread_local! {
    pub static TRAP_EARLY_DATA: Cell<Option<Box<dyn Any>>> = Cell::new(None);
}

pub struct Trapper;

impl UserTrapper for Trapper {
    unsafe fn do_early_trap(&self, data: Box<dyn Any>) -> ! {
        TRAP_EARLY_DATA.with(|cell| cell.set(Some(data)));
        trigger_trap()
    }
}

pub struct Caller {
    func_export_set: HashSet<FuncIndex>,
    handler_data: HandlerData,
    trampolines: Arc<Trampolines>,
    resolver: FuncResolver,
}

impl Caller {
    pub fn new(
        module: &ModuleInfo,
        handler_data: HandlerData,
        trampolines: Arc<Trampolines>,
        resolver: FuncResolver,
    ) -> Self {
        let mut func_export_set = HashSet::new();
        for export_index in module.exports.values() {
            if let ExportIndex::Func(func_index) = export_index {
                func_export_set.insert(*func_index);
            }
        }
        if let Some(start_func_index) = module.start_func {
            func_export_set.insert(start_func_index);
        }

        Self {
            func_export_set,
            handler_data,
            trampolines,
            resolver,
        }
    }
}

impl RunnableModule for Caller {
    fn get_func(&self, _: &ModuleInfo, func_index: LocalFuncIndex) -> Option<NonNull<vm::Func>> {
        self.resolver.lookup(func_index)
    }

    fn get_trampoline(&self, _: &ModuleInfo, sig_index: SigIndex) -> Option<Wasm> {
        unsafe extern "C" fn invoke(
            trampoline: unsafe extern "C" fn(*mut vm::Ctx, NonNull<vm::Func>, *const u64, *mut u64),
            ctx: *mut vm::Ctx,
            func: NonNull<vm::Func>,
            args: *const u64,
            rets: *mut u64,
            trap_info: *mut WasmTrapInfo,
            invoke_env: Option<NonNull<c_void>>,
        ) -> bool {
            let handler_data = &*invoke_env.unwrap().cast().as_ptr();

            #[cfg(not(target_os = "windows"))]
            let res = call_protected(handler_data, || unsafe {
                // Leap of faith.
                trampoline(ctx, func, args, rets);
            })
            .is_ok();

            // the trampoline is called from C on windows
            #[cfg(target_os = "windows")]
            let res = call_protected(handler_data, trampoline, ctx, func, args, rets).is_ok();

            res
        }

        let trampoline = self
            .trampolines
            .lookup(sig_index)
            .expect("that trampoline doesn't exist");

        Some(unsafe {
            Wasm::from_raw_parts(
                trampoline,
                invoke,
                Some(NonNull::from(&self.handler_data).cast()),
            )
        })
    }

    fn get_early_trapper(&self) -> Box<dyn UserTrapper> {
        Box::new(Trapper)
    }
}

fn get_func_from_index<'a>(
    module: &'a ModuleInner,
    import_backing: &ImportBacking,
    func_index: FuncIndex,
) -> (NonNull<vm::Func>, Context, &'a FuncSig, SigIndex) {
    let sig_index = *module
        .info
        .func_assoc
        .get(func_index)
        .expect("broken invariant, incorrect func index");

    let (func_ptr, ctx) = match func_index.local_or_import(&module.info) {
        LocalOrImport::Local(local_func_index) => (
            module
                .runnable_module
                .get_func(&module.info, local_func_index)
                .expect("broken invariant, func resolver not synced with module.exports")
                .cast(),
            Context::Internal,
        ),
        LocalOrImport::Import(imported_func_index) => {
            let imported_func = import_backing.imported_func(imported_func_index);
            (
                NonNull::new(imported_func.func as *mut _).unwrap(),
                Context::External(imported_func.vmctx),
            )
        }
    };

    let signature = &module.info.signatures[sig_index];

    (func_ptr, ctx, signature, sig_index)
}

unsafe impl Send for HandlerData {}
unsafe impl Sync for HandlerData {}

#[derive(Clone)]
pub struct HandlerData {
    pub trap_data: Arc<TrapSink>,
    exec_buffer_ptr: *const c_void,
    exec_buffer_size: usize,
}

impl HandlerData {
    pub fn new(
        trap_data: Arc<TrapSink>,
        exec_buffer_ptr: *const c_void,
        exec_buffer_size: usize,
    ) -> Self {
        Self {
            trap_data,
            exec_buffer_ptr,
            exec_buffer_size,
        }
    }

    pub fn lookup(&self, ip: *const c_void) -> Option<TrapData> {
        let ip = ip as usize;
        let buffer_ptr = self.exec_buffer_ptr as usize;

        if buffer_ptr <= ip && ip < buffer_ptr + self.exec_buffer_size {
            let offset = ip - buffer_ptr;
            self.trap_data.lookup(offset)
        } else {
            None
        }
    }
}