bevy_mod_reqwest 0.22.1

use std::{
    ops::{Deref, DerefMut},
    sync::Mutex,
};

use bevy::{
    ecs::{
        system::{Commands, IntoObserverSystem, SystemParam},
        world::{EntityWorldMut, World},
    },
    prelude::*,
    tasks::IoTaskPool,
};

pub use reqwest;

#[cfg(target_family = "wasm")]
use crossbeam_channel::{bounded, Receiver};

#[cfg(feature = "json")]
pub use json::*;

pub use reqwest::header::HeaderMap;
pub use reqwest::{StatusCode, Version};

#[cfg(not(target_family = "wasm"))]
use {bevy::tasks::Task, futures_lite::future};

/// The [`SystemSet`] that Reqwest systems are added to.
#[derive(Debug, Hash, PartialEq, Eq, Clone, SystemSet)]
pub struct ReqwestSet;

/// Plugin that allows to send http request using the [reqwest](https://crates.io/crates/reqwest) library from
/// inside bevy.
///
/// The plugin uses [`Observer`] systems to provide callbacks when the http requests finishes.
///
/// Supports both wasm and native.
pub struct ReqwestPlugin {
    /// this enables the plugin to insert a new [`Name`] component onto the entity used to drive
    /// the http request to completion, if no such component already exists
    pub automatically_name_requests: bool,
}
impl Default for ReqwestPlugin {
    fn default() -> Self {
        Self {
            automatically_name_requests: true,
        }
    }
}
impl Plugin for ReqwestPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<ReqwestClient>()
            .init_resource::<PendingReqwestEntityCommands>();

        if self.automatically_name_requests {
            // register a hook on the component to add a name to the entity if it doesnt have one already
            app.world_mut()
                .register_component_hooks::<ReqwestInflight>()
                .on_insert(|mut world, ctx| {
                    let url = world
                        .get::<ReqwestInflight>(ctx.entity)
                        .unwrap()
                        .url
                        .clone();

                    if world.get::<Name>(ctx.entity).is_none() {
                        let mut commands = world.commands();
                        let mut entity = commands.get_entity(ctx.entity).unwrap();
                        entity.insert(Name::new(format!("http: {url}")));
                    }
                });
        }
        //
        app.add_systems(
            PreUpdate,
            (
                // These systems are chained, since the poll_inflight_requests will trigger the callback and mark the entity for deletion

                // So if remove_finished_requests runs after poll_inflight_requests_to_bytes
                // the entity will be removed before the callback is triggered.
                Self::remove_finished_requests,
                Self::poll_inflight_requests_to_bytes,
            )
                .chain()
                .in_set(ReqwestSet),
        )
        .add_systems(PostUpdate, Self::apply_pending_entity_commands);
    }
}

//TODO: Make type generic, and we can create systems for JSON and TEXT requests
impl ReqwestPlugin {
    fn apply_pending_entity_commands(world: &mut World) {
        let Some(mut pending) = world.remove_resource::<PendingReqwestEntityCommands>() else {
            return;
        };

        let mut still_pending = Vec::new();
        for mut command in std::mem::take(pending.0.get_mut().unwrap()) {
            if let Ok(mut entity) = world.get_entity_mut(command.entity) {
                if let Some(action) = command.action.take() {
                    action(&mut entity);
                }
            } else {
                still_pending.push(command);
            }
        }

        world.insert_resource(PendingReqwestEntityCommands(Mutex::new(still_pending)));
    }

    /// despawns finished reqwests if marked to be despawned and does not contain 'ReqwestInflight' component
    fn remove_finished_requests(
        mut commands: Commands,
        q: Query<Entity, (With<DespawnReqwestEntity>, Without<ReqwestInflight>)>,
    ) {
        for e in q.iter() {
            if let Ok(mut ec) = commands.get_entity(e) {
                ec.despawn();
            }
        }
    }

    /// Polls any requests in flight to completion, and then removes the 'ReqwestInflight' component.
    fn poll_inflight_requests_to_bytes(
        mut commands: Commands,
        mut requests: Query<(Entity, &mut ReqwestInflight)>,
    ) {
        for (entity, mut request) in requests.iter_mut() {
            debug!("polling: {entity:?}");
            if let Some((result, parts)) = request.poll() {
                match result {
                    Ok(body) => {
                        // if the response is ok, the other values are already gotten, its safe to unwrap
                        let parts = parts.unwrap();

                        commands.trigger(ReqwestResponseEvent::new(
                            entity,
                            body.clone(),
                            parts.status,
                            parts.headers,
                        ));
                    }
                    Err(err) => {
                        commands.trigger(ReqwestErrorEvent { entity, error: err });
                    }
                }
                if let Ok(mut ec) = commands.get_entity(entity) {
                    ec.remove::<ReqwestInflight>();
                }
            }
        }
    }
}

/// Wrapper around Commands to create the on_response and on_error
pub struct BevyReqwestBuilder<'a> {
    entity: Entity,
    commands: Commands<'a, 'a>,
}

type PendingEntityAction = Box<dyn FnOnce(&mut EntityWorldMut) + Send + 'static>;

#[derive(Resource, Default)]
struct PendingReqwestEntityCommands(Mutex<Vec<PendingReqwestEntityCommand>>);

struct PendingReqwestEntityCommand {
    entity: Entity,
    action: Option<PendingEntityAction>,
}

fn queue_when_entity_exists(
    commands: &mut Commands,
    entity: Entity,
    f: impl FnOnce(&mut EntityWorldMut) + Send + 'static,
) {
    commands.queue(move |world: &mut World| {
        world
            .get_resource_or_insert_with(PendingReqwestEntityCommands::default)
            .0
            .lock()
            .unwrap()
            .push(PendingReqwestEntityCommand {
                entity,
                action: Some(Box::new(f)),
            });
    });
}

impl<'a> BevyReqwestBuilder<'a> {
    /// Provide a system where the first argument is [`Trigger`] [`ReqwestResponseEvent`] that will run on the
    /// response from the http request
    ///
    /// # Examples
    ///
    /// ```
    /// use bevy::prelude::On;
    /// use bevy_mod_reqwest::ReqwestResponseEvent;
    /// |trigger: On<ReqwestResponseEvent>|  {
    ///   bevy::log::info!("response: {:?}", trigger.event());
    /// };
    /// ```
    pub fn on_response<RB: Bundle, RM, OR: IntoObserverSystem<ReqwestResponseEvent, RB, RM>>(
        mut self,
        onresponse: OR,
    ) -> Self {
        queue_when_entity_exists(&mut self.commands, self.entity, |entity| {
            entity.observe(onresponse);
        });
        self
    }

    /// Provide a system where the first argument is [`Trigger`] [`JsonResponse`] that will run on the
    /// response from the http request, skipping some boilerplate of having to manually doing the JSON
    /// parsing
    ///
    /// # Examples
    /// ```
    /// use bevy::prelude::On;
    /// use bevy_mod_reqwest::JsonResponse;
    /// use serde::Deserialize;
    /// #[derive(Deserialize, Debug)]
    /// struct MyResponse;
    /// |trigger: On<JsonResponse<MyResponse>>|  {
    ///   bevy::log::info!("response: {:?}", trigger.event().data);
    /// };
    /// ```
    #[cfg(feature = "json")]
    pub fn on_json_response<
        T: std::marker::Sync + std::marker::Send + serde::de::DeserializeOwned + 'static,
        RB: Bundle,
        RM,
        OR: IntoObserverSystem<json::JsonResponse<T>, RB, RM>,
    >(
        mut self,
        onresponse: OR,
    ) -> Self {
        queue_when_entity_exists(&mut self.commands, self.entity, |entity| {
            entity.observe(|evt: On<ReqwestResponseEvent>, mut commands: Commands| {
                let entity = evt.event().entity;
                let evt = evt.event();
                let data = evt.deserialize_json::<T>();

                match data {
                    Ok(data) => {
                        // retrigger a new event with the serialized data
                        commands.trigger(json::JsonResponse { entity, data });
                    }
                    Err(e) => {
                        bevy::log::error!("deserialization error: {e}");
                        bevy::log::debug!(
                            "tried serializing: {}",
                            evt.as_str().unwrap_or("failed getting event data")
                        );
                    }
                }
            });
        });
        queue_when_entity_exists(&mut self.commands, self.entity, |entity| {
            entity.observe(onresponse);
        });
        self
    }

    /// Provide a system where the first argument is [`Trigger`] [`ReqwestErrorEvent`] that will run on the
    /// response from the http request
    ///
    /// # Examples
    ///
    /// ```
    /// use bevy::prelude::On;
    /// use bevy_mod_reqwest::ReqwestErrorEvent;
    /// |trigger: On<ReqwestErrorEvent>|  {
    ///   bevy::log::info!("response: {:?}", trigger.event());
    /// };
    /// ```
    pub fn on_error<EB: Bundle, EM, OE: IntoObserverSystem<ReqwestErrorEvent, EB, EM>>(
        mut self,
        onerror: OE,
    ) -> Self {
        queue_when_entity_exists(&mut self.commands, self.entity, |entity| {
            entity.observe(onerror);
        });
        self
    }
}

#[derive(SystemParam)]
/// Systemparam to have a shorthand for creating http calls in systems
pub struct BevyReqwest<'w, 's> {
    commands: Commands<'w, 's>,
    client: Res<'w, ReqwestClient>,
}

impl<'w, 's> BevyReqwest<'w, 's> {
    /// Starts sending and processing the supplied [`reqwest::Request`]
    /// then use the [`BevyReqwestBuilder`] to add handlers for responses and errors
    pub fn send(&mut self, req: reqwest::Request) -> BevyReqwestBuilder<'_> {
        let inflight = self.create_inflight_task(req);
        let entity = self.commands.spawn((inflight, DespawnReqwestEntity)).id();
        BevyReqwestBuilder {
            entity,
            commands: self.commands.reborrow(),
        }
    }

    /// Starts sending and processing the supplied [`reqwest::Request`] on the supplied [`Entity`] if it exists
    /// and then use the [`BevyReqwestBuilder`] to add handlers for responses and errors
    pub fn send_using_entity(
        &mut self,
        entity: Entity,
        req: reqwest::Request,
    ) -> Result<BevyReqwestBuilder<'_>, Box<dyn std::error::Error>> {
        let inflight = self.create_inflight_task(req);
        self.commands.get_entity(entity)?;
        info!("inserting request on entity: {:?}", entity);
        queue_when_entity_exists(&mut self.commands, entity, |entity| {
            entity.insert(inflight);
        });
        Ok(BevyReqwestBuilder {
            entity,
            commands: self.commands.reborrow(),
        })
    }

    /// get access to the underlying ReqwestClient
    pub fn client(&self) -> &reqwest::Client {
        &self.client.0
    }

    fn create_inflight_task(&self, request: reqwest::Request) -> ReqwestInflight {
        let thread_pool = IoTaskPool::get();
        // bevy::log::debug!("Creating: {entity:?}");
        // if we take the data, we can use it
        let client = self.client.0.clone();
        let url = request.url().to_string();

        // wasm implementation
        #[cfg(target_family = "wasm")]
        let task = {
            let (tx, task) = bounded(1);
            thread_pool
                .spawn(async move {
                    let r = client.execute(request).await;
                    let r = match r {
                        Ok(res) => {
                            let parts = Parts {
                                status: res.status(),
                                headers: res.headers().clone(),
                            };
                            (res.bytes().await, Some(parts))
                        }
                        Err(r) => (Err(r), None),
                    };
                    tx.send(r).ok();
                })
                .detach();
            task
        };

        // otherwise
        #[cfg(all(not(test), not(target_family = "wasm")))]
        let task = {
            thread_pool.spawn(async move {
                let task_res = async_compat::Compat::new(async {
                    let p = client.execute(request).await;
                    match p {
                        Ok(res) => {
                            let parts = Parts {
                                status: res.status(),
                                headers: res.headers().clone(),
                            };
                            (res.bytes().await, Some(parts))
                        }
                        Err(e) => (Err(e), None),
                    }
                })
                .await;
                task_res
            })
        };

        #[cfg(all(test, not(target_family = "wasm")))]
        let task = {
            let _ = request;
            let _ = client;
            thread_pool.spawn(async { future::pending::<Resp>().await })
        };
        // put it as a component to be polled, and remove the request, it has been handled
        ReqwestInflight::new(task, url)
    }
}

impl<'w, 's> Deref for BevyReqwest<'w, 's> {
    type Target = reqwest::Client;

    fn deref(&self) -> &Self::Target {
        self.client()
    }
}

#[derive(Component)]
/// Marker component that is used to despawn an entity if the reqwest is finshed
pub struct DespawnReqwestEntity;

#[derive(Resource)]
/// Wrapper around the ReqwestClient, that when inserted as a resource will start connection pools towards
/// the hosts, and also allows all the configuration from the ReqwestLibrary such as setting default headers etc
/// to be used inside the bevy application
pub struct ReqwestClient(pub reqwest::Client);
impl Default for ReqwestClient {
    fn default() -> Self {
        Self(reqwest::Client::new())
    }
}

impl std::ops::Deref for ReqwestClient {
    type Target = reqwest::Client;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
impl DerefMut for ReqwestClient {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

type Resp = (reqwest::Result<bytes::Bytes>, Option<Parts>);

/// Dont touch these, its just to poll once every request, can be used to detect if there is an active request on the entity
/// but should otherwise NOT be added/removed/changed by a user of this Crate
#[derive(Component)]
#[component(storage = "SparseSet")]
pub struct ReqwestInflight {
    // the url this request is handling as a string
    pub(crate) url: String,
    #[cfg(not(target_family = "wasm"))]
    res: Task<Resp>,

    #[cfg(target_family = "wasm")]
    res: Receiver<Resp>,
}

impl ReqwestInflight {
    fn poll(&mut self) -> Option<Resp> {
        #[cfg(target_family = "wasm")]
        {
            self.res.try_recv().ok()
        }

        #[cfg(not(target_family = "wasm"))]
        {
            future::block_on(future::poll_once(&mut self.res))
        }
    }

    #[cfg(target_family = "wasm")]
    pub(crate) fn new(res: Receiver<Resp>, url: String) -> Self {
        Self { url, res }
    }

    #[cfg(not(target_family = "wasm"))]
    pub(crate) fn new(res: Task<Resp>, url: String) -> Self {
        Self { url, res }
    }
}

#[derive(Component, Debug)]
/// information about the response used to transfer headers between different stages in the async code
struct Parts {
    /// the `StatusCode`
    pub(crate) status: StatusCode,

    /// the headers of the response
    pub(crate) headers: HeaderMap,
}

#[derive(Clone, EntityEvent, Debug)]
/// the resulting data from a finished request is found here
pub struct ReqwestResponseEvent {
    entity: Entity,
    bytes: bytes::Bytes,
    status: StatusCode,
    headers: HeaderMap,
}

#[derive(EntityEvent, Debug)]
pub struct ReqwestErrorEvent {
    pub entity: Entity,
    pub error: reqwest::Error,
}

impl ReqwestResponseEvent {
    /// retrieve a reference to the body of the response
    #[inline]
    pub fn body(&self) -> &bytes::Bytes {
        &self.bytes
    }

    /// try to get the body of the response as_str
    pub fn as_str(&self) -> anyhow::Result<&str> {
        let s = std::str::from_utf8(&self.bytes)?;
        Ok(s)
    }
    /// try to get the body of the response as an owned string
    pub fn as_string(&self) -> anyhow::Result<String> {
        Ok(self.as_str()?.to_string())
    }
    #[cfg(feature = "json")]
    /// try to deserialize the body of the response using json
    pub fn deserialize_json<'de, T: serde::Deserialize<'de>>(&'de self) -> anyhow::Result<T> {
        Ok(serde_json::from_str(self.as_str()?)?)
    }

    #[cfg(feature = "msgpack")]
    /// try to deserialize the body of the response using msgpack
    pub fn deserialize_msgpack<'de, T: serde::Deserialize<'de>>(&'de self) -> anyhow::Result<T> {
        Ok(rmp_serde::decode::from_slice(self.body())?)
    }
    #[inline]
    /// Get the `StatusCode` of this `Response`.
    pub fn status(&self) -> StatusCode {
        self.status
    }

    #[inline]
    /// Get the `Headers` of this `Response`.
    pub fn response_headers(&self) -> &HeaderMap {
        &self.headers
    }
}

#[cfg(feature = "json")]
pub mod json {
    use bevy::{ecs::entity::Entity, prelude::EntityEvent};
    #[derive(EntityEvent)]
    pub struct JsonResponse<T> {
        pub entity: Entity,
        pub data: T,
    }
}

impl ReqwestResponseEvent {
    pub(crate) fn new(
        entity: Entity,
        bytes: bytes::Bytes,
        status: StatusCode,
        headers: HeaderMap,
    ) -> Self {
        Self {
            entity,
            bytes,
            status,
            headers,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn app_with_send_system<M>(system: impl IntoSystem<(), (), M> + 'static) -> App {
        let mut app = App::new();
        app.add_plugins((TaskPoolPlugin::default(), ReqwestPlugin::default()))
            .add_systems(Update, system);
        app
    }

    fn build_request(client: &BevyReqwest) -> reqwest::Request {
        client.get("http://127.0.0.1:9/").build().unwrap()
    }

    fn send_with_bevy_reqwest_before_commands(mut client: BevyReqwest, mut commands: Commands) {
        let entity = commands.spawn(DespawnReqwestEntity).id();
        let request = build_request(&client);

        client.send_using_entity(entity, request).unwrap();
    }

    fn send_with_commands_before_bevy_reqwest(mut commands: Commands, mut client: BevyReqwest) {
        let entity = commands.spawn(DespawnReqwestEntity).id();
        let request = build_request(&client);

        client.send_using_entity(entity, request).unwrap();
    }

    #[test]
    fn sending_on_an_entity_spawned_by_commands_does_not_depend_on_parameter_order() {
        let mut app = app_with_send_system(send_with_bevy_reqwest_before_commands);

        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| app.update()));
        std::mem::forget(app);

        assert!(
            result.is_ok(),
            "sending a request using an entity spawned via Commands should not depend on system parameter order"
        );
    }

    #[test]
    fn sending_on_an_entity_spawned_by_commands_does_not_panic_when_commands_param_is_first() {
        let mut app = app_with_send_system(send_with_commands_before_bevy_reqwest);

        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| app.update()));
        std::mem::forget(app);

        assert!(
            result.is_ok(),
            "Commands first is the currently working ordering and should not panic"
        );
    }
}