//! WebSocket message types for the LiveSpeech SDK

use serde::{Deserialize, Serialize};
use super::config::{SessionDurationConfig, Tool};

/// Tool response payload for function calling
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolResponsePayload {
    /// Tool call ID from the toolCall event
    pub id: String,
    /// Function execution result
    pub response: serde_json::Value,
}

/// System message payload
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SystemMessagePayload {
    /// The message text (max 500 characters)
    pub text: String,
    /// If true, AI responds immediately (default: true)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub trigger_response: Option<bool>,
}

/// 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>,
        #[serde(rename = "outputLanguage", skip_serializing_if = "Option::is_none")]
        output_language: Option<String>,
        #[serde(rename = "pipelineMode", skip_serializing_if = "Option::is_none")]
        pipeline_mode: Option<String>,
        #[serde(rename = "aiSpeaksFirst", skip_serializing_if = "Option::is_none")]
        ai_speaks_first: Option<bool>,
        #[serde(rename = "allowHarmCategory", skip_serializing_if = "Option::is_none")]
        allow_harm_category: Option<bool>,
        #[serde(skip_serializing_if = "Option::is_none")]
        tools: Option<Vec<Tool>>,
        #[serde(rename = "sessionMaxDurationSeconds", skip_serializing_if = "Option::is_none")]
        session_max_duration_seconds: Option<u32>,
        #[serde(rename = "enableSessionWarning", skip_serializing_if = "Option::is_none")]
        enable_session_warning: Option<bool>,
        #[serde(rename = "enableSessionGoodbye", skip_serializing_if = "Option::is_none")]
        enable_session_goodbye: Option<bool>,
    },
    /// 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,
    /// Send a system message to the AI during a live session
    #[serde(rename = "systemMessage")]
    SystemMessage {
        /// System message payload
        payload: SystemMessagePayload,
    },
    /// Send tool response (function execution result) back to AI
    #[serde(rename = "toolResponse")]
    ToolResponse {
        /// Tool response payload
        payload: ToolResponsePayload,
    },
    /// Update user ID (guest-to-user migration)
    #[serde(rename = "updateUserId")]
    UpdateUserId {
        /// The authenticated user's unique identifier
        #[serde(rename = "userId")]
        user_id: String,
    },
    /// Explicit interrupt (stop AI response)
    #[serde(rename = "interrupt")]
    Interrupt,
    /// 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,
    },
    /// Ready - session is ready for audio input
    #[serde(rename = "ready")]
    Ready { timestamp: String },
    /// User transcript - user's speech transcription
    #[serde(rename = "userTranscript")]
    UserTranscript {
        text: String,
        timestamp: String,
    },
    /// Response - AI's text response
    #[serde(rename = "response")]
    Response {
        text: String,
        #[serde(rename = "isFinal")]
        is_final: bool,
        timestamp: String,
    },
    /// Audio - AI's audio response
    #[serde(rename = "audio")]
    Audio {
        /// Base64 encoded PCM16 audio data
        data: String,
        format: String,
        #[serde(rename = "sampleRate")]
        sample_rate: u32,
        timestamp: String,
    },
    /// Turn complete - AI has finished its response turn
    #[serde(rename = "turnComplete")]
    TurnComplete { timestamp: String },
    /// Error message
    #[serde(rename = "error")]
    Error {
        code: String,
        message: String,
        timestamp: String,
    },
    /// Tool call from AI (function calling)
    #[serde(rename = "toolCall")]
    ToolCall {
        /// Unique ID for this tool call
        id: String,
        /// Function name to execute
        name: String,
        /// Function arguments
        #[serde(default)]
        args: serde_json::Value,
        timestamp: String,
    },
    /// User ID updated (guest-to-user migration complete)
    #[serde(rename = "userIdUpdated")]
    UserIdUpdated {
        /// The new user ID that was set
        #[serde(rename = "userId")]
        user_id: String,
        /// Number of messages migrated from guest to user partition
        #[serde(rename = "migratedMessages", default)]
        migrated_messages: usize,
        timestamp: String,
    },
    /// Interrupted - AI response was interrupted by user speech (barge-in)
    #[serde(rename = "interrupted")]
    Interrupted {
        timestamp: String,
    },
    /// Pong response
    #[serde(rename = "pong")]
    Pong { timestamp: String },
    /// Session warning - nearing duration limit
    #[serde(rename = "session_warning")]
    SessionWarning {
        #[serde(rename = "remainingSeconds")]
        remaining_seconds: u32,
        #[serde(default)]
        timestamp: Option<String>,
    },
    /// Session goodbye - about to end
    #[serde(rename = "session_goodbye")]
    SessionGoodbye {
        #[serde(rename = "remainingSeconds")]
        remaining_seconds: u32,
        #[serde(default)]
        timestamp: Option<String>,
    },
}

/// Parameters for creating a StartSession message.
#[derive(Debug, Clone, Default)]
pub struct StartSessionParams {
    pub pre_prompt: Option<String>,
    pub language: Option<String>,
    pub output_language: Option<String>,
    pub pipeline_mode: Option<String>,
    pub ai_speaks_first: Option<bool>,
    pub allow_harm_category: Option<bool>,
    pub tools: Option<Vec<Tool>>,
    pub session_duration: Option<SessionDurationConfig>,
}

impl ClientMessage {
    /// Create a start session message
    pub fn start_session(params: StartSessionParams) -> Self {
        let (session_max_duration_seconds, enable_session_warning, enable_session_goodbye) =
            match params.session_duration {
                Some(duration) => (
                    Some(duration.max_seconds),
                    Some(duration.enable_warning.unwrap_or(false)),
                    Some(duration.enable_goodbye.unwrap_or(false)),
                ),
                None => (None, None, None),
            };
        ClientMessage::StartSession {
            pre_prompt: params.pre_prompt,
            language: params.language,
            output_language: params.output_language,
            pipeline_mode: params.pipeline_mode,
            ai_speaks_first: params.ai_speaks_first,
            allow_harm_category: params.allow_harm_category,
            tools: params.tools,
            session_max_duration_seconds,
            enable_session_warning,
            enable_session_goodbye,
        }
    }

    /// 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
    }

    /// Create a system message (AI responds immediately)
    pub fn system_message(text: impl Into<String>) -> Self {
        ClientMessage::SystemMessage {
            payload: SystemMessagePayload {
                text: text.into(),
                trigger_response: None, // defaults to true on server
            },
        }
    }

    /// Create a system message with explicit trigger_response option
    pub fn system_message_with_options(text: impl Into<String>, trigger_response: bool) -> Self {
        ClientMessage::SystemMessage {
            payload: SystemMessagePayload {
                text: text.into(),
                trigger_response: Some(trigger_response),
            },
        }
    }

    /// Create a tool response message (function execution result)
    pub fn tool_response(id: impl Into<String>, response: serde_json::Value) -> Self {
        ClientMessage::ToolResponse {
            payload: ToolResponsePayload {
                id: id.into(),
                response,
            },
        }
    }

    /// Create an update user ID message (guest-to-user migration)
    pub fn update_user_id(user_id: impl Into<String>) -> Self {
        ClientMessage::UpdateUserId {
            user_id: user_id.into(),
        }
    }

    /// Create an interrupt message (explicit stop)
    pub fn interrupt() -> Self {
        ClientMessage::Interrupt
    }

    /// 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(StartSessionParams {
            pre_prompt: Some("You are helpful".to_string()),
            ..Default::default()
        });
        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(StartSessionParams {
            pre_prompt: Some("You are helpful".to_string()),
            language: Some("ko-KR".to_string()),
            ..Default::default()
        });
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("prePrompt"));
        assert!(json.contains("ko-KR"));
    }

    #[test]
    fn test_serialize_start_session_with_ai_speaks_first() {
        let msg = ClientMessage::start_session(StartSessionParams {
            pre_prompt: Some("You are helpful. Greet the user.".to_string()),
            pipeline_mode: Some("live".to_string()),
            ai_speaks_first: Some(true),
            ..Default::default()
        });
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("aiSpeaksFirst"));
        assert!(json.contains("true"));
    }

    #[test]
    fn test_serialize_start_session_with_allow_harm_category() {
        let msg = ClientMessage::start_session(StartSessionParams {
            pre_prompt: Some("You are helpful.".to_string()),
            allow_harm_category: Some(false),
            ..Default::default()
        });
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("allowHarmCategory"));
        assert!(json.contains("false"));
    }

    #[test]
    fn test_serialize_start_session_with_duration_limits() {
        let msg = ClientMessage::start_session(StartSessionParams {
            pre_prompt: Some("You are helpful.".to_string()),
            session_duration: Some(SessionDurationConfig {
                max_seconds: 3600,
                enable_warning: None,
                enable_goodbye: None,
            }),
            ..Default::default()
        });
        let json = msg.to_json().unwrap();
        assert!(json.contains("startSession"));
        assert!(json.contains("sessionMaxDurationSeconds"));
        assert!(json.contains("\"enableSessionWarning\":false"));
        assert!(json.contains("\"enableSessionGoodbye\":false"));
    }

    #[test]
    fn test_serialize_start_session_with_duration_enables() {
        let msg = ClientMessage::start_session(StartSessionParams {
            pre_prompt: Some("You are helpful.".to_string()),
            session_duration: Some(SessionDurationConfig {
                max_seconds: 600,
                enable_warning: Some(true),
                enable_goodbye: Some(false),
            }),
            ..Default::default()
        });
        let json = msg.to_json().unwrap();
        assert!(json.contains("enableSessionWarning"));
        assert!(json.contains("enableSessionGoodbye"));
    }

    #[test]
    fn test_serialize_tool_response() {
        let msg = ClientMessage::tool_response("call_123", serde_json::json!({"success": true}));
        let json = msg.to_json().unwrap();
        assert!(json.contains("toolResponse"));
        assert!(json.contains("call_123"));
        assert!(json.contains("success"));
    }

    #[test]
    fn test_deserialize_tool_call() {
        let json = r#"{"type":"toolCall","id":"call_abc","name":"open_login","args":{},"timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::ToolCall { id, name, .. } => {
                assert_eq!(id, "call_abc");
                assert_eq!(name, "open_login");
            }
            _ => panic!("Expected ToolCall message"),
        }
    }

    #[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_ready() {
        let json = r#"{"type":"ready","timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::Ready { .. } => {}
            _ => panic!("Expected Ready message"),
        }
    }

    #[test]
    fn test_deserialize_user_transcript() {
        let json = r#"{"type":"userTranscript","text":"Hello world","timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::UserTranscript { text, .. } => {
                assert_eq!(text, "Hello world");
            }
            _ => panic!("Expected UserTranscript message"),
        }
    }

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

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

    #[test]
    fn test_serialize_update_user_id() {
        let msg = ClientMessage::update_user_id("user_abc123");
        let json = msg.to_json().unwrap();
        assert!(json.contains("updateUserId"));
        assert!(json.contains("user_abc123"));
    }

    #[test]
    fn test_deserialize_user_id_updated() {
        let json = r#"{"type":"userIdUpdated","userId":"user_abc123","migratedMessages":12,"timestamp":"2024-01-01T00:00:00Z"}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::UserIdUpdated { user_id, migrated_messages, .. } => {
                assert_eq!(user_id, "user_abc123");
                assert_eq!(migrated_messages, 12);
            }
            _ => panic!("Expected UserIdUpdated message"),
        }
    }

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

    #[test]
    fn test_deserialize_session_warning() {
        let json = r#"{"type":"session_warning","remainingSeconds":30}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::SessionWarning { remaining_seconds, .. } => {
                assert_eq!(remaining_seconds, 30);
            }
            _ => panic!("Expected SessionWarning message"),
        }
    }

    #[test]
    fn test_deserialize_session_goodbye() {
        let json = r#"{"type":"session_goodbye","remainingSeconds":20}"#;
        let msg = ServerMessage::from_json(json).unwrap();
        match msg {
            ServerMessage::SessionGoodbye { remaining_seconds, .. } => {
                assert_eq!(remaining_seconds, 20);
            }
            _ => panic!("Expected SessionGoodbye message"),
        }
    }
}
