//! 🦹 Behavior module
//!
//! Behavior instances are identified through a densely packed identifier called `BehaviorInstanceId`.
//!
//! `BehaviorInstanceId` uniquely identifies a triple of:
//!
//! * An external behavior module handle
//! * A behavior type in this behavior module
//! * A behavior instance of this behavior type in this module
//!
//! Different users of this identifier need different parts:
//!
//! * `BehaviorModuleHandle`, while internal to Ark, consists of the high 16 bits of this id,
//!   and can be used to index into an array of behavior modules
//! * `BehaviorTypeId` consists of the module handle as well as the local type id, that is, the
//!   high 32-bits. It is used to create new behavior instances, through explicit creation or deserialization.
//! * `ForeignBehaviorInstanceId` consists of the parts that are local to a behavior module. Ark will complete
//!   it together with a `BehaviorModuleHandle` into a `BehaviorInstanceId`.
//!
//!  /---------- `ForeignBehaviorInstanceId` ------------\
//! /                                                     \
//!   `LocalBehaviorInstanceId`     `LocalBehaviorTypeId`      `BehaviorModuleHandle`
//! 0.........................31 32......................47 48.........................63
//!                              \                                                     /
//!                               \--------------- `BehaviorTypeId` ------------------/
//!  \                                                                               /
//!   \--------------------------- `BehaviorInstanceId` ----------------------------/

pub use crate::ErrorCode;
pub use ark_api::behavior::ForeignBehaviorInstanceId;
pub use ark_api::behavior::IncomingMessage;
pub use ark_api::behavior::IncomingMessageAddr;
pub use ark_api::behavior::LocalBehaviorInstanceId;
pub use ark_api::behavior::LocalBehaviorTypeId;
pub use ark_api_ffi::behavior_v0::ActorId;
pub use ark_api_ffi::behavior_v0::IncomingMessage as IncomingMessageRaw;
pub use ark_api_ffi::behavior_v0::LocalBehaviorRegistration;
pub use ark_api_ffi::behavior_v0::LocalModuleRegistration;

/// Ergonomic wrapper around a behavior module's entry points.
///
/// See the module level documentation for an explanation on the various identifiers
/// used in this trait
pub trait BehaviorEntry {
    /// Should return a serialized JSON vector of bytes of local registration info.
    ///
    /// # Example
    ///
    /// ```rust
    /// use ark_module::behavior::{LocalBehaviorRegistration, LocalModuleRegistration};
    ///
    /// fn register(&self) -> Vec<u8> {
    ///     let behavior_info = LocalBehaviorRegistration{
    ///         type_id: 0,
    ///         name: String::from("test123"),
    ///     };
    ///
    ///     let registration_info = LocalModuleRegistration{
    ///         behavior_infos: vec![behavior_info],
    ///     };
    ///
    ///     serde_json::to_vec(&registration_info).unwrap()
    /// }
    /// ```
    /// Users can return custom registration data by including `LocalModuleRegistration`
    /// and tagging it with `#[serde(flatten)]`.
    ///
    /// ```rust
    /// use ark_module::behavior::{LocalBehaviorRegistration, LocalModuleRegistration};
    ///
    /// #[derive(serde::Deserialize)]
    /// pub struct CustomRegistrationInfo {
    ///     #[serde(flatten)]
    ///     pub registration_info: LocalModuleRegistration,
    ///     pub custom_string: String,
    ///     pub custom_uint: u32,
    /// }
    ///
    /// fn register(&self) -> Vec<u8> {
    ///     let behavior_info = LocalBehaviorRegistration{
    ///         type_id: 0,
    ///         name: String::from("test123"),
    ///     };
    ///
    ///     let registration_info = LocalModuleRegistration{
    ///         behavior_infos: vec![behavior_info],
    ///     };
    ///
    ///     let custom_registration_info = CustomRegistrationInfo {
    ///         registration_info,
    ///         custom_string: String::from("test456"),
    ///         custom_uint: 1,
    ///     };
    /// }
    /// ```
    fn register(&mut self) -> Vec<u8>;

    /// Persists a whole module after it failed, before an attempt to reload it and re-run it.
    /// In particular, the module should keep an account of all the cleanups it should run after
    /// restoring is done.
    fn persist(&mut self) -> Vec<u8>;

    /// Restores a module serialized with `persist`. The module should run additional cleanups at
    /// this point.
    fn restore(&mut self, state: &[u8]);

    /// Construct a new [`LocalBehaviorInstanceId`] given the type id and ron parameters,
    /// which can then be packed into a [`ForeignBehaviorInstanceId`] and returned here.
    fn instance_create(
        &mut self,
        local_type_id: LocalBehaviorTypeId,
        ron_params: &str,
    ) -> ForeignBehaviorInstanceId;

    /// Destroys the provided instance in the behavior module
    fn instance_destroy(&mut self, instance_id: ForeignBehaviorInstanceId);

    /// Clones the provided instance in the behavior module
    fn instance_clone(
        &mut self,
        instance_id: ForeignBehaviorInstanceId,
    ) -> ForeignBehaviorInstanceId;

    /// Serialize the instance's state into a binary blob
    fn instance_persist(&self, instance_id: ForeignBehaviorInstanceId) -> Vec<u8>;

    /// Restore an instance given the provided type id and serialized state
    fn instance_restore(
        &mut self,
        local_type_id: LocalBehaviorTypeId,
        blob: &[u8],
    ) -> ForeignBehaviorInstanceId;

    /// Have the provided instance handle the incoming serialized message
    fn instance_handle_message(&mut self, addr: IncomingMessageAddr, serialized_message: &[u8]);

    /// Process the incoming messages
    fn instances_handle_messages<'a>(
        &mut self,
        messages: impl ExactSizeIterator<Item = IncomingMessage<'a>>,
    ) {
        for msg in messages {
            self.instance_handle_message(msg.addr, msg.serialized_message);
        }
    }
}

/// An iterator over a collection of [`IncomingMessage`]s encoded in a byte buffer
#[doc(hidden)]
pub struct IncomingMessageIter<'a> {
    bytes: &'a [u8],
    offset: usize,
    len: usize,
}

impl<'a> IncomingMessageIter<'a> {
    #[doc(hidden)]
    pub fn new(bytes: &'a [u8]) -> Result<Self, ErrorCode> {
        let mut offset = 0;
        let len = Self::read_u32_len(bytes, &mut offset).ok_or(ErrorCode::InternalError)?;

        Ok(Self { bytes, offset, len })
    }

    /// Read a u32 value from our byte stream, but return as usize as we use it for length
    fn read_u32_len(bytes: &[u8], offset: &mut usize) -> Option<usize> {
        let result = bytes.get((*offset)..(*offset + 4)).map(|slice| {
            *offset += 4;
            u32::from_le_bytes(slice.try_into().unwrap()) as usize
        });

        result
    }

    fn read_incoming_msg_ffi(bytes: &[u8], offset: &mut usize) -> Option<IncomingMessageRaw> {
        const INCOMING_MESSAGE_SIZE: usize = std::mem::size_of::<IncomingMessageRaw>();

        let result = bytes
            .get((*offset)..(*offset + INCOMING_MESSAGE_SIZE))
            .map(|slice| {
                *offset += INCOMING_MESSAGE_SIZE;
                // SAFETY: Safe due to the layout enforcement of `repr(C)` and being `Copy`
                unsafe {
                    std::mem::transmute::<[u8; INCOMING_MESSAGE_SIZE], IncomingMessageRaw>(
                        slice.try_into().unwrap(),
                    )
                }
            });

        result
    }

    fn read_incoming_message(bytes: &'a [u8], offset: &mut usize) -> Option<IncomingMessage<'a>> {
        Self::read_incoming_msg_ffi(bytes, offset).map(|incoming| {
            *offset += incoming.serialized_message_len as usize;
            IncomingMessage::from(incoming)
        })
    }
}

impl<'a> Iterator for IncomingMessageIter<'a> {
    type Item = IncomingMessage<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        IncomingMessageIter::read_incoming_message(self.bytes, &mut self.offset)
    }
}

impl<'a> ExactSizeIterator for IncomingMessageIter<'a> {
    fn len(&self) -> usize {
        self.len
    }
}

/// By passing a type that implements the [`BehaviorEntry`] trait, this macro will hook up
/// and route the behavior module's raw entry points to the trait methods.
#[macro_export]
macro_rules! impl_behavior_module {
    ($module:ty) => {
        use $crate::behavior::BehaviorEntry as _;

        static mut MODULE: Option<$module> = None;

        #[no_mangle]
        pub fn ark_behavior_register() {
            $crate::init();

            let mut module = <$module>::default();
            let behavior_infos = module.register();

            // SAFETY: Sound to access as are no threads in Wasm and this is the first behavior function that will be called in the module
            let module = unsafe { MODULE.replace(module) };

            ark::core::return_slice(&behavior_infos);
        }

        #[no_mangle]
        pub fn ark_behavior_persist() {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            let state = module.persist();
            ark::core::return_slice(&state);
        }

        #[no_mangle]
        pub unsafe fn ark_behavior_restore(state_ptr: *const u8, state_size: u32) {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            // SAFETY: Usage is kept within lifetime of function call only
            let state_slice = unsafe { $crate::util::param_byte_slice(state_ptr, state_size) };
            module.restore(&state_slice);
        }

        #[no_mangle]
        pub unsafe fn ark_behavior_instance_create(
            local_type_id: u32,
            params_ptr: *const u8,
            params_len: u32,
        ) -> u64 {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            // SAFETY: Usage is kept within lifetime of function call only
            let params_slice = unsafe { $crate::util::param_byte_slice(params_ptr, params_len) };
            // TODO: should we require UTF-8 strings for the parameters, this constraint is not described in the entrypoint definition
            let params_str = std::str::from_utf8(params_slice).unwrap();

            assert!(local_type_id <= u16::max_value() as u32);
            let local_type_id = $crate::behavior::LocalBehaviorTypeId(local_type_id as u16);
            module.instance_create(local_type_id, params_str).0
        }

        #[no_mangle]
        pub fn ark_behavior_instance_destroy(instance_id: u64) {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            module.instance_destroy($crate::behavior::ForeignBehaviorInstanceId(instance_id));
        }

        #[no_mangle]
        pub fn ark_behavior_instance_clone(instance_id: u64) -> u64 {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            module
                .instance_clone($crate::behavior::ForeignBehaviorInstanceId(instance_id))
                .0
        }

        #[no_mangle]
        pub fn ark_behavior_instance_persist(instance_id: u64) {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            let bytes =
                module.instance_persist($crate::behavior::ForeignBehaviorInstanceId(instance_id));
            ark::core::return_slice(&bytes);
        }

        #[no_mangle]
        pub unsafe fn ark_behavior_instance_restore(
            local_type_id: u32,
            blob_ptr: *const u8,
            blob_size: u32,
        ) -> u64 {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            // SAFETY: Usage is kept within lifetime of function call only
            let blob_slice = unsafe { $crate::util::param_byte_slice(blob_ptr, blob_size) };
            assert!(local_type_id <= u16::max_value() as u32);
            let local_type_id = $crate::behavior::LocalBehaviorTypeId(local_type_id as u16);

            module.instance_restore(local_type_id, blob_slice).0
        }

        #[no_mangle]
        pub unsafe fn ark_behavior_instance_handle_message(
            instance_id: u64,
            actor_id: $crate::behavior::ActorId,
            message_blob_ptr: *const u8,
            message_blob_size: u32,
        ) {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            // SAFETY: Usage is kept within lifetime of function call only
            let message_blob_slice =
                unsafe { $crate::util::param_byte_slice(message_blob_ptr, message_blob_size) };

            let instance_id = $crate::behavior::ForeignBehaviorInstanceId(instance_id);

            let addr = $crate::behavior::IncomingMessageAddr::from_raw(actor_id, instance_id);

            module.instance_handle_message(addr, message_blob_slice);
        }

        #[no_mangle]
        pub unsafe fn ark_behavior_instances_handle_messages(
            encoded_ptr: *const u8,
            encoded_len: u32,
        ) {
            // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
            let module = unsafe { MODULE.as_mut().unwrap() };
            // SAFETY: Usage is kept within lifetime of function call only
            let encoded_slice = unsafe { $crate::util::param_byte_slice(encoded_ptr, encoded_len) };
            if let Ok(iter) = $crate::behavior::IncomingMessageIter::new(encoded_slice) {
                module.instances_handle_messages(iter);
            } else {
                $crate::error!("Invalid encoded `IncomingMessage` buffer");
            }
        }
    };
}