1use crate::domains::Domain;
9use crate::persona::{PersonaGenerator, PersonaProfile, PersonaRegistry};
10use crate::persona_graph::PersonaGraph;
11use crate::Result;
12use serde_json::Value;
13use std::collections::HashMap;
14use std::sync::{Arc, RwLock};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
21pub enum EntityType {
22 User,
24 Device,
26 Organization,
28 Generic,
30}
31
32impl EntityType {
33 pub fn as_str(&self) -> &'static str {
35 match self {
36 EntityType::User => "user",
37 EntityType::Device => "device",
38 EntityType::Organization => "organization",
39 EntityType::Generic => "generic",
40 }
41 }
42
43 pub fn parse(s: &str) -> Self {
45 match s.to_lowercase().as_str() {
46 "user" | "users" => EntityType::User,
47 "device" | "devices" => EntityType::Device,
48 "organization" | "organizations" | "org" | "orgs" => EntityType::Organization,
49 _ => EntityType::Generic,
50 }
51 }
52}
53
54#[derive(Debug)]
61pub struct ConsistencyStore {
62 persona_registry: Arc<PersonaRegistry>,
64 generators: Arc<RwLock<HashMap<Domain, PersonaGenerator>>>,
66 default_domain: Option<Domain>,
68 persona_graph: Arc<PersonaGraph>,
70}
71
72impl ConsistencyStore {
73 pub fn new() -> Self {
75 Self {
76 persona_registry: Arc::new(PersonaRegistry::new()),
77 generators: Arc::new(RwLock::new(HashMap::new())),
78 default_domain: None,
79 persona_graph: Arc::new(PersonaGraph::new()),
80 }
81 }
82
83 pub fn with_default_domain(default_domain: Domain) -> Self {
85 Self {
86 persona_registry: Arc::new(PersonaRegistry::new()),
87 generators: Arc::new(RwLock::new(HashMap::new())),
88 default_domain: Some(default_domain),
89 persona_graph: Arc::new(PersonaGraph::new()),
90 }
91 }
92
93 pub fn with_registry_and_domain(
95 persona_registry: Arc<PersonaRegistry>,
96 default_domain: Option<Domain>,
97 ) -> Self {
98 Self {
99 persona_registry,
100 generators: Arc::new(RwLock::new(HashMap::new())),
101 default_domain,
102 persona_graph: Arc::new(PersonaGraph::new()),
103 }
104 }
105
106 pub fn with_persona_graph(persona_graph: Arc<PersonaGraph>) -> Self {
108 Self {
109 persona_registry: Arc::new(PersonaRegistry::new()),
110 generators: Arc::new(RwLock::new(HashMap::new())),
111 default_domain: None,
112 persona_graph,
113 }
114 }
115
116 pub fn get_entity_persona(&self, entity_id: &str, domain: Option<Domain>) -> PersonaProfile {
121 let domain = domain.or(self.default_domain).unwrap_or(Domain::General);
122 self.persona_registry.get_or_create_persona(entity_id.to_string(), domain)
123 }
124
125 pub fn get_or_create_persona_by_type(
134 &self,
135 entity_id: &str,
136 entity_type: EntityType,
137 domain: Option<Domain>,
138 ) -> PersonaProfile {
139 let domain = domain.or(self.default_domain).unwrap_or(Domain::General);
140 let persona_id = format!("{}:{}", entity_type.as_str(), entity_id);
141 let persona = self.persona_registry.get_or_create_persona(persona_id.clone(), domain);
142
143 let entity_type_str = entity_type.as_str();
145 self.persona_graph
146 .get_or_create_node_with_links(&persona_id, entity_type_str, None, None);
147
148 if entity_type != EntityType::Generic {
150 let base_persona = self.get_entity_persona(entity_id, Some(domain));
152 let base_persona_id = base_persona.id.clone();
153
154 self.persona_graph
156 .get_or_create_node_with_links(&base_persona_id, "base", None, None);
157
158 let mut base_persona_mut = base_persona.clone();
160 match entity_type {
161 EntityType::User => {
162 self.persona_graph.link_entity_types(
166 &base_persona_id,
167 "base",
168 &persona_id,
169 entity_type_str,
170 );
171 }
172 EntityType::Device => {
173 base_persona_mut
175 .add_relationship("owns_devices".to_string(), persona_id.clone());
176 self.persona_graph.link_entity_types(
178 &base_persona_id,
179 "base",
180 &persona_id,
181 entity_type_str,
182 );
183 }
184 EntityType::Organization => {
185 base_persona_mut.add_relationship("has_users".to_string(), persona_id.clone());
187 self.persona_graph.link_entity_types(
189 &base_persona_id,
190 "base",
191 &persona_id,
192 entity_type_str,
193 );
194 }
195 EntityType::Generic => {}
196 }
197
198 for (rel_type, related_ids) in &base_persona_mut.relationships {
201 for related_id in related_ids {
202 if let Some(existing) = self.persona_registry.get_persona(entity_id) {
204 if !existing.get_related_personas(rel_type).contains(related_id) {
205 self.persona_registry
206 .add_relationship(entity_id, rel_type.clone(), related_id.clone())
207 .ok();
208 }
209 }
210 }
211 }
212 }
213
214 persona
215 }
216
217 pub fn link_personas(
222 &self,
223 from_entity_id: &str,
224 from_entity_type: &str,
225 to_entity_id: &str,
226 to_entity_type: &str,
227 ) {
228 let from_persona_id = format!("{}:{}", from_entity_type, from_entity_id);
229 let to_persona_id = format!("{}:{}", to_entity_type, to_entity_id);
230
231 self.persona_graph.get_or_create_node_with_links(
233 &from_persona_id,
234 from_entity_type,
235 None,
236 None,
237 );
238 self.persona_graph.get_or_create_node_with_links(
239 &to_persona_id,
240 to_entity_type,
241 None,
242 None,
243 );
244
245 self.persona_graph.link_entity_types(
247 &from_persona_id,
248 from_entity_type,
249 &to_persona_id,
250 to_entity_type,
251 );
252 }
253
254 pub fn persona_graph(&self) -> &Arc<PersonaGraph> {
256 &self.persona_graph
257 }
258
259 pub fn get_personas_for_base_id(
263 &self,
264 base_id: &str,
265 domain: Option<Domain>,
266 ) -> Vec<PersonaProfile> {
267 let domain = domain.or(self.default_domain).unwrap_or(Domain::General);
268 let mut personas = Vec::new();
269
270 let base_persona = self.get_entity_persona(base_id, Some(domain));
272 personas.push(base_persona);
273
274 for entity_type in [
276 EntityType::User,
277 EntityType::Device,
278 EntityType::Organization,
279 ] {
280 let persona_id = format!("{}:{}", entity_type.as_str(), base_id);
281 if let Some(persona) = self.persona_registry.get_persona(&persona_id) {
282 personas.push(persona);
283 }
284 }
285
286 personas
287 }
288
289 pub fn generate_consistent_value(
294 &self,
295 entity_id: &str,
296 field_type: &str,
297 domain: Option<Domain>,
298 ) -> Result<Value> {
299 self.generate_consistent_value_with_reality(entity_id, field_type, domain, 0.0, None, None)
301 }
302
303 pub fn generate_consistent_value_with_reality(
316 &self,
317 entity_id: &str,
318 field_type: &str,
319 domain: Option<Domain>,
320 reality_ratio: f64,
321 recorded_data: Option<&Value>,
322 real_data: Option<&Value>,
323 ) -> Result<Value> {
324 let persona = self.get_entity_persona(entity_id, domain);
326 let domain = persona.domain;
327
328 let generator = {
332 let generators = self.generators.read().map_err(|e| {
333 crate::Error::LockPoisoned(format!("consistency generators read lock: {}", e))
334 })?;
335 if generators.contains_key(&domain) {
336 PersonaGenerator::new(domain)
339 } else {
340 drop(generators);
341 let mut generators = self.generators.write().map_err(|e| {
342 crate::Error::LockPoisoned(format!("consistency generators write lock: {}", e))
343 })?;
344 generators.insert(domain, PersonaGenerator::new(domain));
345 PersonaGenerator::new(domain)
346 }
347 };
348
349 generator.generate_for_persona_with_reality(
351 &persona,
352 field_type,
353 reality_ratio,
354 recorded_data,
355 real_data,
356 )
357 }
358
359 pub fn persona_registry(&self) -> &Arc<PersonaRegistry> {
361 &self.persona_registry
362 }
363
364 pub fn set_default_domain(&mut self, domain: Option<Domain>) {
366 self.default_domain = domain;
367 }
368
369 pub fn default_domain(&self) -> Option<Domain> {
371 self.default_domain
372 }
373
374 pub fn clear(&self) {
376 self.persona_registry.clear();
377 let mut generators =
378 self.generators.write().expect("consistency generators write lock poisoned");
379 generators.clear();
380 }
381
382 pub fn persona_count(&self) -> usize {
384 self.persona_registry.count()
385 }
386}
387
388impl Default for ConsistencyStore {
389 fn default() -> Self {
390 Self::new()
391 }
392}
393
394pub struct EntityIdExtractor;
399
400impl EntityIdExtractor {
401 pub fn from_field_name(field_name: &str) -> Option<String> {
406 let field_lower = field_name.to_lowercase();
407
408 let patterns = [
410 "user_id",
411 "userid",
412 "user-id",
413 "device_id",
414 "deviceid",
415 "device-id",
416 "transaction_id",
417 "transactionid",
418 "transaction-id",
419 "order_id",
420 "orderid",
421 "order-id",
422 "customer_id",
423 "customerid",
424 "customer-id",
425 "patient_id",
426 "patientid",
427 "patient-id",
428 "account_id",
429 "accountid",
430 "account-id",
431 "id", ];
433
434 for pattern in &patterns {
436 if field_lower == *pattern {
437 return Some(field_name.to_string());
438 }
439 }
440
441 for pattern in &patterns {
443 if field_lower.ends_with(&format!("_{}", pattern))
444 || field_lower.ends_with(&format!("-{}", pattern))
445 {
446 return Some(field_name.to_string());
447 }
448 }
449
450 None
451 }
452
453 pub fn from_path(path: &str) -> Option<(String, EntityType)> {
458 let segments: Vec<&str> = path.split('/').collect();
463 if segments.len() >= 3 {
464 let resource = segments[segments.len() - 2].to_lowercase();
465 let id = segments[segments.len() - 1];
466
467 let entity_type = EntityType::parse(&resource);
469
470 if entity_type != EntityType::Generic && !id.is_empty() {
471 return Some((id.to_string(), entity_type));
472 }
473
474 let entity_types = [
476 "transaction",
477 "transactions",
478 "order",
479 "orders",
480 "customer",
481 "customers",
482 "patient",
483 "patients",
484 "account",
485 "accounts",
486 ];
487
488 if entity_types.contains(&resource.as_str()) && !id.is_empty() {
489 return Some((id.to_string(), EntityType::Generic));
490 }
491 }
492
493 None
494 }
495
496 pub fn from_path_id_only(path: &str) -> Option<String> {
500 Self::from_path(path).map(|(id, _)| id)
501 }
502
503 pub fn from_json_value(value: &Value) -> Option<String> {
507 if let Some(obj) = value.as_object() {
508 let id_fields = [
510 "user_id",
511 "userId",
512 "user-id",
513 "device_id",
514 "deviceId",
515 "device-id",
516 "transaction_id",
517 "transactionId",
518 "transaction-id",
519 "order_id",
520 "orderId",
521 "order-id",
522 "customer_id",
523 "customerId",
524 "customer-id",
525 "patient_id",
526 "patientId",
527 "patient-id",
528 "account_id",
529 "accountId",
530 "account-id",
531 "id",
532 ];
533
534 for field in &id_fields {
535 if let Some(id_value) = obj.get(*field) {
536 if let Some(id_str) = id_value.as_str() {
537 return Some(id_str.to_string());
538 } else if let Some(id_num) = id_value.as_u64() {
539 return Some(id_num.to_string());
540 }
541 }
542 }
543 }
544
545 None
546 }
547
548 pub fn from_multiple_sources(
552 field_name: Option<&str>,
553 path: Option<&str>,
554 json_value: Option<&Value>,
555 ) -> Option<String> {
556 if let Some(field) = field_name {
558 if let Some(id) = Self::from_field_name(field) {
559 return Some(id);
560 }
561 }
562
563 if let Some(p) = path {
565 if let Some((id, _)) = Self::from_path(p) {
566 return Some(id);
567 }
568 }
569
570 if let Some(json) = json_value {
572 if let Some(id) = Self::from_json_value(json) {
573 return Some(id);
574 }
575 }
576
577 None
578 }
579}
580
581#[cfg(test)]
582mod tests {
583 use super::*;
584 use serde_json::json;
585
586 #[test]
587 fn test_consistency_store_new() {
588 let store = ConsistencyStore::new();
589 assert_eq!(store.persona_count(), 0);
590 }
591
592 #[test]
593 fn test_consistency_store_with_default_domain() {
594 let store = ConsistencyStore::with_default_domain(Domain::Finance);
595 assert_eq!(store.default_domain(), Some(Domain::Finance));
596 }
597
598 #[test]
599 fn test_get_entity_persona() {
600 let store = ConsistencyStore::with_default_domain(Domain::Finance);
601 let persona1 = store.get_entity_persona("user123", None);
602 let persona2 = store.get_entity_persona("user123", None);
603
604 assert_eq!(persona1.id, persona2.id);
606 assert_eq!(persona1.seed, persona2.seed);
607 }
608
609 #[test]
610 fn test_generate_consistent_value() {
611 let store = ConsistencyStore::with_default_domain(Domain::Finance);
612
613 let value1 = store.generate_consistent_value("user123", "amount", None).unwrap();
615 let value2 = store.generate_consistent_value("user123", "amount", None).unwrap();
616
617 assert!(value1.is_string() || value1.is_number());
619 assert!(value2.is_string() || value2.is_number());
620 }
621
622 #[test]
623 fn test_entity_id_extractor_from_field_name() {
624 assert_eq!(EntityIdExtractor::from_field_name("user_id"), Some("user_id".to_string()));
625 assert_eq!(EntityIdExtractor::from_field_name("deviceId"), Some("deviceId".to_string()));
626 assert_eq!(
627 EntityIdExtractor::from_field_name("transaction_id"),
628 Some("transaction_id".to_string())
629 );
630 assert_eq!(EntityIdExtractor::from_field_name("name"), None);
631 }
632
633 #[test]
634 fn test_entity_id_extractor_from_path() {
635 assert_eq!(
636 EntityIdExtractor::from_path("/users/123"),
637 Some(("123".to_string(), EntityType::User))
638 );
639 assert_eq!(
640 EntityIdExtractor::from_path("/devices/abc-123"),
641 Some(("abc-123".to_string(), EntityType::Device))
642 );
643 assert_eq!(
644 EntityIdExtractor::from_path("/organizations/org1"),
645 Some(("org1".to_string(), EntityType::Organization))
646 );
647 assert_eq!(EntityIdExtractor::from_path("/api/health"), None);
648 }
649
650 #[test]
651 fn test_entity_id_extractor_from_path_id_only() {
652 assert_eq!(EntityIdExtractor::from_path_id_only("/users/123"), Some("123".to_string()));
653 assert_eq!(
654 EntityIdExtractor::from_path_id_only("/devices/abc-123"),
655 Some("abc-123".to_string())
656 );
657 assert_eq!(
658 EntityIdExtractor::from_path_id_only("/organizations/org1"),
659 Some("org1".to_string())
660 );
661 assert_eq!(EntityIdExtractor::from_path_id_only("/api/health"), None);
662 }
663
664 #[test]
665 fn test_entity_type() {
666 assert_eq!(EntityType::User.as_str(), "user");
667 assert_eq!(EntityType::Device.as_str(), "device");
668 assert_eq!(EntityType::Organization.as_str(), "organization");
669 assert_eq!(EntityType::parse("users"), EntityType::User);
670 assert_eq!(EntityType::parse("devices"), EntityType::Device);
671 assert_eq!(EntityType::parse("organizations"), EntityType::Organization);
672 }
673
674 #[test]
675 fn test_get_or_create_persona_by_type() {
676 let store = ConsistencyStore::with_default_domain(Domain::Finance);
677
678 let user_persona = store.get_or_create_persona_by_type("user123", EntityType::User, None);
680 let device_persona =
681 store.get_or_create_persona_by_type("user123", EntityType::Device, None);
682
683 assert_ne!(user_persona.id, device_persona.id);
685 assert!(user_persona.id.contains("user:user123"));
686 assert!(device_persona.id.contains("device:user123"));
687 }
688
689 #[test]
690 fn test_get_personas_for_base_id() {
691 let store = ConsistencyStore::with_default_domain(Domain::Finance);
692
693 store.get_or_create_persona_by_type("user123", EntityType::User, None);
695 store.get_or_create_persona_by_type("user123", EntityType::Device, None);
696
697 let personas = store.get_personas_for_base_id("user123", None);
698 assert!(personas.len() >= 2); }
700
701 #[test]
702 fn test_entity_id_extractor_from_json_value() {
703 let json = json!({
704 "user_id": "user123",
705 "name": "John Doe"
706 });
707 assert_eq!(EntityIdExtractor::from_json_value(&json), Some("user123".to_string()));
708
709 let json2 = json!({
710 "id": 456,
711 "name": "Device"
712 });
713 assert_eq!(EntityIdExtractor::from_json_value(&json2), Some("456".to_string()));
714 }
715
716 #[test]
717 fn test_entity_id_extractor_from_multiple_sources() {
718 let id1 = EntityIdExtractor::from_multiple_sources(Some("user_id"), None, None);
720 assert_eq!(id1, Some("user_id".to_string()));
721
722 let id2 = EntityIdExtractor::from_multiple_sources(Some("name"), Some("/users/123"), None);
724 assert_eq!(id2, Some("123".to_string()));
725
726 let json = json!({"user_id": "user456"});
728 let id3 = EntityIdExtractor::from_multiple_sources(
729 Some("name"),
730 Some("/api/health"),
731 Some(&json),
732 );
733 assert_eq!(id3, Some("user456".to_string()));
734 }
735}