1use chrono::{DateTime, Utc};
7use parking_lot::Mutex;
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::HashMap;
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct StateTransition {
15 pub key: String,
16 pub old_value: Option<Value>,
17 pub new_value: Option<Value>,
18 pub action_id: String,
19 pub timestamp: DateTime<Utc>,
20}
21
22pub struct StateStore {
27 state: Mutex<HashMap<String, Value>>,
28 transitions: Mutex<Vec<StateTransition>>,
29}
30
31impl StateStore {
32 pub fn new() -> Self {
33 Self {
34 state: Mutex::new(HashMap::new()),
35 transitions: Mutex::new(Vec::new()),
36 }
37 }
38
39 pub fn get(&self, key: &str) -> Option<Value> {
40 self.state.lock().get(key).cloned()
41 }
42
43 pub fn get_or(&self, key: &str, default: Value) -> Value {
44 self.state.lock().get(key).cloned().unwrap_or(default)
45 }
46
47 pub fn exists(&self, key: &str) -> bool {
48 self.state.lock().contains_key(key)
49 }
50
51 pub fn set(&self, key: &str, value: Value, action_id: &str) -> StateTransition {
52 let mut state = self.state.lock();
53 let old = state.get(key).cloned();
54 state.insert(key.to_string(), value.clone());
55
56 let t = StateTransition {
57 key: key.to_string(),
58 old_value: old,
59 new_value: Some(value),
60 action_id: action_id.to_string(),
61 timestamp: Utc::now(),
62 };
63
64 self.transitions.lock().push(t.clone());
65 t
66 }
67
68 pub fn delete(&self, key: &str, action_id: &str) -> Option<StateTransition> {
69 let mut state = self.state.lock();
70 let old = state.remove(key)?;
71
72 let t = StateTransition {
73 key: key.to_string(),
74 old_value: Some(old),
75 new_value: None,
76 action_id: action_id.to_string(),
77 timestamp: Utc::now(),
78 };
79
80 self.transitions.lock().push(t.clone());
81 Some(t)
82 }
83
84 pub fn snapshot(&self) -> HashMap<String, Value> {
86 self.state.lock().clone()
87 }
88
89 pub fn restore(&self, snapshot: HashMap<String, Value>, transition_count: usize) {
91 *self.state.lock() = snapshot;
92 self.transitions.lock().truncate(transition_count);
93 }
94
95 pub fn transition_count(&self) -> usize {
96 self.transitions.lock().len()
97 }
98
99 pub fn transitions(&self) -> Vec<StateTransition> {
100 self.transitions.lock().clone()
101 }
102
103 pub fn transitions_since(&self, index: usize) -> Vec<StateTransition> {
104 self.transitions.lock()[index..].to_vec()
105 }
106
107 pub fn keys(&self) -> Vec<String> {
108 self.state.lock().keys().cloned().collect()
109 }
110}
111
112impl Default for StateStore {
113 fn default() -> Self {
114 Self::new()
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 #[test]
123 fn set_and_get() {
124 let store = StateStore::new();
125 store.set("x", Value::from(42), "test");
126 assert_eq!(store.get("x"), Some(Value::from(42)));
127 }
128
129 #[test]
130 fn exists() {
131 let store = StateStore::new();
132 assert!(!store.exists("x"));
133 store.set("x", Value::from(1), "test");
134 assert!(store.exists("x"));
135 }
136
137 #[test]
138 fn delete() {
139 let store = StateStore::new();
140 store.set("x", Value::from(1), "test");
141 let t = store.delete("x", "test");
142 assert!(t.is_some());
143 assert!(!store.exists("x"));
144 }
145
146 #[test]
147 fn delete_nonexistent() {
148 let store = StateStore::new();
149 assert!(store.delete("x", "test").is_none());
150 }
151
152 #[test]
153 fn snapshot_and_restore() {
154 let store = StateStore::new();
155 store.set("x", Value::from(1), "a");
156 let snap = store.snapshot();
157 let tc = store.transition_count();
158
159 store.set("y", Value::from(2), "b");
160 assert!(store.exists("y"));
161
162 store.restore(snap, tc);
163 assert!(store.exists("x"));
164 assert!(!store.exists("y"));
165 assert_eq!(store.transition_count(), 1);
166 }
167
168 #[test]
169 fn transitions_logged() {
170 let store = StateStore::new();
171 store.set("a", Value::from(1), "act1");
172 store.set("b", Value::from(2), "act2");
173
174 let transitions = store.transitions();
175 assert_eq!(transitions.len(), 2);
176 assert_eq!(transitions[0].key, "a");
177 assert_eq!(transitions[1].key, "b");
178 }
179
180 #[test]
181 fn transitions_since() {
182 let store = StateStore::new();
183 store.set("a", Value::from(1), "act1");
184 let idx = store.transition_count();
185 store.set("b", Value::from(2), "act2");
186
187 let since = store.transitions_since(idx);
188 assert_eq!(since.len(), 1);
189 assert_eq!(since[0].key, "b");
190 }
191
192 #[test]
193 fn transition_records_old_value() {
194 let store = StateStore::new();
195 store.set("x", Value::from(1), "first");
196 store.set("x", Value::from(2), "second");
197
198 let transitions = store.transitions();
199 assert_eq!(transitions[1].old_value, Some(Value::from(1)));
200 assert_eq!(transitions[1].new_value, Some(Value::from(2)));
201 }
202
203 #[test]
204 fn keys() {
205 let store = StateStore::new();
206 store.set("a", Value::from(1), "t");
207 store.set("b", Value::from(2), "t");
208 let mut keys = store.keys();
209 keys.sort();
210 assert_eq!(keys, vec!["a", "b"]);
211 }
212}