//! LiveSpeech client implementation

use crate::audio::AudioEncoder;
use crate::error::{LiveSpeechError, Result};
use crate::types::*;

use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::sync::{broadcast, mpsc, Mutex, RwLock};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use tracing::{debug, error, info, warn};
use url::Url;

/// Connection state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    Disconnected,
    Connecting,
    Connected,
    Reconnecting,
}

/// LiveSpeech client for real-time speech-to-speech AI conversations
pub struct LiveSpeechClient {
    config: Config,
    state: Arc<RwLock<ConnectionState>>,
    connection_id: Arc<RwLock<Option<String>>>,
    session_id: Arc<RwLock<Option<String>>>,
    is_streaming: Arc<RwLock<bool>>,
    audio_encoder: AudioEncoder,
    
    // Message sender to the WebSocket task
    ws_sender: Arc<Mutex<Option<mpsc::Sender<ClientMessage>>>>,
    
    // Event handlers
    response_handler: Arc<RwLock<Option<ResponseHandler>>>,
    audio_handler: Arc<RwLock<Option<AudioHandler>>>,
    error_handler: Arc<RwLock<Option<ErrorHandler>>>,
    
    // Broadcast channel for external subscribers (allows multiple receivers)
    event_sender: broadcast::Sender<LiveSpeechEvent>,
}

impl LiveSpeechClient {
    /// Create a new LiveSpeech client
    pub fn new(config: Config) -> Self {
        let (event_sender, _) = broadcast::channel(100);
        
        Self {
            config,
            state: Arc::new(RwLock::new(ConnectionState::Disconnected)),
            connection_id: Arc::new(RwLock::new(None)),
            session_id: Arc::new(RwLock::new(None)),
            is_streaming: Arc::new(RwLock::new(false)),
            audio_encoder: AudioEncoder::new(),
            ws_sender: Arc::new(Mutex::new(None)),
            response_handler: Arc::new(RwLock::new(None)),
            audio_handler: Arc::new(RwLock::new(None)),
            error_handler: Arc::new(RwLock::new(None)),
            event_sender,
        }
    }

    /// Get current connection state
    pub async fn connection_state(&self) -> ConnectionState {
        *self.state.read().await
    }

    /// Get connection ID
    pub async fn connection_id(&self) -> Option<String> {
        self.connection_id.read().await.clone()
    }

    /// Get current session ID
    pub async fn session_id(&self) -> Option<String> {
        self.session_id.read().await.clone()
    }

    /// Check if connected
    pub async fn is_connected(&self) -> bool {
        *self.state.read().await == ConnectionState::Connected
    }

    /// Check if session is active
    pub async fn has_active_session(&self) -> bool {
        self.session_id.read().await.is_some()
    }

    /// Check if audio streaming is active
    pub async fn is_audio_streaming(&self) -> bool {
        *self.is_streaming.read().await
    }

    /// Set response handler (AI's text response)
    pub async fn on_response<F>(&self, handler: F)
    where
        F: Fn(&str, bool) + Send + Sync + 'static,
    {
        *self.response_handler.write().await = Some(Box::new(handler));
    }

    /// Set audio handler (AI's audio response)
    pub async fn on_audio<F>(&self, handler: F)
    where
        F: Fn(&[u8]) + Send + Sync + 'static,
    {
        *self.audio_handler.write().await = Some(Box::new(handler));
    }

    /// Set error handler
    pub async fn on_error<F>(&self, handler: F)
    where
        F: Fn(&ErrorEvent) + Send + Sync + 'static,
    {
        *self.error_handler.write().await = Some(Box::new(handler));
    }

    /// Subscribe to all events (recommended for full control)
    pub fn subscribe(&self) -> broadcast::Receiver<LiveSpeechEvent> {
        self.event_sender.subscribe()
    }

    /// Connect to the server
    pub async fn connect(&self) -> Result<()> {
        let current_state = *self.state.read().await;
        if current_state == ConnectionState::Connected || current_state == ConnectionState::Connecting {
            warn!("Already connected or connecting");
            return Ok(());
        }

        *self.state.write().await = ConnectionState::Connecting;

        // Parse and modify URL to include API key and optional user ID
        let mut url = Url::parse(&self.config.endpoint)?;
        url.query_pairs_mut()
            .append_pair("apiKey", &self.config.api_key);
        
        // Add userId if provided (for conversation memory persistence)
        if let Some(ref user_id) = self.config.user_id {
            url.query_pairs_mut().append_pair("userId", user_id);
        }

        info!("Connecting to {}", url.host_str().unwrap_or("unknown"));

        // Connect with timeout
        let connect_future = connect_async(url.as_str());
        let (ws_stream, _response) = tokio::time::timeout(
            self.config.connection_timeout,
            connect_future,
        )
        .await
        .map_err(|_| LiveSpeechError::ConnectionTimeout)?
        .map_err(LiveSpeechError::WebSocket)?;

        info!("WebSocket connected");

        *self.state.write().await = ConnectionState::Connected;
        let conn_id = generate_connection_id();
        *self.connection_id.write().await = Some(conn_id.clone());

        // Emit connected event
        let timestamp = chrono::Utc::now().to_rfc3339();
        let _ = self.event_sender.send(LiveSpeechEvent::Connected(ConnectedEvent {
            connection_id: conn_id,
            timestamp,
        }));

        let (write, read) = ws_stream.split();
        let write = Arc::new(Mutex::new(write));

        // Create channel for sending messages
        let (msg_sender, mut msg_receiver) = mpsc::channel::<ClientMessage>(100);
        *self.ws_sender.lock().await = Some(msg_sender);

        // Clone references for the tasks
        let state = self.state.clone();
        let session_id = self.session_id.clone();
        let is_streaming = self.is_streaming.clone();
        let response_handler = self.response_handler.clone();
        let audio_handler = self.audio_handler.clone();
        let error_handler = self.error_handler.clone();
        let event_sender = self.event_sender.clone();
        let audio_encoder = self.audio_encoder.clone();

        // Spawn write task
        let write_clone = write.clone();
        tokio::spawn(async move {
            while let Some(msg) = msg_receiver.recv().await {
                if let Ok(json) = msg.to_json() {
                    debug!("Sending message: {:?}", msg);
                    let mut writer = write_clone.lock().await;
                    if let Err(e) = writer.send(Message::Text(json)).await {
                        error!("Failed to send message: {}", e);
                        break;
                    }
                }
            }
        });

        // Spawn read task
        tokio::spawn(async move {
            let mut read = read;
            while let Some(result) = read.next().await {
                match result {
                    Ok(Message::Text(text)) => {
                        debug!("Received message: {}", text);
                        match ServerMessage::from_json(&text) {
                            Ok(msg) => {
                                Self::handle_message(
                                    msg,
                                    &state,
                                    &session_id,
                                    &is_streaming,
                                    &response_handler,
                                    &audio_handler,
                                    &error_handler,
                                    &event_sender,
                                    &audio_encoder,
                                )
                                .await;
                            }
                            Err(e) => {
                                warn!("Failed to parse message: {}", e);
                            }
                        }
                    }
                    Ok(Message::Close(_)) => {
                        info!("WebSocket closed by server");
                        *state.write().await = ConnectionState::Disconnected;
                        break;
                    }
                    Err(e) => {
                        error!("WebSocket error: {}", e);
                        *state.write().await = ConnectionState::Disconnected;
                        break;
                    }
                    _ => {}
                }
            }
        });

        // Start ping task
        let ws_sender = self.ws_sender.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(30));
            loop {
                interval.tick().await;
                if let Some(sender) = ws_sender.lock().await.as_ref() {
                    if sender.send(ClientMessage::Ping).await.is_err() {
                        break;
                    }
                } else {
                    break;
                }
            }
        });

        Ok(())
    }

    /// Disconnect from the server
    pub async fn disconnect(&self) {
        info!("Disconnecting");
        *self.ws_sender.lock().await = None;
        *self.state.write().await = ConnectionState::Disconnected;
        *self.connection_id.write().await = None;
        *self.session_id.write().await = None;
        *self.is_streaming.write().await = false;
    }

    /// Start a new session
    pub async fn start_session(&self, config: Option<SessionConfig>) -> Result<String> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if self.session_id.read().await.is_some() {
            return Err(LiveSpeechError::SessionAlreadyActive);
        }

        let (
            pre_prompt,
            language,
            pipeline_mode,
            ai_speaks_first,
            allow_harm_category,
            tools,
            session_duration,
            output_language,
        ) = config
            .map(|c| (
                c.pre_prompt,
                c.language,
                Some(c.pipeline_mode.as_str().to_string()),
                if c.ai_speaks_first { Some(true) } else { None },
                // allow_harm_category defaults to false; only send Some(true) if explicitly enabled
                if c.allow_harm_category { Some(true) } else { None },
                c.tools,
                c.session_duration,
                c.output_language,
            ))
            .unwrap_or((None, None, None, None, None, None, None, None));
        let msg = ClientMessage::start_session(StartSessionParams {
            pre_prompt,
            language,
            output_language,
            pipeline_mode,
            ai_speaks_first,
            allow_harm_category,
            tools,
            session_duration,
        });
        
        // Subscribe to events before sending message
        let mut events = self.event_sender.subscribe();
        
        self.send_message(msg).await?;

        // Wait for SessionStarted or Error event with timeout
        let timeout_duration = self.config.connection_timeout;
        let result = tokio::time::timeout(timeout_duration, async {
            while let Ok(event) = events.recv().await {
                match event {
                    LiveSpeechEvent::SessionStarted(e) => {
                        return Ok(e.session_id);
                    }
                    LiveSpeechEvent::Error(e) if matches!(e.code, ErrorCode::SessionError) => {
                        return Err(LiveSpeechError::SessionError(e.message));
                    }
                    _ => continue,
                }
            }
            Err(LiveSpeechError::ChannelReceive)
        })
        .await
        .map_err(|_| LiveSpeechError::ConnectionTimeout)?;
        
        result
    }

    /// End the current session
    pub async fn end_session(&self) -> Result<()> {
        if self.session_id.read().await.is_none() {
            warn!("No active session to end");
            return Ok(());
        }

        // Stop streaming if active
        if *self.is_streaming.read().await {
            self.audio_end().await?;
        }

        // Subscribe to events before sending message
        let mut events = self.event_sender.subscribe();

        self.send_message(ClientMessage::end_session()).await?;
        
        // Wait for SessionEnded event with timeout
        let timeout_duration = self.config.connection_timeout;
        tokio::time::timeout(timeout_duration, async {
            while let Ok(event) = events.recv().await {
                if matches!(event, LiveSpeechEvent::SessionEnded(_)) {
                    return;
                }
            }
        })
        .await
        .map_err(|_| LiveSpeechError::ConnectionTimeout)?;
        
        Ok(())
    }

    /// Start audio streaming session
    pub async fn audio_start(&self) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if self.session_id.read().await.is_none() {
            return Err(LiveSpeechError::NoActiveSession);
        }

        if *self.is_streaming.read().await {
            return Err(LiveSpeechError::AlreadyStreaming);
        }

        // Subscribe to events before sending message to ensure we don't miss Ready
        let mut events = self.event_sender.subscribe();
        
        self.send_message(ClientMessage::audio_start()).await?;

        // Wait for Ready or Error event with timeout
        // This ensures Gemini Live session is fully established before returning
        let timeout_duration = self.config.connection_timeout;
        let result = tokio::time::timeout(timeout_duration, async {
            while let Ok(event) = events.recv().await {
                match event {
                    LiveSpeechEvent::Ready(_) => {
                        return Ok(());
                    }
                    LiveSpeechEvent::Error(e) if matches!(e.code, ErrorCode::StreamingError) => {
                        return Err(LiveSpeechError::SessionError(e.message));
                    }
                    _ => continue,
                }
            }
            Err(LiveSpeechError::ChannelReceive)
        })
        .await
        .map_err(|_| LiveSpeechError::ConnectionTimeout)?;
        
        result?;
        
        *self.is_streaming.write().await = true;
        info!("LiveSpeech audio stream started");
        Ok(())
    }

    /// Send audio chunk (PCM16 bytes)
    pub async fn send_audio_chunk(&self, data: &[u8]) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if !*self.is_streaming.read().await {
            return Err(LiveSpeechError::NotStreaming);
        }

        let base64_data = self.audio_encoder.encode(data);
        self.send_message(ClientMessage::audio_chunk(base64_data)).await
    }

    /// End audio streaming session
    pub async fn audio_end(&self) -> Result<()> {
        if !*self.is_streaming.read().await {
            warn!("Not streaming");
            return Ok(());
        }

        self.send_message(ClientMessage::audio_end()).await?;
        *self.is_streaming.write().await = false;
        Ok(())
    }

    /// Send a system message to the AI (AI responds immediately)
    ///
    /// This allows your application to inject context or trigger AI responses
    /// based on external events (e.g., game events, app state changes, timers).
    ///
    /// # Arguments
    /// * `text` - The system message text (max 500 characters)
    ///
    /// # Errors
    /// Returns an error if not connected or not streaming.
    pub async fn send_system_message(&self, text: &str) -> Result<()> {
        self.send_system_message_with_options(text, true).await
    }

    /// Send a system message to the AI with explicit trigger_response option
    ///
    /// # Arguments
    /// * `text` - The system message text (max 500 characters)
    /// * `trigger_response` - If true, AI responds immediately; if false, context only
    ///
    /// # Errors
    /// Returns an error if not connected, not streaming, or text exceeds 500 chars.
    pub async fn send_system_message_with_options(&self, text: &str, trigger_response: bool) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if !*self.is_streaming.read().await {
            return Err(LiveSpeechError::NotStreaming);
        }

        if text.len() > 500 {
            return Err(LiveSpeechError::InvalidParameter("System message too long (max 500 characters)".to_string()));
        }

        info!("Sending system message: {} (trigger_response: {})", text, trigger_response);
        self.send_message(ClientMessage::system_message_with_options(text, trigger_response)).await
    }

    /// Send a tool response (function execution result) back to AI
    ///
    /// After receiving a `ToolCall` event, execute the function and send
    /// the result back using this method. The AI will then continue the
    /// conversation based on the result.
    ///
    /// # Arguments
    /// * `id` - The tool call ID from the `ToolCallEvent`
    /// * `response` - The function execution result (will be JSON serialized)
    ///
    /// # Example
    /// ```ignore
    /// client.on_tool_call(|event| {
    ///     let result = execute_function(&event.name, &event.args);
    ///     client.send_tool_response(&event.id, result).await?;
    /// });
    /// ```
    ///
    /// # Errors
    /// Returns an error if not connected or not streaming.
    pub async fn send_tool_response(&self, id: &str, response: serde_json::Value) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if !*self.is_streaming.read().await {
            return Err(LiveSpeechError::NotStreaming);
        }

        info!("Sending tool response for id: {}", id);
        self.send_message(ClientMessage::tool_response(id, response)).await
    }

    /// Explicitly interrupt the current AI response
    ///
    /// Use this method for:
    /// - UI "Stop" button functionality
    /// - Programmatic control to stop AI mid-response
    ///
    /// Note: In most cases, simply speaking will trigger automatic
    /// interruption via Gemini's voice activity detection (VAD).
    /// This method is for explicit programmatic control.
    ///
    /// # Example
    /// ```ignore
    /// // User clicks "Stop" button
    /// client.interrupt().await?;
    /// ```
    ///
    /// # Errors
    /// Returns an error if not connected or not streaming.
    pub async fn interrupt(&self) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if !*self.is_streaming.read().await {
            return Err(LiveSpeechError::NotStreaming);
        }

        info!("Sending explicit interrupt");
        self.send_message(ClientMessage::interrupt()).await
    }

    /// Update user ID for the current connection (guest-to-user migration)
    ///
    /// When a guest user logs in during a session, call this method to migrate
    /// their conversation history to a persistent user partition. This enables:
    /// - Entity extraction and long-term memory for the user
    /// - Conversation continuity across sessions
    /// - Personalization based on past interactions
    ///
    /// # Arguments
    /// * `user_id` - The authenticated user's unique identifier
    ///
    /// # Example
    /// ```ignore
    /// // After successful login
    /// let user = authenticate(credentials).await?;
    /// client.update_user_id(&user.id).await?;
    /// ```
    ///
    /// # Errors
    /// Returns an error if not connected or user_id is empty.
    pub async fn update_user_id(&self, user_id: &str) -> Result<()> {
        if !self.is_connected().await {
            return Err(LiveSpeechError::NotConnected);
        }

        if user_id.trim().is_empty() {
            return Err(LiveSpeechError::InvalidParameter("userId cannot be empty".to_string()));
        }

        info!("Updating user ID: {}", user_id);
        self.send_message(ClientMessage::update_user_id(user_id)).await
    }

    /// Send a client message
    async fn send_message(&self, msg: ClientMessage) -> Result<()> {
        let sender = self.ws_sender.lock().await;
        if let Some(sender) = sender.as_ref() {
            sender
                .send(msg)
                .await
                .map_err(|_| LiveSpeechError::ChannelSend)?;
            Ok(())
        } else {
            Err(LiveSpeechError::NotConnected)
        }
    }

    /// Handle incoming server message
    async fn handle_message(
        msg: ServerMessage,
        state: &Arc<RwLock<ConnectionState>>,
        session_id: &Arc<RwLock<Option<String>>>,
        is_streaming: &Arc<RwLock<bool>>,
        response_handler: &Arc<RwLock<Option<ResponseHandler>>>,
        audio_handler: &Arc<RwLock<Option<AudioHandler>>>,
        error_handler: &Arc<RwLock<Option<ErrorHandler>>>,
        event_sender: &broadcast::Sender<LiveSpeechEvent>,
        audio_encoder: &AudioEncoder,
    ) {
        let timestamp = chrono::Utc::now().to_rfc3339();
        
        match msg {
            ServerMessage::SessionStarted { session_id: sess_id, .. } => {
                *session_id.write().await = Some(sess_id.clone());
                let _ = event_sender.send(LiveSpeechEvent::SessionStarted(SessionStartedEvent {
                    session_id: sess_id,
                    timestamp,
                }));
            }
            
            ServerMessage::SessionEnded { session_id: sess_id, .. } => {
                *session_id.write().await = None;
                *is_streaming.write().await = false;
                let _ = event_sender.send(LiveSpeechEvent::SessionEnded(SessionEndedEvent {
                    session_id: sess_id,
                    timestamp,
                }));
            }
            
            ServerMessage::Ready { .. } => {
                info!("Session ready for audio input");
                let _ = event_sender.send(LiveSpeechEvent::Ready(ReadyEvent { timestamp }));
            }
            
            ServerMessage::UserTranscript { text, .. } => {
                info!("User transcript: {}", text);
                let _ = event_sender.send(LiveSpeechEvent::UserTranscript(UserTranscriptEvent {
                    text,
                    timestamp,
                }));
            }
            
            ServerMessage::Response { text, is_final, .. } => {
                if let Some(handler) = response_handler.read().await.as_ref() {
                    handler(&text, is_final);
                }
                let _ = event_sender.send(LiveSpeechEvent::Response(ResponseEvent {
                    text,
                    is_final,
                    timestamp,
                }));
            }
            
            ServerMessage::Audio { data, format, sample_rate, .. } => {
                if let Ok(audio_data) = audio_encoder.decode(&data) {
                    if let Some(handler) = audio_handler.read().await.as_ref() {
                        handler(&audio_data);
                    }
                    let _ = event_sender.send(LiveSpeechEvent::Audio(AudioEvent {
                        data: audio_data,
                        format,
                        sample_rate,
                        timestamp,
                    }));
                }
            }
            
            ServerMessage::TurnComplete { .. } => {
                info!("Turn complete");
                let _ = event_sender.send(LiveSpeechEvent::TurnComplete(TurnCompleteEvent { timestamp }));
            }

            ServerMessage::SessionWarning { remaining_seconds, .. } => {
                info!("Session warning: {}s remaining", remaining_seconds);
                let _ = event_sender.send(LiveSpeechEvent::SessionWarning(SessionWarningEvent {
                    remaining_seconds,
                    timestamp,
                }));
            }

            ServerMessage::SessionGoodbye { remaining_seconds, .. } => {
                info!("Session goodbye: {}s remaining", remaining_seconds);
                let _ = event_sender.send(LiveSpeechEvent::SessionGoodbye(SessionGoodbyeEvent {
                    remaining_seconds,
                    timestamp,
                }));
            }
            
            ServerMessage::ToolCall { id, name, args, .. } => {
                info!("Tool call received: {} (id: {})", name, id);
                let _ = event_sender.send(LiveSpeechEvent::ToolCall(ToolCallEvent {
                    id,
                    name,
                    args,
                    timestamp,
                }));
            }
            
            ServerMessage::UserIdUpdated { user_id, migrated_messages, .. } => {
                info!("User ID updated: {}, migrated {} messages", user_id, migrated_messages);
                let _ = event_sender.send(LiveSpeechEvent::UserIdUpdated(UserIdUpdatedEvent {
                    user_id,
                    migrated_messages,
                    timestamp,
                }));
            }
            
            ServerMessage::Interrupted { .. } => {
                info!("AI response interrupted (barge-in)");
                let _ = event_sender.send(LiveSpeechEvent::Interrupted(InterruptedEvent {
                    timestamp,
                }));
            }
            
            ServerMessage::Error { code, message, .. } => {
                let error_code = match code.as_str() {
                    "connection_failed" => ErrorCode::ConnectionFailed,
                    "authentication_failed" => ErrorCode::AuthenticationFailed,
                    "session_error" => ErrorCode::SessionError,
                    "audio_error" => ErrorCode::AudioError,
                    "streaming_error" => ErrorCode::StreamingError,
                    "stt_error" => ErrorCode::SttError,
                    "llm_error" => ErrorCode::LlmError,
                    "tts_error" => ErrorCode::TtsError,
                    "rate_limit" => ErrorCode::RateLimit,
                    "user_id_update_error" => ErrorCode::UserIdUpdateError,
                    _ => ErrorCode::InternalError,
                };
                
                let error_event = ErrorEvent {
                    code: error_code,
                    message: message.clone(),
                    details: None,
                    timestamp: timestamp.clone(),
                };
                
                if let Some(handler) = error_handler.read().await.as_ref() {
                    handler(&error_event);
                }
                let _ = event_sender.send(LiveSpeechEvent::Error(error_event));
            }
            
            ServerMessage::Pong { .. } => {
                debug!("Pong received");
            }
        }
        
        // Suppress unused warning
        let _ = state;
    }
}

/// Generate a client-side connection ID
fn generate_connection_id() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis();
    format!("client_{}_{:x}", timestamp, rand::random::<u32>())
}
