1use std::collections::VecDeque;
2use std::sync::Arc;
3use parking_lot::RwLock;
4use chrono::{DateTime, Local};
5use crate::ai::Message;
6
7const MAX_CONTEXT_MESSAGES: usize = 20;
8const CONTEXT_EXPIRY_HOURS: i64 = 2;
9
10pub struct ContextManager {
11 messages: Arc<RwLock<VecDeque<ContextEntry>>>,
12 max_messages: usize,
13 expiry_hours: i64,
14}
15
16#[derive(Debug, Clone)]
17pub struct ContextEntry {
18 pub message: Message,
19 pub timestamp: DateTime<Local>,
20 pub player: Option<String>,
21}
22
23impl ContextManager {
24 pub fn new() -> Self {
25 Self {
26 messages: Arc::new(RwLock::new(VecDeque::with_capacity(MAX_CONTEXT_MESSAGES))),
27 max_messages: MAX_CONTEXT_MESSAGES,
28 expiry_hours: CONTEXT_EXPIRY_HOURS,
29 }
30 }
31
32 pub fn add_message(&self, role: &str, content: &str, player: Option<&str>) {
33 let entry = ContextEntry {
34 message: Message {
35 role: role.to_string(),
36 content: content.to_string(),
37 },
38 timestamp: Local::now(),
39 player: player.map(|s| s.to_string()),
40 };
41
42 let mut messages = self.messages.write();
43 messages.push_back(entry);
44
45 while messages.len() > self.max_messages {
46 messages.pop_front();
47 }
48 }
49
50 pub fn add_user_message(&self, content: &str, player: &str) {
51 self.add_message("user", content, Some(player));
52 }
53
54 pub fn add_assistant_message(&self, content: &str) {
55 self.add_message("assistant", content, None);
56 }
57
58 pub fn add_system_message(&self, content: &str) {
59 self.add_message("system", content, None);
60 }
61
62 pub fn get_messages(&self) -> Vec<Message> {
63 let messages = self.messages.read();
64 let now = Local::now();
65
66 messages
67 .iter()
68 .filter(|entry| {
69 let elapsed = now.signed_duration_since(entry.timestamp);
70 elapsed.num_hours() < self.expiry_hours
71 })
72 .map(|entry| entry.message.clone())
73 .collect()
74 }
75
76 pub fn get_messages_for_player(&self, player: &str) -> Vec<Message> {
77 let messages = self.messages.read();
78 let now = Local::now();
79
80 messages
81 .iter()
82 .filter(|entry| {
83 let elapsed = now.signed_duration_since(entry.timestamp);
84 elapsed.num_hours() < self.expiry_hours
85 })
86 .filter(|entry| {
87 entry.player.as_deref() == Some(player) || entry.player.is_none()
88 })
89 .map(|entry| entry.message.clone())
90 .collect()
91 }
92
93 pub fn clear(&self) {
94 let mut messages = self.messages.write();
95 messages.clear();
96 }
97
98 pub fn len(&self) -> usize {
99 self.messages.read().len()
100 }
101
102 pub fn is_empty(&self) -> bool {
103 self.messages.read().is_empty()
104 }
105
106 pub fn get_recent_players(&self, count: usize) -> Vec<String> {
107 let messages = self.messages.read();
108 let mut players = Vec::new();
109
110 for entry in messages.iter().rev() {
111 if let Some(ref player) = entry.player {
112 if !players.contains(player) {
113 players.push(player.clone());
114 if players.len() >= count {
115 break;
116 }
117 }
118 }
119 }
120
121 players
122 }
123}
124
125impl Default for ContextManager {
126 fn default() -> Self {
127 Self::new()
128 }
129}