use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use chrono::{DateTime, Utc};
use crate::version::{Version, ToolInfo, ToolRegistry};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum TrafficLevel {
Green,
Yellow,
Red,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum UpdateStatus {
Available,
Downloading,
Applying,
Applied,
Failed(String),
RequiresManual,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateNotification {
pub tool_name: String,
pub current_version: Version,
pub new_version: Version,
pub traffic_level: TrafficLevel,
pub status: UpdateStatus,
pub timestamp: DateTime<Utc>,
pub message: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Backup {
pub id: String,
pub tool_name: String,
pub version: Version,
pub created_at: DateTime<Utc>,
pub backup_path: PathBuf,
}
pub struct AutoUpdateManager {
registry: ToolRegistry,
backups: HashMap<String, Vec<Backup>>,
notifications: Vec<UpdateNotification>,
auto_update_enabled: bool,
}
impl AutoUpdateManager {
pub fn new(forge_dir: &Path) -> Result<Self> {
let registry = ToolRegistry::new(forge_dir)?;
Ok(Self {
registry,
backups: HashMap::new(),
notifications: Vec::new(),
auto_update_enabled: true,
})
}
pub fn set_auto_update(&mut self, enabled: bool) {
self.auto_update_enabled = enabled;
}
pub fn check_update(&self, tool_name: &str, latest_version: Version) -> Option<UpdateNotification> {
if let Some(current_version) = self.registry.version(tool_name) {
if &latest_version > current_version {
return Some(UpdateNotification {
tool_name: tool_name.to_string(),
current_version: current_version.clone(),
new_version: latest_version.clone(),
traffic_level: self.determine_traffic_level(current_version, &latest_version),
status: UpdateStatus::Available,
timestamp: Utc::now(),
message: format!("Update available: {} -> {}", current_version, latest_version),
});
}
}
None
}
fn determine_traffic_level(&self, current: &Version, new: &Version) -> TrafficLevel {
if current.major == new.major && current.minor == new.minor {
return TrafficLevel::Green;
}
if current.major == new.major {
return TrafficLevel::Yellow;
}
TrafficLevel::Red
}
pub fn process_update(&mut self, notification: &UpdateNotification) -> Result<()> {
if !self.auto_update_enabled {
return Ok(());
}
match notification.traffic_level {
TrafficLevel::Green => {
println!("🟢 Auto-applying green traffic update: {} {} -> {}",
notification.tool_name,
notification.current_version,
notification.new_version
);
self.apply_update(notification)?;
},
TrafficLevel::Yellow => {
println!("🟡 Yellow traffic update available: {} {} -> {} (requires review)",
notification.tool_name,
notification.current_version,
notification.new_version
);
self.add_notification(notification.clone());
},
TrafficLevel::Red => {
println!("🔴 Red traffic update available: {} {} -> {} (requires manual intervention)",
notification.tool_name,
notification.current_version,
notification.new_version
);
self.add_notification(notification.clone());
},
}
Ok(())
}
fn apply_update(&mut self, notification: &UpdateNotification) -> Result<()> {
self.create_backup(¬ification.tool_name, ¬ification.current_version)?;
println!(" ✓ Created backup");
println!(" ⬇ Downloading update...");
println!(" ✓ Applying update...");
let _tool_info = ToolInfo {
name: notification.tool_name.clone(),
version: notification.new_version.clone(),
installed_at: Utc::now(),
source: crate::version::ToolSource::Crate {
version: notification.new_version.to_string(),
},
dependencies: HashMap::new(),
};
println!(" ✓ Update applied successfully");
self.add_notification(UpdateNotification {
status: UpdateStatus::Applied,
timestamp: Utc::now(),
message: format!("Successfully updated to {}", notification.new_version),
..notification.clone()
});
Ok(())
}
fn create_backup(&mut self, tool_name: &str, version: &Version) -> Result<()> {
let backup_id = uuid::Uuid::new_v4().to_string();
let backup_path = PathBuf::from(format!(".dx/forge/backups/{}/{}", tool_name, backup_id));
std::fs::create_dir_all(&backup_path)?;
let backup = Backup {
id: backup_id,
tool_name: tool_name.to_string(),
version: version.clone(),
created_at: Utc::now(),
backup_path,
};
self.backups
.entry(tool_name.to_string())
.or_insert_with(Vec::new)
.push(backup);
Ok(())
}
pub fn rollback(&mut self, tool_name: &str) -> Result<()> {
let backups = self.backups
.get_mut(tool_name)
.ok_or_else(|| anyhow!("No backups found for {}", tool_name))?;
let backup = backups.pop()
.ok_or_else(|| anyhow!("No backups available for {}", tool_name))?;
println!("🔄 Rolling back {} to version {}", tool_name, backup.version);
println!(" ✓ Restoring from backup: {}", backup.id);
println!(" ✓ Rollback complete");
self.add_notification(UpdateNotification {
tool_name: tool_name.to_string(),
current_version: backup.version.clone(),
new_version: backup.version.clone(),
traffic_level: TrafficLevel::Green,
status: UpdateStatus::Applied,
timestamp: Utc::now(),
message: format!("Rolled back to version {}", backup.version),
});
Ok(())
}
fn add_notification(&mut self, notification: UpdateNotification) {
self.notifications.push(notification);
if self.notifications.len() > 100 {
self.notifications.remove(0);
}
}
pub fn get_notifications(&self) -> &[UpdateNotification] {
&self.notifications
}
pub fn get_pending_updates(&self) -> Vec<&UpdateNotification> {
self.notifications
.iter()
.filter(|n| n.status == UpdateStatus::Available || n.status == UpdateStatus::RequiresManual)
.collect()
}
pub fn clear_old_notifications(&mut self, days: i64) {
let cutoff = Utc::now() - chrono::Duration::days(days);
self.notifications.retain(|n| n.timestamp > cutoff);
}
pub fn detect_conflicts(&self) -> Vec<String> {
let mut conflicts = Vec::new();
for tool_info in self.registry.list() {
if let Ok(missing_deps) = self.registry.check_dependencies(&tool_info.name) {
conflicts.extend(missing_deps);
}
}
conflicts
}
pub fn get_backups(&self, tool_name: &str) -> Vec<&Backup> {
self.backups
.get(tool_name)
.map(|backups| backups.iter().collect())
.unwrap_or_default()
}
pub fn clean_old_backups(&mut self, days: i64) -> Result<()> {
let cutoff = Utc::now() - chrono::Duration::days(days);
for backups in self.backups.values_mut() {
backups.retain(|backup| {
if backup.created_at < cutoff {
let _ = std::fs::remove_dir_all(&backup.backup_path);
false
} else {
true
}
});
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdatePreference {
pub auto_update_green: bool,
pub notify_yellow: bool,
pub notify_red: bool,
pub email_notifications: bool,
}
impl Default for UpdatePreference {
fn default() -> Self {
Self {
auto_update_green: true,
notify_yellow: true,
notify_red: true,
email_notifications: false,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_determine_traffic_level() {
let manager = AutoUpdateManager {
registry: ToolRegistry::new(Path::new(".dx/forge")).unwrap(),
backups: HashMap::new(),
notifications: Vec::new(),
auto_update_enabled: true,
};
let current = Version::new(1, 2, 3);
let new = Version::new(1, 2, 4);
assert_eq!(manager.determine_traffic_level(¤t, &new), TrafficLevel::Green);
let new = Version::new(1, 3, 0);
assert_eq!(manager.determine_traffic_level(¤t, &new), TrafficLevel::Yellow);
let new = Version::new(2, 0, 0);
assert_eq!(manager.determine_traffic_level(¤t, &new), TrafficLevel::Red);
}
}