use crate::{arch::generate_process_pass, rhai_shell::KEYBUFF};
use wasmi::{
    Error, Externals, FuncInstance, FuncRef, ModuleImportResolver, RuntimeArgs, RuntimeValue,
    Signature, Trap, ValueType,
};

const ADD_FUNC_INDEX: usize = 0;
const GET_TIME_INDEX: usize = 2;
const GET_RANDOM_INDEX: usize = 3;
const GET_INPUT_INDEX: usize = 4;
const PRINT_CHAR: usize = 6;
const SEND_SIGNAL_INDEX: usize = 1;

pub struct HostExternals;
impl HostExternals {
    fn check_signature(&self, index: usize, signature: &Signature) -> bool {
        let (params, ret_ty): (&[_], _) = match index {
            ADD_FUNC_INDEX => (&[ValueType::I32, ValueType::I32], Some(ValueType::I32)),
            SEND_SIGNAL_INDEX => (&[ValueType::I32, ValueType::I32], Some(ValueType::I32)),
            GET_TIME_INDEX => (&[], Some(ValueType::I32)),
            GET_RANDOM_INDEX => (&[], Some(ValueType::I32)),
            GET_INPUT_INDEX => (&[], Some(ValueType::I32)),
            PRINT_CHAR => (&[ValueType::I32], None),
            _ => return false,
        };

        if params.len() != signature.params().len() || ret_ty != signature.return_type() {
            false
        } else {
            params
                .iter()
                .zip(signature.params())
                .find(|(ty, param)| ty != param)
                .is_none()
        }
    }
}

impl Externals for HostExternals {
    fn invoke_index(
        &mut self,
        index: usize,
        args: RuntimeArgs,
    ) -> Result<Option<RuntimeValue>, Trap> {
        match index {
            ADD_FUNC_INDEX => {
                let a: u32 = args.nth_checked(0)?;
                let b: u32 = args.nth_checked(1)?;
                let result = a + b;
                trace!("SYSCALL: {} + {} = {}", a, b, result);
                Ok(Some(RuntimeValue::I32(result as i32)))
            }

            SEND_SIGNAL_INDEX => {
                let pid: u32 = args.nth_checked(0)?;
                let signal: u32 = args.nth_checked(1)?;

                trace!("SYSCALL: send signal {} to pid {}", signal, pid);
                let ret = RuntimeValue::I32(0);
                Ok(Some(ret))
            }

            GET_TIME_INDEX => {
                use core::sync::atomic::Ordering::*;
                trace!("SYSCALL: get time");

                x86_64::instructions::interrupts::disable();
                let tick_time = kernel::TICK.load(Relaxed);
                x86_64::instructions::interrupts::enable();

                let ret = RuntimeValue::I32(tick_time.try_into().unwrap());
                Ok(Some(ret))
            }

            GET_RANDOM_INDEX => {
                trace!("SYSCALL: get random");
                let rand = generate_process_pass();

                let ret = RuntimeValue::I32(rand as i32);

                // let ret = RuntimeValue::I32(rand.try_into().unwrap());
                Ok(Some(ret))
            }

            GET_INPUT_INDEX => {
                let input = None;
                x86_64::instructions::interrupts::without_interrupts(|| KEYBUFF.lock().pop());
                if let Some(chr) = input {
                    trace!("SYSCALL: input: {}", chr);
                }

                let ret = RuntimeValue::I32(input.unwrap_or(0x00 as char) as i32);

                Ok(Some(ret))
            }

            PRINT_CHAR => {
                let chr: u8 = args.nth_checked(0)?;
                trace!("SYSCALL: print: {}", chr);
                print!("{}", char::from(chr));
                Ok(None)
            }
            _ => {
                error!("Unimplemented function at {}", index);
                Err(Trap::new(wasmi::TrapKind::Unreachable))
            }
        }
    }
}

impl ModuleImportResolver for HostExternals {
    fn resolve_func(&self, field_name: &str, signature: &Signature) -> Result<FuncRef, Error> {
        let index = match field_name {
            "add" => ADD_FUNC_INDEX,
            "send_signal" => SEND_SIGNAL_INDEX,
            "get_time" => GET_TIME_INDEX,
            "get_random" => GET_RANDOM_INDEX,
            "get_input" => GET_INPUT_INDEX,
            "print_char" => PRINT_CHAR,
            _ => {
                return Err(Error::Instantiation(format!(
                    "Export {} not found",
                    field_name
                )))
            }
        };

        if !self.check_signature(index, signature) {
            return Err(Error::Instantiation(format!(
                "Export {} has a bad signature {:?}",
                field_name, signature
            )));
        }

        trace!("Resolved export {} as func {}", field_name, index);
        Ok(FuncInstance::alloc_host(signature.clone(), index))
    }
}