use anyhow::Result;
use rustc_hash::FxHashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::RwLock;
#[derive(Debug, Clone)]
pub struct LazyProjectContext {
pub root_path: String,
pub file_paths: Arc<Vec<PathBuf>>,
pub files: Arc<RwLock<FxHashMap<String, String>>>,
pub token_count: Arc<AtomicUsize>,
pub loading_queue: Arc<Mutex<Vec<PathBuf>>>,
pub cache: Option<Arc<crate::cache::CacheManager>>,
}
impl LazyProjectContext {
pub fn new(root_path: String, file_paths: Vec<PathBuf>) -> Self {
let cache = crate::cache::CacheManager::new().ok().map(Arc::new);
Self {
root_path,
file_paths: Arc::new(file_paths),
files: Arc::new(RwLock::new(FxHashMap::default())),
token_count: Arc::new(AtomicUsize::new(0)),
loading_queue: Arc::new(Mutex::new(Vec::new())),
cache,
}
}
pub async fn get_file(&self, path: &str) -> Result<Option<String>> {
{
let files = self.files.read().await;
if let Some(content) = files.get(path) {
return Ok(Some(content.clone()));
}
}
let full_path = if path.starts_with(&self.root_path) {
PathBuf::from(path)
} else {
PathBuf::from(&self.root_path).join(path)
};
if full_path.exists() {
let content = tokio::fs::read_to_string(&full_path).await?;
if let Some(ref cache) = self.cache {
if let Ok(tokens) = cache.get_or_compute_tokens(&full_path, &content, "cl100k_base")
{
self.token_count.fetch_add(tokens, Ordering::Relaxed);
}
}
let mut files = self.files.write().await;
files.insert(path.to_string(), content.clone());
Ok(Some(content))
} else {
Ok(None)
}
}
pub async fn load_files_batch(&self, paths: Vec<String>) -> Result<()> {
use futures::future::join_all;
let futures = paths.into_iter().map(|path| {
let self_clone = self.clone();
async move {
let _ = self_clone.get_file(&path).await;
}
});
join_all(futures).await;
Ok(())
}
pub fn get_file_list(&self) -> Vec<String> {
self.file_paths
.iter()
.filter_map(|p| {
p.strip_prefix(&self.root_path)
.ok()
.and_then(|p| p.to_str())
.map(|s| s.to_string())
})
.collect()
}
pub async fn loaded_file_count(&self) -> usize {
self.files.read().await.len()
}
pub fn total_file_count(&self) -> usize {
self.file_paths.len()
}
pub async fn is_fully_loaded(&self) -> bool {
self.loaded_file_count().await >= self.total_file_count()
}
pub async fn to_project_context(&self) -> crate::models::ProjectContext {
let files = self.files.read().await;
let mut context = crate::models::ProjectContext::new(self.root_path.clone());
context.token_count = self.token_count.load(Ordering::Relaxed);
for (path, content) in files.iter() {
context.add_file(path.clone(), content.clone());
}
context
}
}
pub fn get_priority_files(root_path: &str) -> Vec<String> {
vec![
"README.md",
"readme.md",
"README.rst",
"README.txt",
"CLAUDE.md", "Cargo.toml",
"package.json",
"pyproject.toml",
"go.mod",
".gitignore",
"LICENSE",
]
.into_iter()
.filter_map(|f| {
let path = Path::new(root_path).join(f);
if path.exists() {
Some(f.to_string())
} else {
None
}
})
.collect()
}