use crate::colony::{Colony, ColonyConfig};
use crate::topology_impl::PetTopologyGraph;
use phago_core::topology::TopologyGraph;
use phago_core::types::*;
use std::path::{Path, PathBuf};
#[cfg(feature = "sqlite")]
use crate::sqlite_topology::SqliteTopologyGraph;
#[derive(Debug)]
pub enum BuilderError {
SqliteNotEnabled,
DatabaseError(String),
PersistenceError(String),
}
impl std::fmt::Display for BuilderError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BuilderError::SqliteNotEnabled => {
write!(f, "SQLite feature not enabled. Add features = [\"sqlite\"] to Cargo.toml")
}
BuilderError::DatabaseError(msg) => write!(f, "Database error: {}", msg),
BuilderError::PersistenceError(msg) => write!(f, "Persistence error: {}", msg),
}
}
}
impl std::error::Error for BuilderError {}
pub struct ColonyBuilder {
persistence_path: Option<PathBuf>,
auto_save: bool,
cache_size: usize,
colony_config: ColonyConfig,
}
impl Default for ColonyBuilder {
fn default() -> Self {
Self::new()
}
}
impl ColonyBuilder {
pub fn new() -> Self {
Self {
persistence_path: None,
auto_save: false,
cache_size: 1000,
colony_config: ColonyConfig::default(),
}
}
pub fn with_config(mut self, config: ColonyConfig) -> Self {
self.colony_config = config;
self
}
#[cfg(feature = "sqlite")]
pub fn with_persistence<P: AsRef<Path>>(mut self, path: P) -> Self {
self.persistence_path = Some(path.as_ref().to_path_buf());
self
}
#[cfg(not(feature = "sqlite"))]
pub fn with_persistence<P: AsRef<Path>>(self, _path: P) -> Self {
self
}
pub fn auto_save(mut self, enabled: bool) -> Self {
self.auto_save = enabled;
self
}
pub fn cache_size(mut self, size: usize) -> Self {
self.cache_size = size;
self
}
pub fn build_simple(self) -> Colony {
Colony::from_config(self.colony_config)
}
#[cfg(feature = "sqlite")]
pub fn build(self) -> Result<PersistentColony, BuilderError> {
let mut colony = Colony::from_config(self.colony_config);
let persistence = if let Some(path) = self.persistence_path {
let db = SqliteTopologyGraph::open(&path)
.map_err(|e| BuilderError::DatabaseError(e.to_string()))?
.with_cache_size(self.cache_size);
load_from_sqlite(&db, colony.substrate_mut().graph_mut())?;
Some(PersistenceState {
db,
path,
auto_save: self.auto_save,
})
} else {
None
};
Ok(PersistentColony {
colony,
persistence,
})
}
#[cfg(not(feature = "sqlite"))]
pub fn build(self) -> Result<PersistentColony, BuilderError> {
if self.persistence_path.is_some() {
return Err(BuilderError::SqliteNotEnabled);
}
Ok(PersistentColony {
colony: Colony::from_config(self.colony_config),
persistence: None,
})
}
}
#[cfg(feature = "sqlite")]
struct PersistenceState {
db: SqliteTopologyGraph,
path: PathBuf,
auto_save: bool,
}
#[cfg(not(feature = "sqlite"))]
struct PersistenceState;
pub struct PersistentColony {
colony: Colony,
#[cfg(feature = "sqlite")]
persistence: Option<PersistenceState>,
#[cfg(not(feature = "sqlite"))]
persistence: Option<PersistenceState>,
}
impl PersistentColony {
pub fn colony(&self) -> &Colony {
&self.colony
}
pub fn colony_mut(&mut self) -> &mut Colony {
&mut self.colony
}
pub fn into_inner(mut self) -> Colony {
#[cfg(feature = "sqlite")]
if let Some(ref mut state) = self.persistence {
state.auto_save = false;
}
let colony = std::mem::replace(&mut self.colony, Colony::new());
std::mem::forget(self); colony
}
pub fn has_persistence(&self) -> bool {
self.persistence.is_some()
}
#[cfg(feature = "sqlite")]
pub fn save(&mut self) -> Result<(), BuilderError> {
if let Some(ref mut state) = self.persistence {
save_to_sqlite(self.colony.substrate().graph(), &mut state.db)?;
}
Ok(())
}
#[cfg(not(feature = "sqlite"))]
pub fn save(&mut self) -> Result<(), BuilderError> {
Ok(())
}
#[cfg(feature = "sqlite")]
pub fn persistence_path(&self) -> Option<&Path> {
self.persistence.as_ref().map(|s| s.path.as_path())
}
#[cfg(not(feature = "sqlite"))]
pub fn persistence_path(&self) -> Option<&Path> {
None
}
}
impl PersistentColony {
pub fn run(&mut self, ticks: u64) -> Vec<Vec<crate::colony::ColonyEvent>> {
self.colony.run(ticks)
}
pub fn tick(&mut self) -> Vec<crate::colony::ColonyEvent> {
self.colony.tick()
}
pub fn ingest_document(&mut self, title: &str, content: &str, position: Position) -> DocumentId {
self.colony.ingest_document(title, content, position)
}
pub fn stats(&self) -> crate::colony::ColonyStats {
self.colony.stats()
}
pub fn snapshot(&self) -> crate::colony::ColonySnapshot {
self.colony.snapshot()
}
pub fn spawn(
&mut self,
agent: Box<dyn phago_core::agent::Agent<Input = String, Fragment = String, Presentation = Vec<String>>>,
) -> AgentId {
self.colony.spawn(agent)
}
pub fn alive_count(&self) -> usize {
self.colony.alive_count()
}
}
#[cfg(feature = "sqlite")]
impl Drop for PersistentColony {
fn drop(&mut self) {
if let Some(ref state) = self.persistence {
if state.auto_save {
let _ = save_to_sqlite(self.colony.substrate().graph(), &mut self.persistence.as_mut().unwrap().db);
}
}
}
}
#[cfg(feature = "sqlite")]
fn load_from_sqlite(
source: &SqliteTopologyGraph,
target: &mut PetTopologyGraph,
) -> Result<(), BuilderError> {
let mut node_count = 0;
for node in source.iter_nodes() {
target.add_node(node);
node_count += 1;
}
let mut edge_count = 0;
for (from, to, edge) in source.iter_edges() {
target.set_edge(from, to, edge);
edge_count += 1;
}
if node_count > 0 || edge_count > 0 {
eprintln!(
"Loaded {} nodes and {} edges from SQLite database.",
node_count, edge_count
);
}
Ok(())
}
#[cfg(feature = "sqlite")]
fn save_to_sqlite(
source: &PetTopologyGraph,
target: &mut SqliteTopologyGraph,
) -> Result<(), BuilderError> {
for node_id in source.all_nodes() {
if let Some(node) = source.get_node(&node_id) {
target.add_node(node.clone());
}
}
for (from, to, edge) in source.all_edges() {
target.set_edge(from, to, edge.clone());
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_simple_colony() {
let colony = ColonyBuilder::new().build_simple();
assert_eq!(colony.alive_count(), 0);
}
#[test]
fn build_without_persistence() {
let colony = ColonyBuilder::new().build().unwrap();
assert!(!colony.has_persistence());
}
#[cfg(feature = "sqlite")]
#[test]
fn build_with_persistence() {
let tmp = std::env::temp_dir().join("phago_builder_test.db");
let _ = std::fs::remove_file(&tmp);
let mut colony = ColonyBuilder::new()
.with_persistence(&tmp)
.build()
.unwrap();
assert!(colony.has_persistence());
assert_eq!(colony.persistence_path(), Some(tmp.as_path()));
colony.ingest_document("Test", "Content", Position::new(0.0, 0.0));
colony.run(5);
colony.save().unwrap();
let _ = std::fs::remove_file(&tmp);
}
#[cfg(feature = "sqlite")]
#[test]
fn auto_save_on_drop() {
let tmp = std::env::temp_dir().join("phago_autosave_test.db");
let _ = std::fs::remove_file(&tmp);
{
let mut colony = ColonyBuilder::new()
.with_persistence(&tmp)
.auto_save(true)
.build()
.unwrap();
colony.ingest_document("Test", "Content", Position::new(0.0, 0.0));
colony.run(5);
}
assert!(tmp.exists());
let _ = std::fs::remove_file(&tmp);
}
#[cfg(feature = "sqlite")]
#[test]
fn roundtrip_save_load() {
use phago_agents::digester::Digester;
let tmp = std::env::temp_dir().join("phago_roundtrip_test.db");
let _ = std::fs::remove_file(&tmp);
let (node_count, edge_count) = {
let mut colony = ColonyBuilder::new()
.with_persistence(&tmp)
.build()
.unwrap();
colony.ingest_document("Biology 101", "Cell membrane proteins transport molecules", Position::new(0.0, 0.0));
colony.spawn(Box::new(Digester::new(Position::new(0.0, 0.0)).with_max_idle(50)));
colony.run(15);
let stats = colony.stats();
colony.save().unwrap();
(stats.graph_nodes, stats.graph_edges)
};
let colony2 = ColonyBuilder::new()
.with_persistence(&tmp)
.build()
.unwrap();
let stats2 = colony2.stats();
assert_eq!(stats2.graph_nodes, node_count, "Node count should match after reload");
assert_eq!(stats2.graph_edges, edge_count, "Edge count should match after reload");
let _ = std::fs::remove_file(&tmp);
}
}