bevy_persistence_database 0.3.0

use std::marker::PhantomData;

use bevy::ecs::system::SystemParam;
use bevy::prelude::{Mut, Resource, World};

use crate::bevy::plugins::persistence_plugin::{PersistencePluginConfig, TokioRuntime};
use crate::bevy::world_access::ImmediateWorldPtr;
use crate::core::db::connection::DatabaseConnectionResource;
use crate::core::persist::Persist;
use crate::core::session::PersistenceSession;
use crate::core::versioning::version_manager::VersionKey;
use super::resource_thread_local::{
    set_resource_force_refresh, set_resource_store, take_resource_force_refresh,
    take_resource_store,
};
use serde_json::Value;

/// Read-only persisted resource accessor. Hydrates on first use if absent.
#[derive(SystemParam)]
pub struct PersistentRes<'w, T: Resource + Persist> {
    session: bevy::prelude::ResMut<'w, PersistenceSession>,
    db: bevy::prelude::Res<'w, DatabaseConnectionResource>,
    config: bevy::prelude::Res<'w, PersistencePluginConfig>,
    runtime: bevy::prelude::Res<'w, TokioRuntime>,
    world_ptr: Option<bevy::prelude::Res<'w, ImmediateWorldPtr>>,
    _marker: PhantomData<T>,
}

/// Mutable persisted resource accessor. Marks dirty on first mutation.
#[derive(SystemParam)]
pub struct PersistentResMut<'w, T: Resource + Persist> {
    session: bevy::prelude::ResMut<'w, PersistenceSession>,
    db: bevy::prelude::Res<'w, DatabaseConnectionResource>,
    config: bevy::prelude::Res<'w, PersistencePluginConfig>,
    runtime: bevy::prelude::Res<'w, TokioRuntime>,
    world_ptr: Option<bevy::prelude::Res<'w, ImmediateWorldPtr>>,
    _marker: PhantomData<T>,
}

impl<'w, T: Resource + Persist> PersistentRes<'w, T> {
    /// Override the store for the next load of this resource.
    pub fn store(self, store: impl Into<String>) -> Self {
        set_resource_store(store);
        self
    }

    /// Force a refresh from persistence even if the resource already exists in the world.
    pub fn force_refresh(self) -> Self {
        set_resource_force_refresh();
        self
    }

    /// Returns an immutable reference, hydrating from persistence if necessary.
    pub fn get(&mut self) -> Option<&T> {
        self.ensure_loaded();
        self.resource_ref()
    }

    fn ensure_loaded(&mut self) {
        let force_refresh = take_resource_force_refresh();

        if !force_refresh && self.resource_ref().is_some() {
            return;
        }

        let Some(world_ptr) = self.world_ptr.as_ref() else {
            bevy::log::error!("PersistentRes missing ImmediateWorldPtr; cannot hydrate resource");
            return;
        };

        let store = take_resource_store().unwrap_or_else(|| self.config.default_store.clone());

        let result = self
            .runtime
            .block_on(async { self.db.connection.fetch_resource(&store, T::name()).await });

        match result {
            Ok(Some((value, version))) => {
                if let Err(err) = self.deserialize_into_world(world_ptr.ptr, value) {
                    bevy::log::error!(
                        "failed to deserialize persisted resource {}: {}",
                        T::name(),
                        err
                    );
                    return;
                }
                if let Some(type_id) = self.session.resource_type_id(T::name()) {
                    self.session
                        .version_manager_mut()
                        .set_version(VersionKey::Resource(type_id), version);
                }
            }
            Ok(None) => {}
            Err(err) => {
                bevy::log::error!(%err, store, "failed to load persisted resource {}", T::name());
            }
        }
    }

    fn deserialize_into_world(
        &mut self,
        world_ptr: *mut World,
        value: Value,
    ) -> Result<(), String> {
        let world: &mut World = unsafe { &mut *world_ptr };
        self.session
            .deserialize_resource_by_name(world, T::name(), value)
            .map_err(|e| e.to_string())
    }

    fn resource_ref(&self) -> Option<&T> {
        let ptr = self.world_ptr.as_ref()?;
        let world: &World = unsafe { &*ptr.ptr.cast_const() };
        world.get_resource::<T>()
    }
}

impl<'w, T: Resource + Persist> PersistentResMut<'w, T> {
    /// Override the store for the next load of this resource.
    pub fn store(self, store: impl Into<String>) -> Self {
        set_resource_store(store);
        self
    }

    /// Force a refresh from persistence even if the resource already exists in the world.
    pub fn force_refresh(self) -> Self {
        set_resource_force_refresh();
        self
    }

    /// Returns a mutable reference, hydrating from persistence if necessary.
    /// Marks the resource dirty on first mutation.
    pub fn get_mut(&mut self) -> Option<Mut<'_, T>> {
        self.ensure_loaded();
        self.session.mark_resource_dirty::<T>();
        let Some(value) = self.resource_mut() else {
            return None;
        };
        Some(value)
    }

    fn ensure_loaded(&mut self) {
        let force_refresh = take_resource_force_refresh();

        if !force_refresh && self.resource_ref().is_some() {
            return;
        }

        let Some(world_ptr) = self.world_ptr.as_ref() else {
            bevy::log::error!(
                "PersistentResMut missing ImmediateWorldPtr; cannot hydrate resource"
            );
            return;
        };

        let store = take_resource_store().unwrap_or_else(|| self.config.default_store.clone());

        let result = self
            .runtime
            .block_on(async { self.db.connection.fetch_resource(&store, T::name()).await });

        match result {
            Ok(Some((value, version))) => {
                if let Err(err) = self.deserialize_into_world(world_ptr.ptr, value) {
                    bevy::log::error!(
                        "failed to deserialize persisted resource {}: {}",
                        T::name(),
                        err
                    );
                    return;
                }
                if let Some(type_id) = self.session.resource_type_id(T::name()) {
                    self.session
                        .version_manager_mut()
                        .set_version(VersionKey::Resource(type_id), version);
                }
            }
            Ok(None) => {}
            Err(err) => {
                bevy::log::error!(%err, store, "failed to load persisted resource {}", T::name());
            }
        }
    }

    fn deserialize_into_world(
        &mut self,
        world_ptr: *mut World,
        value: Value,
    ) -> Result<(), String> {
        let world: &mut World = unsafe { &mut *world_ptr };
        self.session
            .deserialize_resource_by_name(world, T::name(), value)
            .map_err(|e| e.to_string())
    }

    fn resource_ref(&self) -> Option<&T> {
        let ptr = self.world_ptr.as_ref()?;
        let world: &World = unsafe { &*ptr.ptr.cast_const() };
        world.get_resource::<T>()
    }

    fn resource_mut(&self) -> Option<Mut<'_, T>> {
        let ptr = self.world_ptr.as_ref()?;
        let world: &mut World = unsafe { &mut *ptr.ptr };
        world.get_resource_mut::<T>()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::db::connection::{DatabaseConnectionResource, MockDatabaseConnection};
    use crate::bevy::plugins::persistence_plugin::TokioRuntime;
    use bevy::prelude::App;
    use bevy::prelude::MinimalPlugins;
    use bevy_persistence_database_derive::persist;
    use serde_json::json;
    use std::sync::Arc;

    #[derive(Clone)]
    #[persist(resource)]
    struct TestResource {
        value: String,
    }

    fn runtime() -> TokioRuntime {
        TokioRuntime {
            runtime: Arc::new(tokio::runtime::Runtime::new().expect("runtime")),
        }
    }

    fn default_config() -> PersistencePluginConfig {
        PersistencePluginConfig {
            default_store: "store".to_string(),
            ..Default::default()
        }
    }

    fn init_world_ptr(app: &mut App) {
        let ptr = app.world_mut() as *mut World;
        app.insert_resource(ImmediateWorldPtr::new(ptr));
    }

    #[test]
    fn loads_resource_on_first_access() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);

        let mut db = MockDatabaseConnection::new();
        db.expect_fetch_resource().returning(|store, name| {
            assert_eq!(store, "store");
            assert_eq!(name, TestResource::name());
            Box::pin(async { Ok(Some((json!({ "value": "persisted" }), 3))) })
        });
        db.expect_document_key_field().return_const("_key");

        app.insert_resource(DatabaseConnectionResource {
            connection: Arc::new(db),
        });
        app.insert_resource(default_config());
        app.insert_resource(runtime());

        let mut session = PersistenceSession::new();
        session.register_resource::<TestResource>();
        app.insert_resource(session);
        init_world_ptr(&mut app);

        app.add_systems(
            bevy::prelude::Update,
            |mut res: PersistentRes<TestResource>| {
                let r = res.get().expect("resource should load");
                assert_eq!(r.value, "persisted");
            },
        );

        app.update();
    }

    #[test]
    fn respects_existing_resource_and_skips_fetch() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);

        let mut db = MockDatabaseConnection::new();
        db.expect_fetch_resource()
            .returning(|_, _| Box::pin(async { unreachable!() }));
        db.expect_document_key_field().return_const("_key");

        app.insert_resource(DatabaseConnectionResource {
            connection: Arc::new(db),
        });
        app.insert_resource(default_config());
        app.insert_resource(runtime());

        let mut session = PersistenceSession::new();
        session.register_resource::<TestResource>();
        app.insert_resource(session);
        init_world_ptr(&mut app);

        app.insert_resource(TestResource {
            value: "existing".to_string(),
        });

        app.add_systems(
            bevy::prelude::Update,
            |mut res: PersistentRes<TestResource>| {
                let r = res.get().expect("resource should already exist");
                assert_eq!(r.value, "existing");
            },
        );

        app.update();
    }

    #[test]
    fn marks_dirty_on_mutation() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);

        let mut db = MockDatabaseConnection::new();
        db.expect_fetch_resource()
            .returning(|_, _| Box::pin(async { Ok(Some((json!({ "value": "persisted" }), 1))) }));
        db.expect_document_key_field().return_const("_key");

        app.insert_resource(DatabaseConnectionResource {
            connection: Arc::new(db),
        });
        app.insert_resource(default_config());
        app.insert_resource(runtime());

        let mut session = PersistenceSession::new();
        session.register_resource::<TestResource>();
        app.insert_resource(session);
        init_world_ptr(&mut app);

        app.add_systems(
            bevy::prelude::Update,
            |mut res: PersistentResMut<TestResource>| {
                let mut r = res.get_mut().expect("resource should load");
                r.value = "changed".to_string();
            },
        );

        app.update();

        let session = app.world().get_resource::<PersistenceSession>().unwrap();
        assert!(
            session.is_resource_dirty(std::any::TypeId::of::<TestResource>())
        );
    }

    #[test]
    fn store_override_is_used() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);

        let mut db = MockDatabaseConnection::new();
        db.expect_fetch_resource().returning(|store, _| {
            assert_eq!(store, "alt-store");
            Box::pin(async { Ok(None) })
        });
        db.expect_document_key_field().return_const("_key");

        app.insert_resource(DatabaseConnectionResource {
            connection: Arc::new(db),
        });
        app.insert_resource(default_config());
        app.insert_resource(runtime());

        let mut session = PersistenceSession::new();
        session.register_resource::<TestResource>();
        app.insert_resource(session);
        init_world_ptr(&mut app);

        app.add_systems(bevy::prelude::Update, |res: PersistentRes<TestResource>| {
            let mut res = res.store("alt-store");
            let _ = res.get();
        });

        app.update();
    }

    #[test]
    fn logs_and_skips_on_deserialize_error() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);

        let mut db = MockDatabaseConnection::new();
        db.expect_fetch_resource().returning(|_, _| {
            // Return an invalid shape for TestResource to force a deserialize error
            Box::pin(async { Ok(Some((json!({ "wrong": 1 }), 2))) })
        });
        db.expect_document_key_field().return_const("_key");

        app.insert_resource(DatabaseConnectionResource {
            connection: Arc::new(db),
        });
        app.insert_resource(default_config());
        app.insert_resource(runtime());

        let mut session = PersistenceSession::new();
        session.register_resource::<TestResource>();
        app.insert_resource(session);
        init_world_ptr(&mut app);

        app.add_systems(bevy::prelude::Update, |mut res: PersistentRes<TestResource>| {
            let _ = res.get();
        });

        app.update();

        // Resource should not be inserted when deserialization fails
        assert!(app.world().get_resource::<TestResource>().is_none());
    }
}