use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use crate::storage::Database;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneratedFileInfo {
pub file_path: PathBuf,
pub tool_name: String,
pub generated_at: DateTime<Utc>,
pub source_hash: String,
pub metadata: HashMap<String, String>,
pub modified: bool,
}
impl GeneratedFileInfo {
pub fn new(
file_path: PathBuf,
tool_name: String,
source_hash: String,
metadata: HashMap<String, String>,
) -> Self {
Self {
file_path,
tool_name,
generated_at: Utc::now(),
source_hash,
metadata,
modified: false,
}
}
pub async fn check_modified(&mut self) -> Result<bool> {
if !self.file_path.exists() {
self.modified = true;
return Ok(true);
}
let content = tokio::fs::read(&self.file_path).await?;
let mut hasher = Sha256::new();
hasher.update(&content);
let current_hash = format!("{:x}", hasher.finalize());
self.modified = current_hash != self.source_hash;
Ok(self.modified)
}
}
pub struct GeneratedCodeTracker {
_db: Database,
tracked_files: HashMap<PathBuf, GeneratedFileInfo>,
index_path: PathBuf,
}
impl GeneratedCodeTracker {
pub fn new(forge_dir: &Path) -> Result<Self> {
let db_path = forge_dir.join("forge.db");
let db = Database::new(&db_path)
.context("Failed to open tracking database")?;
let index_path = forge_dir.join("generated_files.json");
let tracked_files = if index_path.exists() {
let content = std::fs::read_to_string(&index_path)?;
serde_json::from_str(&content).unwrap_or_default()
} else {
HashMap::new()
};
Ok(Self {
_db: db,
tracked_files,
index_path,
})
}
pub fn track_file(
&mut self,
file_path: PathBuf,
tool_name: &str,
metadata: HashMap<String, String>,
) -> Result<()> {
let source_hash = if file_path.exists() {
let content = std::fs::read(&file_path)?;
let mut hasher = Sha256::new();
hasher.update(&content);
format!("{:x}", hasher.finalize())
} else {
String::new()
};
let info = GeneratedFileInfo::new(
file_path.clone(),
tool_name.to_string(),
source_hash,
metadata,
);
self.tracked_files.insert(file_path.clone(), info);
self.save_index()?;
tracing::debug!("Tracked generated file: {:?}", file_path);
Ok(())
}
pub fn get_files_by_tool(&self, tool_name: &str) -> Vec<PathBuf> {
self.tracked_files
.values()
.filter(|info| info.tool_name == tool_name)
.map(|info| info.file_path.clone())
.collect()
}
pub fn get_file_info(&self, path: &Path) -> Option<&GeneratedFileInfo> {
self.tracked_files.get(path)
}
pub fn is_tracked(&self, path: &Path) -> bool {
self.tracked_files.contains_key(path)
}
pub fn get_all_files(&self) -> Vec<PathBuf> {
self.tracked_files.keys().cloned().collect()
}
pub fn untrack_file(&mut self, path: &Path) -> Result<()> {
if self.tracked_files.remove(path).is_some() {
self.save_index()?;
tracing::debug!("Untracked file: {:?}", path);
}
Ok(())
}
pub async fn cleanup_tool_files(&mut self, tool_name: &str) -> Result<Vec<PathBuf>> {
let files_to_remove: Vec<PathBuf> = self.get_files_by_tool(tool_name);
let mut removed = Vec::new();
for file in &files_to_remove {
if file.exists() {
tokio::fs::remove_file(file).await
.context(format!("Failed to remove file: {:?}", file))?;
removed.push(file.clone());
tracing::info!("Removed generated file: {:?}", file);
}
self.tracked_files.remove(file);
}
if !removed.is_empty() {
self.save_index()?;
}
Ok(removed)
}
pub async fn check_modifications(&mut self) -> Result<Vec<PathBuf>> {
let mut modified = Vec::new();
for (path, info) in &mut self.tracked_files {
if info.check_modified().await? {
modified.push(path.clone());
}
}
if !modified.is_empty() {
self.save_index()?;
}
Ok(modified)
}
pub fn stats(&self) -> TrackingStats {
let mut by_tool: HashMap<String, usize> = HashMap::new();
for info in self.tracked_files.values() {
*by_tool.entry(info.tool_name.clone()).or_insert(0) += 1;
}
let modified_count = self.tracked_files
.values()
.filter(|info| info.modified)
.count();
TrackingStats {
total_files: self.tracked_files.len(),
files_by_tool: by_tool,
modified_files: modified_count,
}
}
fn save_index(&self) -> Result<()> {
let content = serde_json::to_string_pretty(&self.tracked_files)?;
std::fs::write(&self.index_path, content)?;
Ok(())
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TrackingStats {
pub total_files: usize,
pub files_by_tool: HashMap<String, usize>,
pub modified_files: usize,
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_track_file() {
let temp_dir = TempDir::new().unwrap();
let mut tracker = GeneratedCodeTracker::new(temp_dir.path()).unwrap();
let test_file = temp_dir.path().join("test.txt");
std::fs::write(&test_file, "test content").unwrap();
let mut metadata = HashMap::new();
metadata.insert("version".to_string(), "1.0.0".to_string());
tracker.track_file(test_file.clone(), "test-tool", metadata).unwrap();
assert!(tracker.is_tracked(&test_file));
let files = tracker.get_files_by_tool("test-tool");
assert_eq!(files.len(), 1);
}
#[tokio::test]
async fn test_cleanup_tool_files() {
let temp_dir = TempDir::new().unwrap();
let mut tracker = GeneratedCodeTracker::new(temp_dir.path()).unwrap();
let test_file = temp_dir.path().join("test.txt");
std::fs::write(&test_file, "test content").unwrap();
tracker.track_file(test_file.clone(), "test-tool", HashMap::new()).unwrap();
let removed = tracker.cleanup_tool_files("test-tool").await.unwrap();
assert_eq!(removed.len(), 1);
assert!(!test_file.exists());
assert!(!tracker.is_tracked(&test_file));
}
#[tokio::test]
async fn test_modification_detection() {
let temp_dir = TempDir::new().unwrap();
let mut tracker = GeneratedCodeTracker::new(temp_dir.path()).unwrap();
let test_file = temp_dir.path().join("test.txt");
std::fs::write(&test_file, "original content").unwrap();
tracker.track_file(test_file.clone(), "test-tool", HashMap::new()).unwrap();
std::fs::write(&test_file, "modified content").unwrap();
let modified = tracker.check_modifications().await.unwrap();
assert_eq!(modified.len(), 1);
assert_eq!(modified[0], test_file);
}
}