1use super::MessageStore;
4use super::types::*;
5use super::user_handle::UserHandle;
6use anyhow::Result;
7use regex::Regex;
8use serde::{Deserialize, Serialize};
9use std::collections::{HashMap, HashSet};
10use std::sync::Arc;
11use tokio::sync::RwLock;
12
13pub struct MessageSearch {
15 store: MessageStore,
16 search_index: Arc<RwLock<SearchIndex>>,
18 recent_searches: Arc<RwLock<Vec<SearchQuery>>>,
20}
21
22impl MessageSearch {
23 pub async fn new(store: MessageStore) -> Result<Self> {
25 let search_index = SearchIndex::new();
26
27 Ok(Self {
28 store,
29 search_index: Arc::new(RwLock::new(search_index)),
30 recent_searches: Arc::new(RwLock::new(Vec::new())),
31 })
32 }
33
34 pub async fn search(&self, query: SearchQuery) -> Result<Vec<RichMessage>> {
36 self.add_to_recent(query.clone()).await;
38
39 let filters = self.build_filters(&query);
41
42 let message_ids = self.search_index(&query, &filters).await?;
44
45 let mut messages = Vec::new();
47 for id in message_ids {
48 if let Ok(msg) = self.store.get_message(id).await
49 && self.matches_query(&msg, &query, &filters)
50 {
51 messages.push(msg);
52 }
53 }
54
55 messages.sort_by(|a, b| b.created_at.cmp(&a.created_at));
57
58 messages.truncate(query.limit);
60
61 Ok(messages)
62 }
63
64 pub async fn quick_search(&self, text: String, limit: usize) -> Result<Vec<RichMessage>> {
66 let query = SearchQuery {
67 text: Some(text),
68 from: None,
69 in_channels: None,
70 has_attachments: None,
71 has_reactions: None,
72 is_thread: None,
73 date_range: None,
74 limit,
75 };
76
77 self.search(query).await
78 }
79
80 pub async fn search_channel(
82 &self,
83 channel_id: ChannelId,
84 text: Option<String>,
85 limit: usize,
86 ) -> Result<Vec<RichMessage>> {
87 let query = SearchQuery {
88 text,
89 from: None,
90 in_channels: Some(vec![channel_id]),
91 has_attachments: None,
92 has_reactions: None,
93 is_thread: None,
94 date_range: None,
95 limit,
96 };
97
98 self.search(query).await
99 }
100
101 pub async fn search_from_users(
103 &self,
104 users: Vec<UserHandle>,
105 text: Option<String>,
106 limit: usize,
107 ) -> Result<Vec<RichMessage>> {
108 let query = SearchQuery {
109 text,
110 from: Some(users),
111 in_channels: None,
112 has_attachments: None,
113 has_reactions: None,
114 is_thread: None,
115 date_range: None,
116 limit,
117 };
118
119 self.search(query).await
120 }
121
122 pub async fn search_attachments(
124 &self,
125 mime_type: Option<String>,
126 limit: usize,
127 ) -> Result<Vec<RichMessage>> {
128 let query = SearchQuery {
129 text: mime_type,
130 from: None,
131 in_channels: None,
132 has_attachments: Some(true),
133 has_reactions: None,
134 is_thread: None,
135 date_range: None,
136 limit,
137 };
138
139 self.search(query).await
140 }
141
142 pub async fn regex_search(&self, pattern: &str, limit: usize) -> Result<Vec<RichMessage>> {
144 let regex = Regex::new(pattern)?;
145
146 let query = SearchQuery {
148 text: Some(pattern.to_string()),
149 from: None,
150 in_channels: None,
151 has_attachments: None,
152 has_reactions: None,
153 is_thread: None,
154 date_range: None,
155 limit: limit * 2, };
157
158 let mut messages = self.search(query).await?;
159
160 messages.retain(|msg| match &msg.content {
162 MessageContent::Text(text) => regex.is_match(text),
163 MessageContent::RichText(rich) => regex.is_match(&rich.raw),
164 _ => false,
165 });
166
167 messages.truncate(limit);
168 Ok(messages)
169 }
170
171 pub async fn get_suggestions(&self, partial: &str) -> Result<Vec<SearchSuggestion>> {
173 let mut suggestions = Vec::new();
174
175 let recent = self.recent_searches.read().await;
177 for query in recent.iter() {
178 if let Some(text) = &query.text
179 && text.starts_with(partial)
180 {
181 suggestions.push(SearchSuggestion {
182 text: text.clone(),
183 category: SuggestionCategory::Recent,
184 icon: "🕐".to_string(),
185 });
186 }
187 }
188
189 let index = self.search_index.read().await;
191 for user in &index.known_users {
192 let user_str = user.to_string();
193 if user_str.contains(partial) {
194 suggestions.push(SearchSuggestion {
195 text: format!("from:{}", user_str),
196 category: SuggestionCategory::User,
197 icon: "👤".to_string(),
198 });
199 }
200 }
201
202 for channel in &index.known_channels {
204 suggestions.push(SearchSuggestion {
205 text: format!("in:channel-{}", channel.0),
206 category: SuggestionCategory::Channel,
207 icon: "#".to_string(),
208 });
209 }
210
211 if "has:".starts_with(partial) {
213 suggestions.push(SearchSuggestion {
214 text: "has:attachment".to_string(),
215 category: SuggestionCategory::Filter,
216 icon: "📎".to_string(),
217 });
218 suggestions.push(SearchSuggestion {
219 text: "has:reaction".to_string(),
220 category: SuggestionCategory::Filter,
221 icon: "😊".to_string(),
222 });
223 }
224
225 Ok(suggestions)
226 }
227
228 pub async fn get_recent_searches(&self) -> Vec<SearchQuery> {
230 let recent = self.recent_searches.read().await;
231 recent.clone()
232 }
233
234 pub async fn clear_history(&self) -> Result<()> {
236 let mut recent = self.recent_searches.write().await;
237 recent.clear();
238 Ok(())
239 }
240
241 pub async fn index_message(&self, message: &RichMessage) -> Result<()> {
243 let mut index = self.search_index.write().await;
244 index.add_message(message);
245 Ok(())
246 }
247
248 fn build_filters(&self, query: &SearchQuery) -> SearchFilters {
250 SearchFilters {
251 text_tokens: query
252 .text
253 .as_ref()
254 .map(|t| t.split_whitespace().map(|s| s.to_lowercase()).collect())
255 .unwrap_or_default(),
256 from_users: query.from.clone().unwrap_or_default(),
257 channels: query.in_channels.clone().unwrap_or_default(),
258 has_attachments: query.has_attachments,
259 has_reactions: query.has_reactions,
260 is_thread: query.is_thread,
261 date_range: query.date_range.clone(),
262 }
263 }
264
265 async fn search_index(
267 &self,
268 query: &SearchQuery,
269 _filters: &SearchFilters,
270 ) -> Result<Vec<MessageId>> {
271 let index = self.search_index.read().await;
272
273 let mut results = Vec::new();
274
275 if let Some(text) = &query.text
277 && let Some(msg_ids) = index.text_index.get(&text.to_lowercase())
278 {
279 results.extend(msg_ids.iter().copied());
280 }
281
282 if results.is_empty() && query.text.is_none() {
284 }
287
288 Ok(results)
289 }
290
291 fn matches_query(
293 &self,
294 message: &RichMessage,
295 _query: &SearchQuery,
296 filters: &SearchFilters,
297 ) -> bool {
298 if !filters.text_tokens.is_empty() {
300 let content_text = match &message.content {
301 MessageContent::Text(t) => t.clone(),
302 MessageContent::RichText(r) => r.raw.clone(),
303 _ => String::new(),
304 };
305
306 let content_lower = content_text.to_lowercase();
307 if !filters
308 .text_tokens
309 .iter()
310 .all(|token| content_lower.contains(token))
311 {
312 return false;
313 }
314 }
315
316 if !filters.from_users.is_empty() && !filters.from_users.contains(&message.sender) {
318 return false;
319 }
320
321 if !filters.channels.is_empty() && !filters.channels.contains(&message.channel_id) {
323 return false;
324 }
325
326 if let Some(has_attach) = filters.has_attachments
328 && has_attach == message.attachments.is_empty()
329 {
330 return false;
331 }
332
333 if let Some(has_react) = filters.has_reactions
335 && has_react == message.reactions.is_empty()
336 {
337 return false;
338 }
339
340 if let Some(is_thread) = filters.is_thread
342 && is_thread != message.thread_id.is_some()
343 {
344 return false;
345 }
346
347 if let Some(range) = &filters.date_range
349 && (message.created_at < range.start || message.created_at > range.end)
350 {
351 return false;
352 }
353
354 true
355 }
356
357 async fn add_to_recent(&self, query: SearchQuery) {
359 let mut recent = self.recent_searches.write().await;
360
361 recent.retain(|q| q.text != query.text);
363
364 recent.insert(0, query);
366
367 recent.truncate(10);
369 }
370}
371
372struct SearchIndex {
374 text_index: HashMap<String, HashSet<MessageId>>,
376 known_users: HashSet<UserHandle>,
378 known_channels: HashSet<ChannelId>,
380}
381
382impl SearchIndex {
383 fn new() -> Self {
384 Self {
385 text_index: HashMap::new(),
386 known_users: HashSet::new(),
387 known_channels: HashSet::new(),
388 }
389 }
390
391 fn add_message(&mut self, message: &RichMessage) {
392 let text = match &message.content {
394 MessageContent::Text(t) => t.clone(),
395 MessageContent::RichText(r) => r.raw.clone(),
396 _ => String::new(),
397 };
398
399 for word in text.split_whitespace() {
401 let token = word.to_lowercase();
402 self.text_index.entry(token).or_default().insert(message.id);
403 }
404
405 self.known_users.insert(message.sender.clone());
407 self.known_channels.insert(message.channel_id);
408 }
409}
410
411struct SearchFilters {
413 text_tokens: Vec<String>,
414 from_users: Vec<UserHandle>,
415 channels: Vec<ChannelId>,
416 has_attachments: Option<bool>,
417 has_reactions: Option<bool>,
418 is_thread: Option<bool>,
419 date_range: Option<DateRange>,
420}
421
422#[derive(Debug, Clone, Serialize, Deserialize)]
424pub struct SearchSuggestion {
425 pub text: String,
426 pub category: SuggestionCategory,
427 pub icon: String,
428}
429
430#[derive(Debug, Clone, Serialize, Deserialize)]
432pub enum SuggestionCategory {
433 Recent,
434 User,
435 Channel,
436 Filter,
437 Command,
438}
439
440#[cfg(test)]
441mod tests {
442 use super::*;
443
444 #[tokio::test]
445 async fn test_search_creation() {
446 #[allow(unused)]
447 let store = super::super::database::DatabaseMessageStore::new(
448 super::super::DhtClient::new_mock(),
449 None,
450 )
451 .await
452 .unwrap();
453 let search = MessageSearch::new(store).await.unwrap();
454
455 let recent = search.get_recent_searches().await;
456 assert_eq!(recent.len(), 0);
457 }
458
459 #[tokio::test]
460 async fn test_quick_search() {
461 #[allow(unused)]
462 let store = super::super::database::DatabaseMessageStore::new(
463 super::super::DhtClient::new_mock(),
464 None,
465 )
466 .await
467 .unwrap();
468 let search = MessageSearch::new(store).await.unwrap();
469
470 let results = search.quick_search("test".to_string(), 10).await.unwrap();
471 assert_eq!(results.len(), 0); }
473
474 #[tokio::test]
475 async fn test_search_suggestions() {
476 #[allow(unused)]
477 let store = super::super::database::DatabaseMessageStore::new(
478 super::super::DhtClient::new_mock(),
479 None,
480 )
481 .await
482 .unwrap();
483 let search = MessageSearch::new(store).await.unwrap();
484
485 let suggestions = search.get_suggestions("has").await.unwrap();
486 assert!(suggestions.iter().any(|s| s.text == "has:attachment"));
487 assert!(suggestions.iter().any(|s| s.text == "has:reaction"));
488 }
489}