use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use chrono::{DateTime, Utc};
use entelix_core::{ExecutionContext, Result};
use serde::{Deserialize, Serialize};
use crate::namespace::Namespace;
use crate::store::Store;
const DEFAULT_KEY: &str = "entities";
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EntityRecord {
pub fact: String,
pub last_seen: DateTime<Utc>,
pub created_at: DateTime<Utc>,
}
pub struct EntityMemory {
store: Arc<dyn Store<HashMap<String, EntityRecord>>>,
namespace: Namespace,
}
impl EntityMemory {
pub fn new(store: Arc<dyn Store<HashMap<String, EntityRecord>>>, namespace: Namespace) -> Self {
Self { store, namespace }
}
pub const fn namespace(&self) -> &Namespace {
&self.namespace
}
pub async fn set_entity(
&self,
ctx: &ExecutionContext,
entity: &str,
fact: impl Into<String>,
) -> Result<()> {
let mut all = self
.store
.get(ctx, &self.namespace, DEFAULT_KEY)
.await?
.unwrap_or_default();
let now = Utc::now();
let fact = fact.into();
match all.entry(entity.to_owned()) {
std::collections::hash_map::Entry::Occupied(mut occ) => {
let existing = occ.get_mut();
existing.fact = fact;
existing.last_seen = now;
}
std::collections::hash_map::Entry::Vacant(vac) => {
vac.insert(EntityRecord {
fact,
last_seen: now,
created_at: now,
});
}
}
self.store.put(ctx, &self.namespace, DEFAULT_KEY, all).await
}
pub async fn touch(&self, ctx: &ExecutionContext, entity: &str) -> Result<bool> {
let Some(mut all) = self.store.get(ctx, &self.namespace, DEFAULT_KEY).await? else {
return Ok(false);
};
let Some(record) = all.get_mut(entity) else {
return Ok(false);
};
record.last_seen = Utc::now();
self.store
.put(ctx, &self.namespace, DEFAULT_KEY, all)
.await?;
Ok(true)
}
pub async fn entity(&self, ctx: &ExecutionContext, entity: &str) -> Result<Option<String>> {
Ok(self
.entity_record(ctx, entity)
.await?
.map(|record| record.fact))
}
pub async fn entity_record(
&self,
ctx: &ExecutionContext,
entity: &str,
) -> Result<Option<EntityRecord>> {
Ok(self
.store
.get(ctx, &self.namespace, DEFAULT_KEY)
.await?
.and_then(|all| all.get(entity).cloned()))
}
pub async fn all(&self, ctx: &ExecutionContext) -> Result<HashMap<String, String>> {
let records = self.all_records(ctx).await?;
Ok(records
.into_iter()
.map(|(name, record)| (name, record.fact))
.collect())
}
pub async fn all_records(
&self,
ctx: &ExecutionContext,
) -> Result<HashMap<String, EntityRecord>> {
Ok(self
.store
.get(ctx, &self.namespace, DEFAULT_KEY)
.await?
.unwrap_or_default())
}
pub async fn prune_older_than(&self, ctx: &ExecutionContext, ttl: Duration) -> Result<usize> {
let Some(mut all) = self.store.get(ctx, &self.namespace, DEFAULT_KEY).await? else {
return Ok(0);
};
let cutoff = Utc::now() - chrono::Duration::from_std(ttl).unwrap_or(chrono::Duration::MAX);
let before = all.len();
all.retain(|_, record| record.last_seen >= cutoff);
let removed = before - all.len();
if removed > 0 {
self.store
.put(ctx, &self.namespace, DEFAULT_KEY, all)
.await?;
}
Ok(removed)
}
pub async fn remove(&self, ctx: &ExecutionContext, entity: &str) -> Result<()> {
let Some(mut all) = self.store.get(ctx, &self.namespace, DEFAULT_KEY).await? else {
return Ok(());
};
all.remove(entity);
self.store.put(ctx, &self.namespace, DEFAULT_KEY, all).await
}
pub async fn clear(&self, ctx: &ExecutionContext) -> Result<()> {
self.store.delete(ctx, &self.namespace, DEFAULT_KEY).await
}
}