coderlib 0.1.0

//! Streaming response handler for Edit integration
//!
//! This module provides real-time streaming response display and user
//! interaction for AI assistant responses within Microsoft Edit.

use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock, broadcast};
use tokio_stream::{Stream, StreamExt};
use tracing::{debug, info, error};

use crate::core::CoderLibError;
use crate::integration::{HostCommand, MessageLevel};

/// Simple streaming chunk for integration
#[derive(Debug, Clone)]
pub struct StreamingChunk {
    /// Content of the chunk
    pub content: Option<String>,
    /// Number of tokens in this chunk
    pub tokens: Option<u32>,
    /// Whether this is the final chunk
    pub is_final: bool,
}

/// Streaming response handler for Edit integration
pub struct StreamingHandler {
    /// Command sender to Edit
    command_sender: Option<mpsc::UnboundedSender<HostCommand>>,
    /// Active streaming sessions
    active_sessions: Arc<RwLock<std::collections::HashMap<String, StreamingSession>>>,
    /// Event broadcaster for UI updates
    event_broadcaster: broadcast::Sender<StreamingEvent>,
    /// Configuration
    config: StreamingConfig,
}

/// Streaming session information
#[derive(Debug, Clone)]
pub struct StreamingSession {
    /// Session ID
    pub id: String,
    /// Current accumulated response
    pub accumulated_response: String,
    /// Session state
    pub state: StreamingState,
    /// Start time
    pub start_time: std::time::Instant,
    /// Last update time
    pub last_update: std::time::Instant,
    /// Token count
    pub token_count: usize,
    /// Estimated completion percentage
    pub completion_percentage: f32,
}

/// Streaming session state
#[derive(Debug, Clone, PartialEq)]
pub enum StreamingState {
    /// Session is starting
    Starting,
    /// Actively streaming
    Streaming,
    /// Stream completed successfully
    Completed,
    /// Stream was cancelled
    Cancelled,
    /// Stream encountered an error
    Error(String),
}

/// Streaming configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamingConfig {
    /// Update interval in milliseconds
    pub update_interval_ms: u64,
    /// Maximum buffer size before forcing update
    pub max_buffer_size: usize,
    /// Whether to show typing indicator
    pub show_typing_indicator: bool,
    /// Whether to show token count
    pub show_token_count: bool,
    /// Whether to show completion percentage
    pub show_completion_percentage: bool,
    /// Debounce time for rapid updates
    pub debounce_ms: u64,
}

impl Default for StreamingConfig {
    fn default() -> Self {
        Self {
            update_interval_ms: 100,
            max_buffer_size: 1000,
            show_typing_indicator: true,
            show_token_count: true,
            show_completion_percentage: true,
            debounce_ms: 50,
        }
    }
}

/// Streaming events for UI updates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StreamingEvent {
    /// Streaming session started
    SessionStarted {
        session_id: String,
    },
    /// New content chunk received
    ContentUpdate {
        session_id: String,
        content: String,
        is_complete: bool,
    },
    /// Token count updated
    TokenUpdate {
        session_id: String,
        token_count: usize,
    },
    /// Completion percentage updated
    ProgressUpdate {
        session_id: String,
        percentage: f32,
    },
    /// Streaming session completed
    SessionCompleted {
        session_id: String,
        final_content: String,
        total_tokens: usize,
        duration_ms: u64,
    },
    /// Streaming session cancelled
    SessionCancelled {
        session_id: String,
    },
    /// Streaming session error
    SessionError {
        session_id: String,
        error: String,
    },
}

/// Streaming response display options
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DisplayOptions {
    /// Whether to show in a separate panel
    pub use_separate_panel: bool,
    /// Whether to replace content inline
    pub replace_inline: bool,
    /// Whether to show diff highlighting
    pub show_diff: bool,
    /// Whether to enable user interaction during streaming
    pub allow_interaction: bool,
    /// Panel position for separate panel
    pub panel_position: PanelPosition,
}

/// Panel position for streaming display
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum PanelPosition {
    /// Right side of editor
    Right,
    /// Bottom of editor
    Bottom,
    /// Floating overlay
    Overlay,
}

impl Default for DisplayOptions {
    fn default() -> Self {
        Self {
            use_separate_panel: true,
            replace_inline: false,
            show_diff: true,
            allow_interaction: true,
            panel_position: PanelPosition::Right,
        }
    }
}

impl StreamingHandler {
    /// Create a new streaming handler
    pub fn new(config: StreamingConfig) -> Self {
        let (event_broadcaster, _) = broadcast::channel(1000);
        
        Self {
            command_sender: None,
            active_sessions: Arc::new(RwLock::new(std::collections::HashMap::new())),
            event_broadcaster,
            config,
        }
    }

    /// Set the command sender for communicating with Edit
    pub fn set_command_sender(&mut self, sender: mpsc::UnboundedSender<HostCommand>) {
        self.command_sender = Some(sender);
    }

    /// Subscribe to streaming events
    pub fn subscribe_to_events(&self) -> broadcast::Receiver<StreamingEvent> {
        self.event_broadcaster.subscribe()
    }

    /// Start handling a streaming response
    pub async fn handle_streaming_response<S>(
        &self,
        session_id: String,
        stream: S,
        display_options: DisplayOptions,
    ) -> Result<String, CoderLibError>
    where
        S: Stream<Item = Result<StreamingChunk, CoderLibError>> + Send + Unpin + 'static,
    {
        info!("Starting streaming response handler for session: {}", session_id);

        // Create streaming session
        let session = StreamingSession {
            id: session_id.clone(),
            accumulated_response: String::new(),
            state: StreamingState::Starting,
            start_time: std::time::Instant::now(),
            last_update: std::time::Instant::now(),
            token_count: 0,
            completion_percentage: 0.0,
        };

        // Add to active sessions
        {
            let mut sessions = self.active_sessions.write().await;
            sessions.insert(session_id.clone(), session);
        }

        // Broadcast session started event
        let _ = self.event_broadcaster.send(StreamingEvent::SessionStarted {
            session_id: session_id.clone(),
        });

        // Show initial UI
        self.show_streaming_ui(&session_id, &display_options).await?;

        // Process the stream
        let final_content = self.process_stream(session_id.clone(), stream).await?;

        // Update session state to completed
        {
            let mut sessions = self.active_sessions.write().await;
            if let Some(session) = sessions.get_mut(&session_id) {
                session.state = StreamingState::Completed;
                session.accumulated_response = final_content.clone();
            }
        }

        // Broadcast completion event
        let duration = {
            let sessions = self.active_sessions.read().await;
            sessions.get(&session_id)
                .map(|s| s.start_time.elapsed().as_millis() as u64)
                .unwrap_or(0)
        };

        let token_count = {
            let sessions = self.active_sessions.read().await;
            sessions.get(&session_id)
                .map(|s| s.token_count)
                .unwrap_or(0)
        };

        let _ = self.event_broadcaster.send(StreamingEvent::SessionCompleted {
            session_id: session_id.clone(),
            final_content: final_content.clone(),
            total_tokens: token_count,
            duration_ms: duration,
        });

        // Clean up session
        {
            let mut sessions = self.active_sessions.write().await;
            sessions.remove(&session_id);
        }

        info!("Streaming response completed for session: {}", session_id);
        Ok(final_content)
    }

    /// Process the streaming chunks
    async fn process_stream<S>(
        &self,
        session_id: String,
        mut stream: S,
    ) -> Result<String, CoderLibError>
    where
        S: Stream<Item = Result<StreamingChunk, CoderLibError>> + Send + Unpin,
    {
        let mut accumulated_content = String::new();
        let mut buffer = String::new();
        let mut last_update = std::time::Instant::now();

        // Update session state to streaming
        {
            let mut sessions = self.active_sessions.write().await;
            if let Some(session) = sessions.get_mut(&session_id) {
                session.state = StreamingState::Streaming;
            }
        }

        while let Some(chunk_result) = stream.next().await {
            match chunk_result {
                Ok(chunk) => {
                    // Add chunk content to buffer
                    if let Some(content) = chunk.content {
                        buffer.push_str(&content);
                        accumulated_content.push_str(&content);
                    }

                    // Update token count if available
                    if let Some(tokens) = chunk.tokens {
                        self.update_token_count(&session_id, tokens).await;
                    }

                    // Check if we should update the UI
                    let should_update = buffer.len() >= self.config.max_buffer_size
                        || last_update.elapsed().as_millis() >= self.config.update_interval_ms as u128;

                    if should_update && !buffer.is_empty() {
                        self.update_streaming_content(&session_id, &buffer, false).await?;
                        buffer.clear();
                        last_update = std::time::Instant::now();
                    }

                    // Update session
                    {
                        let mut sessions = self.active_sessions.write().await;
                        if let Some(session) = sessions.get_mut(&session_id) {
                            session.accumulated_response = accumulated_content.clone();
                            session.last_update = std::time::Instant::now();
                            session.token_count += chunk.tokens.unwrap_or(0) as usize;
                        }
                    }
                }
                Err(e) => {
                    error!("Streaming error for session {}: {}", session_id, e);
                    
                    // Update session state to error
                    {
                        let mut sessions = self.active_sessions.write().await;
                        if let Some(session) = sessions.get_mut(&session_id) {
                            session.state = StreamingState::Error(e.to_string());
                        }
                    }

                    // Broadcast error event
                    let _ = self.event_broadcaster.send(StreamingEvent::SessionError {
                        session_id: session_id.clone(),
                        error: e.to_string(),
                    });

                    return Err(e);
                }
            }
        }

        // Send final update if there's remaining content in buffer
        if !buffer.is_empty() {
            self.update_streaming_content(&session_id, &buffer, true).await?;
        }

        Ok(accumulated_content)
    }

    /// Show streaming UI
    async fn show_streaming_ui(&self, session_id: &str, options: &DisplayOptions) -> Result<(), CoderLibError> {
        if options.use_separate_panel {
            self.send_command(HostCommand::ShowDialog {
                title: format!("AI Response - {}", session_id),
                message: "AI is generating response...".to_string(),
                buttons: vec!["Cancel".to_string()],
            }).await?;
        }

        if self.config.show_typing_indicator {
            self.send_command(HostCommand::ShowMessage {
                message: "AI is typing...".to_string(),
                level: MessageLevel::Info,
            }).await?;
        }

        Ok(())
    }

    /// Update streaming content in the UI
    async fn update_streaming_content(&self, session_id: &str, content: &str, is_complete: bool) -> Result<(), CoderLibError> {
        debug!("Updating streaming content for session: {} (complete: {})", session_id, is_complete);

        // Broadcast content update event
        let _ = self.event_broadcaster.send(StreamingEvent::ContentUpdate {
            session_id: session_id.to_string(),
            content: content.to_string(),
            is_complete,
        });

        // Update UI through Edit commands
        if is_complete {
            self.send_command(HostCommand::ShowMessage {
                message: "AI response completed".to_string(),
                level: MessageLevel::Success,
            }).await?;
        }

        Ok(())
    }

    /// Update token count for a session
    async fn update_token_count(&self, session_id: &str, tokens: u32) {
        if self.config.show_token_count {
            let _ = self.event_broadcaster.send(StreamingEvent::TokenUpdate {
                session_id: session_id.to_string(),
                token_count: tokens as usize,
            });
        }
    }

    /// Cancel a streaming session
    pub async fn cancel_session(&self, session_id: &str) -> Result<(), CoderLibError> {
        info!("Cancelling streaming session: {}", session_id);

        // Update session state
        {
            let mut sessions = self.active_sessions.write().await;
            if let Some(session) = sessions.get_mut(session_id) {
                session.state = StreamingState::Cancelled;
            }
        }

        // Broadcast cancellation event
        let _ = self.event_broadcaster.send(StreamingEvent::SessionCancelled {
            session_id: session_id.to_string(),
        });

        // Clean up UI
        self.send_command(HostCommand::ShowMessage {
            message: "AI response cancelled".to_string(),
            level: MessageLevel::Warning,
        }).await?;

        Ok(())
    }

    /// Get active sessions
    pub async fn get_active_sessions(&self) -> Vec<StreamingSession> {
        let sessions = self.active_sessions.read().await;
        sessions.values().cloned().collect()
    }

    /// Send a command to Edit
    async fn send_command(&self, command: HostCommand) -> Result<(), CoderLibError> {
        if let Some(sender) = &self.command_sender {
            sender.send(command)
                .map_err(|e| CoderLibError::Integration(
                    crate::core::error::IntegrationError::OperationFailed(
                        format!("Failed to send command: {}", e)
                    )
                ))?;
        }
        Ok(())
    }
}