use std::path::Path;
use std::sync::Arc;
use anyhow::Result;
use tokio::sync::{broadcast, RwLock};
use tracing::info;
use crate::core::config::CompanyConfig;
use crate::core::messaging::MessageBus;
use crate::core::store::Store;
use crate::domain::{Message, Organization};
use crate::infrastructure::store::SqliteStore;
use super::company_runtime::{AgentManager, OrganizationManager, ToolCapabilityManager};
use crate::{
CapabilityRegistry, FrameworkToolExecutor, McpProtocolHandler, McpServer, ToolEnvironment,
ToolRegistry,
};
pub struct VirtualCompany {
organization_manager: OrganizationManager,
agent_manager: AgentManager,
tool_capability_manager: ToolCapabilityManager,
message_bus: Arc<MessageBus>,
message_tx: broadcast::Sender<Message>,
store: Arc<dyn Store>,
watchdog_agent: Arc<crate::core::watchdog_agent::WatchdogAgent>,
}
impl VirtualCompany {
pub fn from_config(config: CompanyConfig) -> Result<Self> {
Self::with_sqlite(
config,
std::env::var("DB_PATH").unwrap_or_else(|_| "imitatort.db".to_string()),
)
}
pub fn with_sqlite<P: AsRef<Path>>(config: CompanyConfig, db_path: P) -> Result<Self> {
let store = Arc::new(SqliteStore::new(db_path)?);
Ok(Self::with_store(config, store))
}
pub fn with_store(config: CompanyConfig, store: Arc<dyn Store>) -> Self {
let message_bus = Arc::new(MessageBus::with_store(store.clone()));
let (message_tx, _) = broadcast::channel(1000);
let organization_manager = OrganizationManager::new(config);
let tool_capability_manager = ToolCapabilityManager::new();
let agent_manager = AgentManager::new(message_bus.clone());
let framework_executor = tool_capability_manager.get_framework_tool_executor(
message_bus.clone(),
organization_manager.organization_arc(),
store.clone(),
);
let watchdog_agent = Arc::new(crate::core::watchdog_agent::WatchdogAgent::new(
crate::domain::Agent::new(
"system_watchdog",
"System Watchdog Agent",
crate::domain::Role::simple("System", "System monitoring agent"),
crate::domain::LLMConfig::openai("dummy-key"),
),
Arc::new(framework_executor),
));
Self {
organization_manager,
agent_manager,
tool_capability_manager,
message_bus,
message_tx,
store,
watchdog_agent,
}
}
pub async fn from_sqlite<P: AsRef<Path>>(db_path: P) -> Result<Self> {
let store = Arc::new(SqliteStore::new(db_path)?);
Self::from_store(store).await
}
pub async fn from_store(store: Arc<dyn Store>) -> Result<Self> {
let org = store.load_organization().await?;
if org.agents.is_empty() {
return Err(anyhow::anyhow!(
"No organization found in store. Please create config first."
));
}
let config = CompanyConfig {
name: "Loaded Company".to_string(),
organization: org,
};
Ok(Self::with_store(config, store))
}
pub async fn save(&self) -> Result<()> {
info!("Saving company state to storage...");
let org = self.organization_manager.organization().await;
self.store.save_organization(&org).await?;
info!("Company state saved successfully");
Ok(())
}
pub async fn run(&self) -> Result<()> {
info!(
"Starting virtual company: {}",
self.organization_manager.config().name
);
let org = self.organization_manager.organization().await;
self.agent_manager
.initialize_agents(&org, Some(self.watchdog_agent.clone()))
.await?;
drop(org);
info!(
"All {} agents initialized",
self.agent_manager.get_agents().await?.len()
);
info!("Agents initialized and ready for event-driven activation");
Ok(())
}
pub fn subscribe_messages(&self) -> broadcast::Receiver<Message> {
self.message_tx.subscribe()
}
pub async fn organization(&self) -> tokio::sync::RwLockReadGuard<'_, Organization> {
self.organization_manager.organization().await
}
pub fn organization_arc(&self) -> Arc<RwLock<Organization>> {
self.organization_manager.organization_arc()
}
pub fn store(&self) -> &Arc<dyn Store> {
&self.store
}
pub fn name(&self) -> &str {
&self.organization_manager.config().name
}
pub async fn get_agents(&self) -> Result<Vec<crate::domain::Agent>> {
let org = self.organization_manager.organization().await;
Ok(org.agents.clone())
}
pub fn tool_registry(&self) -> Arc<ToolRegistry> {
self.tool_capability_manager.tool_registry()
}
pub async fn register_app_tool(&self, tool: crate::domain::tool::Tool) -> Result<()> {
self.tool_capability_manager.register_app_tool(tool).await
}
pub fn create_tool_environment(&self) -> ToolEnvironment {
self.tool_capability_manager.create_tool_environment(
self.message_bus.clone(),
self.organization_manager.organization_arc(),
self.store.clone(),
)
}
pub fn get_framework_tool_executor(&self) -> FrameworkToolExecutor {
self.tool_capability_manager.get_framework_tool_executor(
self.message_bus.clone(),
self.organization_manager.organization_arc(),
self.store.clone(),
)
}
pub fn capability_registry(&self) -> Arc<CapabilityRegistry> {
self.tool_capability_manager.capability_registry()
}
pub async fn register_app_capability(
&self,
capability: crate::domain::capability::Capability,
) -> Result<()> {
self.tool_capability_manager
.register_app_capability(capability)
.await
}
pub fn create_mcp_server(&self, bind_addr: String) -> McpServer {
self.tool_capability_manager.create_mcp_server(bind_addr)
}
pub fn get_mcp_protocol_handler(&self) -> McpProtocolHandler {
self.tool_capability_manager.get_mcp_protocol_handler()
}
pub fn register_skill(&self, skill: crate::domain::skill::Skill) -> Result<()> {
self.tool_capability_manager.register_skill(skill)
}
pub fn bind_skill_tool(&self, binding: crate::domain::skill::SkillToolBinding) -> Result<()> {
self.tool_capability_manager.bind_skill_tool(binding)
}
pub fn set_tool_access(
&self,
tool_id: &str,
access_type: crate::domain::skill::ToolAccessType,
) -> Result<()> {
self.tool_capability_manager
.set_tool_access(tool_id, access_type)
}
pub fn skill_manager(&self) -> Arc<crate::core::skill::SkillManager> {
self.tool_capability_manager.skill_manager()
}
}
pub struct CompanyBuilder {
config: Option<CompanyConfig>,
store: Option<Arc<dyn Store>>,
}
impl CompanyBuilder {
pub fn new() -> Result<Self> {
Self::with_sqlite(std::env::var("DB_PATH").unwrap_or_else(|_| "imitatort.db".to_string()))
}
pub fn from_config(config: CompanyConfig) -> Result<Self> {
let db_path = std::env::var("DB_PATH").unwrap_or_else(|_| "imitatort.db".to_string());
let mut builder = Self::with_sqlite(&db_path)?;
builder.config = Some(config);
Ok(builder)
}
pub fn with_sqlite<P: AsRef<Path>>(db_path: P) -> Result<Self> {
let store = Arc::new(SqliteStore::new(db_path)?);
Ok(Self {
config: None,
store: Some(store),
})
}
pub fn with_store(store: Arc<dyn Store>) -> Self {
Self {
config: None,
store: Some(store),
}
}
pub fn config(mut self, config: CompanyConfig) -> Self {
self.config = Some(config);
self
}
pub async fn load(mut self) -> Result<Self> {
if let Some(ref store) = self.store {
let org = store.load_organization().await?;
if !org.agents.is_empty() {
self.config = Some(CompanyConfig {
name: "Loaded Company".to_string(),
organization: org,
});
}
}
Ok(self)
}
pub fn build(self) -> Result<VirtualCompany> {
let config = self
.config
.ok_or_else(|| anyhow::anyhow!("Config not set. Use .config() or .load() first."))?;
let store = self
.store
.ok_or_else(|| anyhow::anyhow!("Store not set."))?;
Ok(VirtualCompany::with_store(config, store))
}
pub async fn build_and_save(self) -> Result<VirtualCompany> {
let company = self.build()?;
company.save().await?;
Ok(company)
}
}