1use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::{mpsc, RwLock};
10use tracing::{debug, info};
11
12use crate::core::{CoderLib, CoderLibError, CodeRequest, CodeResponse};
13use crate::integration::{
14 EditHost, ContextGatherer, ContextConfig, GatheredContext,
15 HostCommand
16};
17use crate::storage::{Message, MessageRole, MessageContent};
18
19pub struct AIAssistant {
21 coderlib: Arc<CoderLib>,
23 edit_host: Arc<EditHost>,
25 context_gatherer: ContextGatherer,
27 sessions: Arc<RwLock<HashMap<String, AssistantSession>>>,
29 ui_state: Arc<RwLock<UIState>>,
31 command_sender: Option<mpsc::UnboundedSender<HostCommand>>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct AssistantSession {
38 pub id: String,
40 pub session_type: SessionType,
42 pub messages: Vec<Message>,
44 pub context: Option<GatheredContext>,
46 pub state: SessionState,
48 pub created_at: std::time::SystemTime,
50 pub last_activity: std::time::SystemTime,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
56pub enum SessionType {
57 CodeAssistance,
59 CodeExplanation,
61 CodeRefactoring,
63 BugFix,
65 CodeReview,
67 Documentation,
69 TestGeneration,
71 Performance,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
77pub enum SessionState {
78 Active,
80 Processing,
82 WaitingForUser,
84 Paused,
86 Completed,
88 Error(String),
90}
91
92#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct UIState {
95 pub panel_visible: bool,
97 pub active_session: Option<String>,
99 pub panel_bounds: PanelBounds,
101 pub preferences: UIPreferences,
103}
104
105#[derive(Debug, Clone, Serialize, Deserialize)]
107pub struct PanelBounds {
108 pub width_percent: f32,
110 pub height_percent: f32,
112 pub position: PanelPosition,
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
118pub enum PanelPosition {
119 Right,
121 Left,
123 Bottom,
125 Floating { x: i32, y: i32 },
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct UIPreferences {
132 pub theme: String,
134 pub font_size: u32,
136 pub show_token_usage: bool,
138 pub auto_apply_simple: bool,
140 pub show_context_preview: bool,
142 pub max_visible_messages: usize,
144}
145
146#[derive(Debug, Clone, Serialize, Deserialize)]
148pub struct AssistantRequest {
149 pub session_id: Option<String>,
151 pub session_type: SessionType,
153 pub message: String,
155 pub include_context: bool,
157 pub context_hints: Vec<String>,
159}
160
161#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct AssistantResponse {
164 pub session_id: String,
166 pub message: String,
168 pub suggested_actions: Vec<SuggestedAction>,
170 pub token_usage: Option<TokenUsage>,
172 pub metadata: HashMap<String, serde_json::Value>,
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize)]
178pub struct SuggestedAction {
179 pub action_type: ActionType,
181 pub description: String,
183 pub data: serde_json::Value,
185 pub confidence: f64,
187}
188
189#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
191pub enum ActionType {
192 ApplyChanges,
194 OpenFile,
196 CreateFile,
198 RunCommand,
200 ShowDocumentation,
202 NavigateToDefinition,
204 AddImport,
206 GenerateTests,
208}
209
210#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct TokenUsage {
213 pub input_tokens: u32,
215 pub output_tokens: u32,
217 pub total_tokens: u32,
219 pub estimated_cost: Option<f64>,
221}
222
223impl Default for UIState {
224 fn default() -> Self {
225 Self {
226 panel_visible: false,
227 active_session: None,
228 panel_bounds: PanelBounds {
229 width_percent: 40.0,
230 height_percent: 60.0,
231 position: PanelPosition::Right,
232 },
233 preferences: UIPreferences {
234 theme: "dark".to_string(),
235 font_size: 14,
236 show_token_usage: true,
237 auto_apply_simple: false,
238 show_context_preview: true,
239 max_visible_messages: 50,
240 },
241 }
242 }
243}
244
245impl AIAssistant {
246 pub fn new(
248 coderlib: Arc<CoderLib>,
249 edit_host: Arc<EditHost>,
250 context_config: ContextConfig,
251 ) -> Self {
252 let context_gatherer = ContextGatherer::new(context_config);
253
254 Self {
255 coderlib,
256 edit_host,
257 context_gatherer,
258 sessions: Arc::new(RwLock::new(HashMap::new())),
259 ui_state: Arc::new(RwLock::new(UIState::default())),
260 command_sender: None,
261 }
262 }
263
264 pub fn set_command_sender(&mut self, sender: mpsc::UnboundedSender<HostCommand>) {
266 self.command_sender = Some(sender);
267 }
268
269 pub async fn show_panel(&self) -> Result<(), CoderLibError> {
271 let mut ui_state = self.ui_state.write().await;
272 ui_state.panel_visible = true;
273
274 self.send_command(HostCommand::ShowDialog {
275 title: "AI Assistant".to_string(),
276 message: "AI Assistant is ready to help!".to_string(),
277 buttons: vec!["Ask Question".to_string(), "Explain Code".to_string(), "Refactor".to_string(), "Close".to_string()],
278 }).await?;
279
280 info!("AI Assistant panel shown");
281 Ok(())
282 }
283
284 pub async fn hide_panel(&self) -> Result<(), CoderLibError> {
286 let mut ui_state = self.ui_state.write().await;
287 ui_state.panel_visible = false;
288
289 info!("AI Assistant panel hidden");
290 Ok(())
291 }
292
293 pub async fn process_request(&self, request: AssistantRequest) -> Result<AssistantResponse, CoderLibError> {
295 info!("Processing AI assistant request: {:?}", request.session_type);
296
297 let session_id = if let Some(ref id) = request.session_id {
299 id.clone()
300 } else {
301 self.create_new_session(request.session_type.clone()).await?
302 };
303
304 self.update_session_state(&session_id, SessionState::Processing).await?;
306
307 let context = if request.include_context {
309 let state = self.edit_host.get_state();
310 Some(self.context_gatherer.gather_context(&state, self.edit_host.as_ref()).await?)
311 } else {
312 None
313 };
314
315 let code_request = self.build_code_request(&request, &context).await?;
317
318 let mut code_response_receiver = self.coderlib.process_request(code_request).await?;
320
321 let code_response = if let Ok(response) = code_response_receiver.recv().await {
323 response
324 } else {
325 return Err(CoderLibError::Integration(
326 crate::core::error::IntegrationError::OperationFailed("No response received".to_string())
327 ));
328 };
329
330 let assistant_response = self.build_assistant_response(&session_id, &code_response, &context).await?;
332
333 self.update_session_messages(&session_id, &request.message, &assistant_response.message).await?;
335
336 self.update_session_state(&session_id, SessionState::WaitingForUser).await?;
338
339 info!("AI assistant request processed successfully");
340 Ok(assistant_response)
341 }
342
343 async fn create_new_session(&self, session_type: SessionType) -> Result<String, CoderLibError> {
345 let session_id = uuid::Uuid::new_v4().to_string();
346 let now = std::time::SystemTime::now();
347
348 let session = AssistantSession {
349 id: session_id.clone(),
350 session_type,
351 messages: Vec::new(),
352 context: None,
353 state: SessionState::Active,
354 created_at: now,
355 last_activity: now,
356 };
357
358 let mut sessions = self.sessions.write().await;
359 sessions.insert(session_id.clone(), session);
360
361 let mut ui_state = self.ui_state.write().await;
363 ui_state.active_session = Some(session_id.clone());
364
365 debug!("Created new AI assistant session: {}", session_id);
366 Ok(session_id)
367 }
368
369 async fn update_session_state(&self, session_id: &str, new_state: SessionState) -> Result<(), CoderLibError> {
371 let mut sessions = self.sessions.write().await;
372 if let Some(session) = sessions.get_mut(session_id) {
373 session.state = new_state;
374 session.last_activity = std::time::SystemTime::now();
375 }
376 Ok(())
377 }
378
379 async fn update_session_messages(&self, session_id: &str, user_message: &str, assistant_message: &str) -> Result<(), CoderLibError> {
381 let mut sessions = self.sessions.write().await;
382 if let Some(session) = sessions.get_mut(session_id) {
383 session.messages.push(Message {
384 id: uuid::Uuid::new_v4().to_string(),
385 session_id: session_id.to_string(),
386 role: MessageRole::User,
387 content: MessageContent::Text(user_message.to_string()),
388 timestamp: chrono::Utc::now(),
389 metadata: serde_json::Value::Null,
390 });
391 session.messages.push(Message {
392 id: uuid::Uuid::new_v4().to_string(),
393 session_id: session_id.to_string(),
394 role: MessageRole::Assistant,
395 content: MessageContent::Text(assistant_message.to_string()),
396 timestamp: chrono::Utc::now(),
397 metadata: serde_json::Value::Null,
398 });
399 session.last_activity = std::time::SystemTime::now();
400 }
401 Ok(())
402 }
403
404 async fn build_code_request(&self, request: &AssistantRequest, context: &Option<GatheredContext>) -> Result<CodeRequest, CoderLibError> {
406 let mut prompt = String::new();
407
408 match request.session_type {
410 SessionType::CodeExplanation => {
411 prompt.push_str("Please explain the following code:\n\n");
412 }
413 SessionType::CodeRefactoring => {
414 prompt.push_str("Please refactor the following code to improve its quality, readability, and performance:\n\n");
415 }
416 SessionType::BugFix => {
417 prompt.push_str("Please help me identify and fix bugs in the following code:\n\n");
418 }
419 SessionType::CodeReview => {
420 prompt.push_str("Please review the following code and provide feedback:\n\n");
421 }
422 SessionType::Documentation => {
423 prompt.push_str("Please generate documentation for the following code:\n\n");
424 }
425 SessionType::TestGeneration => {
426 prompt.push_str("Please generate comprehensive tests for the following code:\n\n");
427 }
428 SessionType::Performance => {
429 prompt.push_str("Please analyze and optimize the performance of the following code:\n\n");
430 }
431 SessionType::CodeAssistance => {
432 prompt.push_str("Please help me with the following code question:\n\n");
433 }
434 }
435
436 prompt.push_str(&request.message);
438 prompt.push_str("\n\n");
439
440 if let Some(ctx) = context {
442 prompt.push_str("## Context\n");
443 prompt.push_str(&self.context_gatherer.format_context_for_ai(ctx));
444 }
445
446 Ok(CodeRequest {
447 content: prompt,
448 attachments: Vec::new(),
449 model: None,
450 context: crate::core::RequestContext {
451 current_file: None,
452 cursor_position: None,
453 selection: None,
454 project_root: None,
455 open_files: Vec::new(),
456 },
457 session_id: uuid::Uuid::new_v4().to_string(),
458 })
459 }
460
461 async fn build_assistant_response(&self, session_id: &str, code_response: &CodeResponse, _context: &Option<GatheredContext>) -> Result<AssistantResponse, CoderLibError> {
463 let suggested_actions = self.extract_suggested_actions(&code_response.content).await;
464
465 Ok(AssistantResponse {
466 session_id: session_id.to_string(),
467 message: code_response.content.clone(),
468 suggested_actions,
469 token_usage: None, metadata: HashMap::new(),
471 })
472 }
473
474 async fn extract_suggested_actions(&self, _response_content: &str) -> Vec<SuggestedAction> {
476 vec![]
479 }
480
481 async fn send_command(&self, command: HostCommand) -> Result<(), CoderLibError> {
483 if let Some(sender) = &self.command_sender {
484 sender.send(command)
485 .map_err(|e| CoderLibError::Integration(
486 crate::core::error::IntegrationError::OperationFailed(
487 format!("Failed to send command: {}", e)
488 )
489 ))?;
490 }
491 Ok(())
492 }
493
494 pub async fn get_ui_state(&self) -> UIState {
496 self.ui_state.read().await.clone()
497 }
498
499 pub async fn update_preferences(&self, preferences: UIPreferences) -> Result<(), CoderLibError> {
501 let mut ui_state = self.ui_state.write().await;
502 ui_state.preferences = preferences;
503 Ok(())
504 }
505
506 pub async fn get_session(&self, session_id: &str) -> Option<AssistantSession> {
508 let sessions = self.sessions.read().await;
509 sessions.get(session_id).cloned()
510 }
511
512 pub async fn list_sessions(&self) -> Vec<AssistantSession> {
514 let sessions = self.sessions.read().await;
515 sessions.values().cloned().collect()
516 }
517
518 pub async fn close_session(&self, session_id: &str) -> Result<(), CoderLibError> {
520 let mut sessions = self.sessions.write().await;
521 sessions.remove(session_id);
522
523 let mut ui_state = self.ui_state.write().await;
525 if ui_state.active_session.as_ref() == Some(&session_id.to_string()) {
526 ui_state.active_session = None;
527 }
528
529 Ok(())
530 }
531}