1use crate::error::{CsmError, Result};
6use crate::models::{ChatSession, ChatSessionIndex, ChatSessionIndexEntry};
7use crate::workspace::{get_empty_window_sessions_path, get_workspace_storage_path};
8use regex::Regex;
9use rusqlite::Connection;
10use std::path::{Path, PathBuf};
11use sysinfo::System;
12
13fn sanitize_json_unicode(content: &str) -> String {
16 let re = Regex::new(r"\\u[dD][89aAbB][0-9a-fA-F]{2}(?!\\u[dD][cCdDeEfF][0-9a-fA-F]{2})|(?<!\\u[dD][89aAbB][0-9a-fA-F]{2})\\u[dD][cCdDeEfF][0-9a-fA-F]{2}")
19 .unwrap();
20 re.replace_all(content, "\\uFFFD").to_string()
21}
22
23pub fn parse_session_json(content: &str) -> std::result::Result<ChatSession, serde_json::Error> {
25 match serde_json::from_str::<ChatSession>(content) {
26 Ok(session) => Ok(session),
27 Err(e) => {
28 if e.to_string().contains("surrogate") || e.to_string().contains("escape") {
30 let sanitized = sanitize_json_unicode(content);
31 serde_json::from_str::<ChatSession>(&sanitized)
32 } else {
33 Err(e)
34 }
35 }
36 }
37}
38
39pub fn get_workspace_storage_db(workspace_id: &str) -> Result<PathBuf> {
41 let storage_path = get_workspace_storage_path()?;
42 Ok(storage_path.join(workspace_id).join("state.vscdb"))
43}
44
45pub fn read_chat_session_index(db_path: &Path) -> Result<ChatSessionIndex> {
47 let conn = Connection::open(db_path)?;
48
49 let result: std::result::Result<String, rusqlite::Error> = conn.query_row(
50 "SELECT value FROM ItemTable WHERE key = ?",
51 ["chat.ChatSessionStore.index"],
52 |row| row.get(0),
53 );
54
55 match result {
56 Ok(json_str) => serde_json::from_str(&json_str)
57 .map_err(|e| CsmError::InvalidSessionFormat(e.to_string())),
58 Err(rusqlite::Error::QueryReturnedNoRows) => Ok(ChatSessionIndex::default()),
59 Err(e) => Err(CsmError::SqliteError(e)),
60 }
61}
62
63pub fn write_chat_session_index(db_path: &Path, index: &ChatSessionIndex) -> Result<()> {
65 let conn = Connection::open(db_path)?;
66 let json_str = serde_json::to_string(index)?;
67
68 let exists: bool = conn.query_row(
70 "SELECT COUNT(*) > 0 FROM ItemTable WHERE key = ?",
71 ["chat.ChatSessionStore.index"],
72 |row| row.get(0),
73 )?;
74
75 if exists {
76 conn.execute(
77 "UPDATE ItemTable SET value = ? WHERE key = ?",
78 [&json_str, "chat.ChatSessionStore.index"],
79 )?;
80 } else {
81 conn.execute(
82 "INSERT INTO ItemTable (key, value) VALUES (?, ?)",
83 ["chat.ChatSessionStore.index", &json_str],
84 )?;
85 }
86
87 Ok(())
88}
89
90pub fn add_session_to_index(
92 db_path: &Path,
93 session_id: &str,
94 title: &str,
95 last_message_date_ms: i64,
96 is_imported: bool,
97 initial_location: &str,
98 is_empty: bool,
99) -> Result<()> {
100 let mut index = read_chat_session_index(db_path)?;
101
102 index.entries.insert(
103 session_id.to_string(),
104 ChatSessionIndexEntry {
105 session_id: session_id.to_string(),
106 title: title.to_string(),
107 last_message_date: last_message_date_ms,
108 is_imported,
109 initial_location: initial_location.to_string(),
110 is_empty,
111 },
112 );
113
114 write_chat_session_index(db_path, &index)
115}
116
117pub fn remove_session_from_index(db_path: &Path, session_id: &str) -> Result<bool> {
119 let mut index = read_chat_session_index(db_path)?;
120 let removed = index.entries.remove(session_id).is_some();
121 if removed {
122 write_chat_session_index(db_path, &index)?;
123 }
124 Ok(removed)
125}
126
127pub fn sync_session_index(
129 workspace_id: &str,
130 chat_sessions_dir: &Path,
131 force: bool,
132) -> Result<(usize, usize)> {
133 let db_path = get_workspace_storage_db(workspace_id)?;
134
135 if !db_path.exists() {
136 return Err(CsmError::WorkspaceNotFound(format!(
137 "Database not found: {}",
138 db_path.display()
139 )));
140 }
141
142 if !force && is_vscode_running() {
144 return Err(CsmError::VSCodeRunning);
145 }
146
147 let mut index = read_chat_session_index(&db_path)?;
149
150 let mut files_on_disk: std::collections::HashSet<String> = std::collections::HashSet::new();
152 if chat_sessions_dir.exists() {
153 for entry in std::fs::read_dir(chat_sessions_dir)? {
154 let entry = entry?;
155 let path = entry.path();
156 if path.extension().map(|e| e == "json").unwrap_or(false) {
157 if let Some(stem) = path.file_stem() {
158 files_on_disk.insert(stem.to_string_lossy().to_string());
159 }
160 }
161 }
162 }
163
164 let stale_ids: Vec<String> = index
166 .entries
167 .keys()
168 .filter(|id| !files_on_disk.contains(*id))
169 .cloned()
170 .collect();
171
172 let removed = stale_ids.len();
173 for id in &stale_ids {
174 index.entries.remove(id);
175 }
176
177 let mut added = 0;
179 for entry in std::fs::read_dir(chat_sessions_dir)? {
180 let entry = entry?;
181 let path = entry.path();
182
183 if path.extension().map(|e| e == "json").unwrap_or(false) {
184 if let Ok(content) = std::fs::read_to_string(&path) {
185 if let Ok(session) = parse_session_json(&content) {
186 let session_id = session.session_id.clone().unwrap_or_else(|| {
187 path.file_stem()
188 .map(|s| s.to_string_lossy().to_string())
189 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
190 });
191
192 let title = session.title();
193 let is_empty = session.is_empty();
194 let last_message_date = session.last_message_date;
195 let initial_location = session.initial_location.clone();
196
197 index.entries.insert(
198 session_id.clone(),
199 ChatSessionIndexEntry {
200 session_id,
201 title,
202 last_message_date,
203 is_imported: session.is_imported,
204 initial_location,
205 is_empty,
206 },
207 );
208 added += 1;
209 }
210 }
211 }
212 }
213
214 write_chat_session_index(&db_path, &index)?;
216
217 Ok((added, removed))
218}
219
220pub fn register_all_sessions_from_directory(
222 workspace_id: &str,
223 chat_sessions_dir: &Path,
224 force: bool,
225) -> Result<usize> {
226 let db_path = get_workspace_storage_db(workspace_id)?;
227
228 if !db_path.exists() {
229 return Err(CsmError::WorkspaceNotFound(format!(
230 "Database not found: {}",
231 db_path.display()
232 )));
233 }
234
235 if !force && is_vscode_running() {
237 return Err(CsmError::VSCodeRunning);
238 }
239
240 let (added, removed) = sync_session_index(workspace_id, chat_sessions_dir, force)?;
242
243 for entry in std::fs::read_dir(chat_sessions_dir)? {
245 let entry = entry?;
246 let path = entry.path();
247
248 if path.extension().map(|e| e == "json").unwrap_or(false) {
249 if let Ok(content) = std::fs::read_to_string(&path) {
250 if let Ok(session) = parse_session_json(&content) {
251 let session_id = session.session_id.clone().unwrap_or_else(|| {
252 path.file_stem()
253 .map(|s| s.to_string_lossy().to_string())
254 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string())
255 });
256
257 let title = session.title();
258
259 println!(
260 "[OK] Registered: {} ({}...)",
261 title,
262 &session_id[..12.min(session_id.len())]
263 );
264 }
265 }
266 }
267 }
268
269 if removed > 0 {
270 println!("[OK] Removed {} stale index entries", removed);
271 }
272
273 Ok(added)
274}
275
276pub fn is_vscode_running() -> bool {
278 let mut sys = System::new();
279 sys.refresh_processes();
280
281 for process in sys.processes().values() {
282 let name = process.name().to_lowercase();
283 if name.contains("code") && !name.contains("codec") {
284 return true;
285 }
286 }
287
288 false
289}
290
291pub fn backup_workspace_sessions(workspace_dir: &Path) -> Result<Option<PathBuf>> {
293 let chat_sessions_dir = workspace_dir.join("chatSessions");
294
295 if !chat_sessions_dir.exists() {
296 return Ok(None);
297 }
298
299 let timestamp = std::time::SystemTime::now()
300 .duration_since(std::time::UNIX_EPOCH)
301 .unwrap()
302 .as_secs();
303
304 let backup_dir = workspace_dir.join(format!("chatSessions-backup-{}", timestamp));
305
306 copy_dir_all(&chat_sessions_dir, &backup_dir)?;
308
309 Ok(Some(backup_dir))
310}
311
312fn copy_dir_all(src: &Path, dst: &Path) -> Result<()> {
314 std::fs::create_dir_all(dst)?;
315
316 for entry in std::fs::read_dir(src)? {
317 let entry = entry?;
318 let src_path = entry.path();
319 let dst_path = dst.join(entry.file_name());
320
321 if src_path.is_dir() {
322 copy_dir_all(&src_path, &dst_path)?;
323 } else {
324 std::fs::copy(&src_path, &dst_path)?;
325 }
326 }
327
328 Ok(())
329}
330
331pub fn read_empty_window_sessions() -> Result<Vec<ChatSession>> {
338 let sessions_path = get_empty_window_sessions_path()?;
339
340 if !sessions_path.exists() {
341 return Ok(Vec::new());
342 }
343
344 let mut sessions = Vec::new();
345
346 for entry in std::fs::read_dir(&sessions_path)? {
347 let entry = entry?;
348 let path = entry.path();
349
350 if path.extension().is_some_and(|e| e == "json") {
351 if let Ok(content) = std::fs::read_to_string(&path) {
352 if let Ok(session) = parse_session_json(&content) {
353 sessions.push(session);
354 }
355 }
356 }
357 }
358
359 sessions.sort_by(|a, b| b.last_message_date.cmp(&a.last_message_date));
361
362 Ok(sessions)
363}
364
365#[allow(dead_code)]
367pub fn get_empty_window_session(session_id: &str) -> Result<Option<ChatSession>> {
368 let sessions_path = get_empty_window_sessions_path()?;
369 let session_path = sessions_path.join(format!("{}.json", session_id));
370
371 if !session_path.exists() {
372 return Ok(None);
373 }
374
375 let content = std::fs::read_to_string(&session_path)?;
376 let session: ChatSession = serde_json::from_str(&content)
377 .map_err(|e| CsmError::InvalidSessionFormat(e.to_string()))?;
378
379 Ok(Some(session))
380}
381
382#[allow(dead_code)]
384pub fn write_empty_window_session(session: &ChatSession) -> Result<PathBuf> {
385 let sessions_path = get_empty_window_sessions_path()?;
386
387 std::fs::create_dir_all(&sessions_path)?;
389
390 let session_id = session.session_id.as_deref().unwrap_or("unknown");
391 let session_path = sessions_path.join(format!("{}.json", session_id));
392 let content = serde_json::to_string_pretty(session)?;
393 std::fs::write(&session_path, content)?;
394
395 Ok(session_path)
396}
397
398#[allow(dead_code)]
400pub fn delete_empty_window_session(session_id: &str) -> Result<bool> {
401 let sessions_path = get_empty_window_sessions_path()?;
402 let session_path = sessions_path.join(format!("{}.json", session_id));
403
404 if session_path.exists() {
405 std::fs::remove_file(&session_path)?;
406 Ok(true)
407 } else {
408 Ok(false)
409 }
410}
411
412pub fn count_empty_window_sessions() -> Result<usize> {
414 let sessions_path = get_empty_window_sessions_path()?;
415
416 if !sessions_path.exists() {
417 return Ok(0);
418 }
419
420 let count = std::fs::read_dir(&sessions_path)?
421 .filter_map(|e| e.ok())
422 .filter(|e| e.path().extension().is_some_and(|ext| ext == "json"))
423 .count();
424
425 Ok(count)
426}