//! WebSocket message types for the LiveSpeech SDK

use serde::{Deserialize, Serialize};

/// Client message types (sent from client to server)
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "action", rename_all = "camelCase")]
pub enum ClientMessage {
    /// Start a new session
    #[serde(rename = "startSession")]
    StartSession {
        #[serde(rename = "prePrompt", skip_serializing_if = "Option::is_none")]
        pre_prompt: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        language: Option<String>,
    },
    /// End the current session
    #[serde(rename = "endSession")]
    EndSession,
    /// Start audio streaming session
    #[serde(rename = "audioStart")]
    AudioStart,
    /// Send audio chunk
    #[serde(rename = "audioChunk")]
    AudioChunk {
        /// Base64 encoded PCM16 audio data
        data: String,
    },
    /// End audio streaming session
    #[serde(rename = "audioEnd")]
    AudioEnd,
    /// Ping for keep-alive
    #[serde(rename = "ping")]
    Ping,
}

/// Server message types (received from server)
#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ServerMessage {
    /// Session started
    #[serde(rename = "sessionStarted")]
    SessionStarted {
        #[serde(rename = "sessionId")]
        session_id: String,
        timestamp: String,
    },
    /// Session ended
    #[serde(rename = "sessionEnded")]
    SessionEnded {
        #[serde(rename = "sessionId")]
        session_id: String,
        timestamp: String,
    },
    /// Streaming started - acknowledgment of audioStart
    #[serde(rename = "streamingStarted")]
    StreamingStarted { timestamp: String },
    /// Speech start - VAD detected speech begin
    #[serde(rename = "speechStart")]
    SpeechStart { timestamp: String },
    /// Speech end - VAD detected speech end
    #[serde(rename = "speechEnd")]
    SpeechEnd { timestamp: String },
    /// Transcript from STT
    #[serde(rename = "transcript")]
    Transcript {
        text: String,
        #[serde(rename = "isFinal")]
        is_final: bool,
        confidence: Option<f32>,
        timestamp: String,
    },
    /// Response from LLM
    #[serde(rename = "response")]
    Response {
        text: String,
        #[serde(rename = "isFinal")]
        is_final: bool,
        timestamp: String,
    },
    /// Audio from TTS
    #[serde(rename = "audio")]
    Audio {
        /// Base64 encoded PCM16 audio data
        data: String,
        format: String,
        #[serde(rename = "sampleRate")]
        sample_rate: u32,
        timestamp: String,
    },
    /// Error message
    #[serde(rename = "error")]
    Error {
        code: String,
        message: String,
        timestamp: String,
    },
    /// Pong response
    #[serde(rename = "pong")]
    Pong { timestamp: String },
}

impl ClientMessage {
    /// Create a start session message
    pub fn start_session(pre_prompt: Option<String>, language: Option<String>) -> Self {
        ClientMessage::StartSession { pre_prompt, language }
    }

    /// Create an end session message
    pub fn end_session() -> Self {
        ClientMessage::EndSession
    }

    /// Create an audio start message
    pub fn audio_start() -> Self {
        ClientMessage::AudioStart
    }

    /// Create an audio chunk message
    pub fn audio_chunk(data: impl Into<String>) -> Self {
        ClientMessage::AudioChunk { data: data.into() }
    }

    /// Create an audio end message
    pub fn audio_end() -> Self {
        ClientMessage::AudioEnd
    }

    /// Create a ping message
    pub fn ping() -> Self {
        ClientMessage::Ping
    }

    /// Serialize to JSON string
    pub fn to_json(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string(self)
    }
}

impl ServerMessage {
    /// Parse from JSON string
    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
        serde_json::from_str(json)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_serialize_start_session() {
        let msg = ClientMessage::start_session(Some("You are helpful".to_string()), None);
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("prePrompt"));
    }

    #[test]
    fn test_serialize_start_session_with_language() {
        let msg = ClientMessage::start_session(Some("You are helpful".to_string()), Some("ko-KR".to_string()));
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("prePrompt"));
        assert!(json.contains("ko-KR"));
    }

    #[test]
    fn test_serialize_audio_chunk() {
        let msg = ClientMessage::audio_chunk("base64data");
        let json = msg.to_json().unwrap();
        assert!(json.contains("audioChunk"));
        assert!(json.contains("base64data"));
    }

    #[test]
    fn test_serialize_audio_start() {
        let msg = ClientMessage::audio_start();
        let json = msg.to_json().unwrap();
        assert!(json.contains("audioStart"));
    }

    #[test]
    fn test_deserialize_session_started() {
        let json = r#"{"type":"sessionStarted","sessionId":"abc123","timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::SessionStarted { session_id, .. } => {
                assert_eq!(session_id, "abc123");
            }
            _ => panic!("Expected SessionStarted message"),
        }
    }

    #[test]
    fn test_deserialize_transcript() {
        let json = r#"{"type":"transcript","text":"hello world","isFinal":true,"timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::Transcript { text, is_final, .. } => {
                assert_eq!(text, "hello world");
                assert!(is_final);
            }
            _ => panic!("Expected Transcript message"),
        }
    }

    #[test]
    fn test_deserialize_streaming_started() {
        let json = r#"{"type":"streamingStarted","timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::StreamingStarted { .. } => {}
            _ => panic!("Expected StreamingStarted message"),
        }
    }

    #[test]
    fn test_deserialize_speech_start() {
        let json = r#"{"type":"speechStart","timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::SpeechStart { .. } => {}
            _ => panic!("Expected SpeechStart message"),
        }
    }
}
