use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentState {
pub files: HashMap<PathBuf, FileState>,
pub errors: Vec<ErrorInfo>,
pub occupied_files: HashMap<PathBuf, u32>,
pub investigated: HashMap<PathBuf, u32>,
pub metrics: HashMap<String, f64>,
pub last_progress_tick: u64,
pub total_changes: usize,
}
impl AgentState {
pub fn new() -> Self {
Self::default()
}
pub fn update_file(&mut self, path: PathBuf, state: FileState) {
self.files.insert(path, state);
}
pub fn get_file(&self, path: &PathBuf) -> Option<&FileState> {
self.files.get(path)
}
pub fn add_error(&mut self, error: ErrorInfo) {
self.errors.push(error);
}
pub fn clear_errors(&mut self) {
self.errors.clear();
}
pub fn has_errors(&self) -> bool {
!self.errors.is_empty()
}
pub fn error_count(&self) -> usize {
self.errors.len()
}
pub fn most_errored_file(&self) -> Option<&PathBuf> {
let mut counts: HashMap<&PathBuf, usize> = HashMap::new();
for error in &self.errors {
*counts.entry(&error.file).or_insert(0) += 1;
}
counts.into_iter().max_by_key(|(_, c)| *c).map(|(f, _)| f)
}
pub fn occupy_file(&mut self, path: PathBuf, agent_id: u32) {
self.occupied_files.insert(path, agent_id);
}
pub fn release_file(&mut self, path: &PathBuf) {
self.occupied_files.remove(path);
}
pub fn is_occupied(&self, path: &PathBuf, exclude_agent: u32) -> bool {
self.occupied_files
.get(path)
.map(|&id| id != exclude_agent)
.unwrap_or(false)
}
pub fn mark_investigated(&mut self, path: PathBuf, agent_id: u32) {
self.investigated.insert(path, agent_id);
}
pub fn is_investigated(&self, path: &PathBuf) -> bool {
self.investigated.contains_key(path)
}
pub fn uninvestigated_files(&self) -> Vec<&PathBuf> {
self.files
.keys()
.filter(|p| !self.investigated.contains_key(*p))
.collect()
}
pub fn available_files(&self, agent_id: u32) -> Vec<&PathBuf> {
self.files
.keys()
.filter(|p| !self.investigated.contains_key(*p) && !self.is_occupied(p, agent_id))
.collect()
}
pub fn set_metric(&mut self, key: impl Into<String>, value: f64) {
self.metrics.insert(key.into(), value);
}
pub fn get_metric(&self, key: &str) -> Option<f64> {
self.metrics.get(key).copied()
}
pub fn record_progress(&mut self, tick: u64, changes: usize) {
if changes > 0 {
self.last_progress_tick = tick;
self.total_changes += changes;
}
}
pub fn is_stalled(&self, current_tick: u64, threshold: u64) -> bool {
current_tick.saturating_sub(self.last_progress_tick) >= threshold
}
pub fn summary(&self) -> String {
format!(
"Files: {}, Errors: {}, Investigated: {}/{}, Changes: {}",
self.files.len(),
self.errors.len(),
self.investigated.len(),
self.files.len(),
self.total_changes,
)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct FileState {
pub lines: usize,
pub size_bytes: usize,
pub last_modified_tick: u64,
pub changes: usize,
pub read: bool,
pub modified: bool,
}
impl FileState {
pub fn new(lines: usize, size_bytes: usize) -> Self {
Self {
lines,
size_bytes,
..Default::default()
}
}
pub fn mark_read(&mut self) {
self.read = true;
}
pub fn mark_modified(&mut self, tick: u64, changes: usize) {
self.modified = true;
self.last_modified_tick = tick;
self.changes += changes;
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorInfo {
pub file: PathBuf,
pub line: usize,
pub column: Option<usize>,
pub message: String,
pub severity: ErrorSeverity,
}
impl ErrorInfo {
pub fn new(file: PathBuf, line: usize, message: impl Into<String>) -> Self {
Self {
file,
line,
column: None,
message: message.into(),
severity: ErrorSeverity::Error,
}
}
pub fn with_column(mut self, column: usize) -> Self {
self.column = Some(column);
self
}
pub fn with_severity(mut self, severity: ErrorSeverity) -> Self {
self.severity = severity;
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ErrorSeverity {
Warning,
Error,
Critical,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_agent_state() {
let mut state = AgentState::new();
state.update_file(PathBuf::from("a.rs"), FileState::new(100, 2000));
state.update_file(PathBuf::from("b.rs"), FileState::new(50, 1000));
assert_eq!(state.files.len(), 2);
state.add_error(ErrorInfo::new(PathBuf::from("a.rs"), 10, "error 1"));
state.add_error(ErrorInfo::new(PathBuf::from("a.rs"), 20, "error 2"));
state.add_error(ErrorInfo::new(PathBuf::from("b.rs"), 5, "error 3"));
assert!(state.has_errors());
assert_eq!(state.error_count(), 3);
assert_eq!(state.most_errored_file(), Some(&PathBuf::from("a.rs")));
}
#[test]
fn test_occupation() {
let mut state = AgentState::new();
state.update_file(PathBuf::from("a.rs"), FileState::new(100, 2000));
let path = PathBuf::from("a.rs");
state.occupy_file(path.clone(), 0);
assert!(!state.is_occupied(&path, 0)); assert!(state.is_occupied(&path, 1));
state.release_file(&path);
assert!(!state.is_occupied(&path, 1));
}
#[test]
fn test_investigation() {
let mut state = AgentState::new();
state.update_file(PathBuf::from("a.rs"), FileState::new(100, 2000));
state.update_file(PathBuf::from("b.rs"), FileState::new(50, 1000));
let path_a = PathBuf::from("a.rs");
state.mark_investigated(path_a.clone(), 0);
assert!(state.is_investigated(&path_a));
assert!(!state.is_investigated(&PathBuf::from("b.rs")));
let uninvestigated = state.uninvestigated_files();
assert_eq!(uninvestigated.len(), 1);
}
}