use std::collections::{BTreeMap, HashSet};
use std::fs;
use std::path::Path;
use std::sync::Mutex;
static CATALOG_WRITE_LOCK: Mutex<()> = Mutex::new(());
use serde::{Deserialize, Serialize};
use crate::error::{Result, StorageError};
use super::{SideChatMetadata, ThreadMetadata, ThreadStorage, WorkspaceMetadata};
#[derive(Debug, Default, Serialize, Deserialize)]
pub(super) struct ThreadIndex {
pub(super) workspaces: Vec<WorkspaceMetadata>,
pub(super) unassigned_threads: BTreeMap<String, ThreadMetadata>,
#[serde(default)]
pub(super) sidechat_threads: BTreeMap<String, SideChatMetadata>,
}
fn canonical_workspace_path(path: &Path) -> Result<std::path::PathBuf> {
if !path.is_dir() {
return Err(StorageError::InvalidWorkspacePath(
path.display().to_string(),
));
}
let canonical = fs::canonicalize(path)?;
if canonical.parent().is_none() {
return Err(StorageError::InvalidWorkspacePath(
path.display().to_string(),
));
}
if dirs::home_dir()
.and_then(|home| fs::canonicalize(home).ok())
.is_some_and(|home| home == canonical)
{
return Err(StorageError::InvalidWorkspacePath(
path.display().to_string(),
));
}
#[cfg(unix)]
{
const PROTECTED_PATHS: &[&str] = &[
"/Applications",
"/Library",
"/System",
"/Users",
"/Volumes",
"/bin",
"/boot",
"/dev",
"/etc",
"/home",
"/lib",
"/lib64",
"/opt",
"/proc",
"/root",
"/run",
"/sbin",
"/sys",
"/usr",
"/var",
];
if PROTECTED_PATHS
.iter()
.any(|protected| canonical == Path::new(protected))
{
return Err(StorageError::InvalidWorkspacePath(
path.display().to_string(),
));
}
}
#[cfg(windows)]
{
let normalized = canonical
.to_string_lossy()
.replace('/', "\\")
.to_lowercase();
let drive_relative = normalized
.strip_prefix(r"\\?\")
.unwrap_or(normalized.as_str());
let components = drive_relative
.split('\\')
.filter(|component| !component.is_empty())
.collect::<Vec<_>>();
let protected = [
"program files",
"program files (x86)",
"programdata",
"users",
"windows",
];
if components.len() <= 1
|| components
.get(1)
.is_some_and(|component| protected.contains(component))
{
return Err(StorageError::InvalidWorkspacePath(
path.display().to_string(),
));
}
}
Ok(canonical)
}
impl ThreadStorage {
pub(super) fn read_index(&self) -> Result<ThreadIndex> {
if !self.index_path.exists() {
return Ok(ThreadIndex::default());
}
let index_json = fs::read_to_string(&self.index_path)?;
serde_json::from_str::<ThreadIndex>(&index_json).map_err(Into::into)
}
fn write_index(&self, index: &ThreadIndex) -> Result<()> {
let json = serde_json::to_string_pretty(index)?;
super::atomic_write(&self.index_path, json.as_bytes())
}
fn validate_workspace_input(
name: &str,
directories: &[String],
) -> Result<(String, Vec<String>)> {
let name = name.trim();
if name.is_empty() {
return Err(StorageError::InvalidWorkspaceName);
}
let directories = directories
.iter()
.map(|path| canonical_workspace_path(Path::new(path)))
.collect::<Result<Vec<_>>>()?;
Ok((
name.to_string(),
directories
.into_iter()
.map(|path| path.to_string_lossy().into_owned())
.collect(),
))
}
pub fn create_workspace(
&self,
name: &str,
directories: &[String],
) -> Result<WorkspaceMetadata> {
let (name, directories) = Self::validate_workspace_input(name, directories)?;
let workspace = WorkspaceMetadata::new(name, directories);
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
index.workspaces.push(workspace.clone());
self.write_index(&index)?;
Ok(workspace)
}
pub fn update_workspace(
&self,
workspace_id: &str,
name: &str,
directories: &[String],
) -> Result<WorkspaceMetadata> {
let name = name.trim();
if name.is_empty() {
return Err(StorageError::InvalidWorkspaceName);
}
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
let workspace = index
.workspaces
.iter_mut()
.find(|workspace| workspace.id == workspace_id)
.ok_or_else(|| StorageError::WorkspaceNotFound(workspace_id.to_string()))?;
let directories = Self::validate_workspace_input(name, directories)?.1;
workspace.name = name.to_string();
workspace.directories = directories;
let updated = workspace.clone();
self.write_index(&index)?;
Ok(updated)
}
pub fn delete_workspace(&self, workspace_id: &str) -> Result<()> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
let workspace_index = index
.workspaces
.iter()
.position(|workspace| workspace.id == workspace_id)
.ok_or_else(|| StorageError::WorkspaceNotFound(workspace_id.to_string()))?;
let removed = index.workspaces.remove(workspace_index);
index.unassigned_threads.extend(removed.threads);
self.write_index(&index)
}
pub(super) fn get_index_metadata(&self, thread_id: &str) -> Result<ThreadMetadata> {
let index = self.read_index()?;
index
.unassigned_threads
.get(thread_id)
.cloned()
.or_else(|| {
index
.workspaces
.iter()
.find_map(|workspace| workspace.threads.get(thread_id).cloned())
})
.ok_or_else(|| StorageError::ThreadNotFound(thread_id.to_string()))
}
pub fn get_thread_workspace_id(&self, thread_id: &str) -> Result<Option<String>> {
let index = self.read_index()?;
if index.unassigned_threads.contains_key(thread_id) {
return Ok(None);
}
index
.workspaces
.iter()
.find(|workspace| workspace.threads.contains_key(thread_id))
.map(|workspace| Some(workspace.id.clone()))
.ok_or_else(|| StorageError::ThreadNotFound(thread_id.to_string()))
}
pub(super) fn update_index(&self, metadata: &ThreadMetadata) -> Result<()> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
if let Some(workspace) = index
.workspaces
.iter_mut()
.find(|workspace| workspace.threads.contains_key(&metadata.id))
{
workspace
.threads
.insert(metadata.id.clone(), metadata.clone());
} else {
index
.unassigned_threads
.insert(metadata.id.clone(), metadata.clone());
}
self.write_index(&index)
}
pub(super) fn get_sidechat_metadata(&self, sidechat_id: &str) -> Result<SideChatMetadata> {
self.read_index()?
.sidechat_threads
.get(sidechat_id)
.cloned()
.ok_or_else(|| StorageError::ThreadNotFound(sidechat_id.to_string()))
}
pub(super) fn update_sidechat_index(&self, metadata: &SideChatMetadata) -> Result<()> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
index
.sidechat_threads
.insert(metadata.id.clone(), metadata.clone());
self.write_index(&index)
}
pub fn list_sidechat_threads(&self) -> Result<Vec<SideChatMetadata>> {
Ok(self.read_index()?.sidechat_threads.into_values().collect())
}
pub(super) fn remove_sidechat_from_index(&self, sidechat_id: &str) -> Result<()> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
index.sidechat_threads.remove(sidechat_id);
self.write_index(&index)
}
pub(super) fn update_index_in_workspace(
&self,
metadata: &ThreadMetadata,
workspace_id: Option<&str>,
) -> Result<()> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
if let Some(id) = workspace_id {
if !index.workspaces.iter().any(|workspace| workspace.id == id) {
return Err(StorageError::WorkspaceNotFound(id.to_string()));
}
}
let metadata = index
.unassigned_threads
.get(&metadata.id)
.or_else(|| {
index
.workspaces
.iter()
.find_map(|workspace| workspace.threads.get(&metadata.id))
})
.cloned()
.unwrap_or_else(|| metadata.clone());
index.unassigned_threads.remove(&metadata.id);
for workspace in &mut index.workspaces {
workspace.threads.remove(&metadata.id);
}
if let Some(id) = workspace_id {
let workspace = index
.workspaces
.iter_mut()
.find(|workspace| workspace.id == id)
.unwrap();
workspace
.threads
.insert(metadata.id.clone(), metadata.clone());
} else {
index
.unassigned_threads
.insert(metadata.id.clone(), metadata.clone());
}
self.write_index(&index)
}
pub fn group_threads(&self, first_id: &str, second_id: &str) -> Result<WorkspaceMetadata> {
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
if first_id == second_id
|| !index.unassigned_threads.contains_key(first_id)
|| !index.unassigned_threads.contains_key(second_id)
{
return Err(StorageError::ThreadNotFound(second_id.to_string()));
}
let mut workspace = WorkspaceMetadata::new("New workspace".to_string(), Vec::new());
for id in [first_id, second_id] {
workspace
.threads
.insert(id.to_string(), index.unassigned_threads.remove(id).unwrap());
}
index.workspaces.push(workspace.clone());
self.write_index(&index)?;
Ok(workspace)
}
pub fn list_unassigned_threads(&self) -> Result<Vec<ThreadMetadata>> {
Ok(self
.read_index()?
.unassigned_threads
.into_values()
.collect())
}
pub(super) fn remove_many_from_index(&self, thread_ids: &[String]) -> Result<()> {
let removed = thread_ids
.iter()
.map(String::as_str)
.collect::<HashSet<_>>();
let _guard = CATALOG_WRITE_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let mut index = self.read_index()?;
index
.unassigned_threads
.retain(|id, _| !removed.contains(id.as_str()));
for workspace in &mut index.workspaces {
workspace
.threads
.retain(|id, _| !removed.contains(id.as_str()));
}
self.write_index(&index)
}
}