//! WebSocket streaming for real-time LLM responses
//!
//! Provides token-by-token streaming of assistant responses via WebSocket.

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Path, State,
    },
    response::IntoResponse,
};
use serde::{Deserialize, Serialize};
use serde_json::json;

use super::types::AppState;
use crate::assistant::ConfigAssistant;

/// WebSocket message from client
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WsMessage {
    /// Message type: "message", "suggestion", "generate", etc.
    pub r#type: String,

    /// Message content (user message or other data)
    pub content: String,

    /// Additional data (e.g., format for config generation)
    pub data: Option<serde_json::Value>,
}

/// WebSocket response sent to client
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WsResponse {
    /// Response type: "start", "chunk", "end", "error", "suggestion", etc.
    pub r#type: String,

    /// Response content
    pub content: String,

    /// Optional metadata (e.g., message_id, confidence)
    pub metadata: Option<serde_json::Value>,
}

impl WsResponse {
    /// Create a streaming chunk response
    pub fn chunk(content: impl Into<String>) -> Self {
        WsResponse {
            r#type: "chunk".to_string(),
            content: content.into(),
            metadata: None,
        }
    }

    /// Create a start response
    pub fn start() -> Self {
        WsResponse {
            r#type: "start".to_string(),
            content: String::new(),
            metadata: None,
        }
    }

    /// Create an end response with metadata
    pub fn end(metadata: Option<serde_json::Value>) -> Self {
        WsResponse {
            r#type: "end".to_string(),
            content: String::new(),
            metadata,
        }
    }

    /// Create an error response
    pub fn error(error: impl Into<String>) -> Self {
        WsResponse {
            r#type: "error".to_string(),
            content: error.into(),
            metadata: None,
        }
    }

    /// Serialize to JSON
    pub fn to_json(&self) -> String {
        serde_json::to_string(self).unwrap_or_else(|_| {
            serde_json::to_string(&WsResponse::error("Serialization failed")).unwrap()
        })
    }
}

/// Handle WebSocket upgrade and streaming
///
/// Accepts a WebSocket connection and streams LLM responses in real-time.
/// The client should send JSON messages with:
/// - type: "message", "generate", or "suggestions"
/// - content: the user input
/// - data: optional additional data (e.g., format for generation)
pub async fn handle_websocket(
    ws: WebSocketUpgrade,
    Path(conversation_id): Path<String>,
    State(state): State<AppState>,
) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_socket(socket, conversation_id, state))
}

async fn handle_socket(mut socket: WebSocket, conversation_id: String, state: AppState) {
    // Get or create assistant for this conversation
    let mut assistant_opt = None;

    while let Some(msg_result) = socket.recv().await {
        match msg_result {
            Ok(Message::Text(text)) => {
                // Parse incoming message
                match serde_json::from_str::<WsMessage>(&text) {
                    Ok(ws_msg) => {
                        // Handle message based on type
                        match ws_msg.r#type.as_str() {
                            "message" => {
                                handle_message_type(
                                    &mut socket,
                                    &mut assistant_opt,
                                    &conversation_id,
                                    &state,
                                    &ws_msg,
                                )
                                .await;
                            }
                            "generate" => {
                                handle_generate_type(
                                    &mut socket,
                                    &mut assistant_opt,
                                    &conversation_id,
                                    &state,
                                    &ws_msg,
                                )
                                .await;
                            }
                            "suggestions" => {
                                handle_suggestions_type(
                                    &mut socket,
                                    &mut assistant_opt,
                                    &conversation_id,
                                    &state,
                                    &ws_msg,
                                )
                                .await;
                            }
                            _ => {
                                let error_resp = WsResponse::error(format!(
                                    "Unknown message type: {}",
                                    ws_msg.r#type
                                ));
                                drop(
                                    socket
                                        .send(Message::Text(error_resp.to_json().into()))
                                        .await,
                                );
                            }
                        }
                    }
                    Err(e) => {
                        let error_resp =
                            WsResponse::error(format!("Failed to parse message: {}", e));
                        drop(
                            socket
                                .send(Message::Text(error_resp.to_json().into()))
                                .await,
                        );
                    }
                }
            }
            Ok(Message::Close(_)) => {
                tracing::debug!(conv_id = %conversation_id, "WebSocket closed by client");
                break;
            }
            Ok(_) => {
                // Ignore other message types
            }
            Err(e) => {
                tracing::error!(conv_id = %conversation_id, error = %e, "WebSocket error");
                break;
            }
        }
    }
}

async fn handle_message_type(
    socket: &mut WebSocket,
    _assistant_opt: &mut Option<ConfigAssistant>,
    _conversation_id: &str,
    _state: &AppState,
    ws_msg: &WsMessage,
) {
    // Send start marker
    let start_resp = WsResponse::start();
    if socket
        .send(Message::Text(start_resp.to_json().into()))
        .await
        .is_err()
    {
        return;
    }

    // Get or create assistant (for now, just send the message back with chunks)
    // In a full implementation, this would get the assistant from state
    // and call stream_message() on it

    // Simulate streaming response by splitting the content into words
    let words: Vec<&str> = ws_msg.content.split_whitespace().collect();

    for (i, word) in words.iter().enumerate() {
        let chunk_resp = WsResponse::chunk(if i == 0 {
            word.to_string()
        } else {
            format!(" {}", word)
        });

        if socket
            .send(Message::Text(chunk_resp.to_json().into()))
            .await
            .is_err()
        {
            return;
        }

        // Small delay to simulate token streaming
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
    }

    // Send end marker
    let end_resp = WsResponse::end(Some(json!({
        "type": "message",
        "status": "complete"
    })));
    drop(socket.send(Message::Text(end_resp.to_json().into())).await);
}

async fn handle_generate_type(
    socket: &mut WebSocket,
    _assistant_opt: &mut Option<ConfigAssistant>,
    _conversation_id: &str,
    _state: &AppState,
    _ws_msg: &WsMessage,
) {
    // Send generation response
    let resp = WsResponse::error("Configuration generation not yet implemented via WebSocket");
    drop(socket.send(Message::Text(resp.to_json().into())).await);
}

async fn handle_suggestions_type(
    socket: &mut WebSocket,
    _assistant_opt: &mut Option<ConfigAssistant>,
    _conversation_id: &str,
    _state: &AppState,
    ws_msg: &WsMessage,
) {
    // Get field name from content
    let field_name = &ws_msg.content;

    // Send suggestions response with mock data
    let suggestions = json!([
        {
            "field": field_name,
            "value": "example_value",
            "reasoning": "Found in examples",
            "confidence": 0.85
        }
    ]);

    let resp = WsResponse {
        r#type: "suggestions".to_string(),
        content: serde_json::to_string(&suggestions).unwrap_or_default(),
        metadata: Some(json!({"field": field_name})),
    };

    drop(socket.send(Message::Text(resp.to_json().into())).await);
}

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

    #[test]
    fn test_ws_message_serialization() {
        let msg = WsMessage {
            r#type: "message".to_string(),
            content: "What port?".to_string(),
            data: None,
        };
        let json = serde_json::to_string(&msg).unwrap();
        let deserialized: WsMessage = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.r#type, "message");
    }

    #[test]
    fn test_ws_response_chunk() {
        let resp = WsResponse::chunk("port");
        assert_eq!(resp.r#type, "chunk");
        assert_eq!(resp.content, "port");
    }

    #[test]
    fn test_ws_response_start() {
        let resp = WsResponse::start();
        assert_eq!(resp.r#type, "start");
    }

    #[test]
    fn test_ws_response_end() {
        let resp = WsResponse::end(None);
        assert_eq!(resp.r#type, "end");
    }

    #[test]
    fn test_ws_response_error() {
        let resp = WsResponse::error("Something failed");
        assert_eq!(resp.r#type, "error");
        assert!(resp.content.contains("failed"));
    }

    #[test]
    fn test_ws_response_to_json() {
        let resp = WsResponse::chunk("test");
        let json = resp.to_json();
        assert!(json.contains("chunk"));
        assert!(json.contains("test"));
    }

    #[test]
    fn test_ws_response_with_metadata() {
        let metadata = Some(serde_json::json!({"message_id": "123"}));
        let resp = WsResponse::end(metadata);
        assert!(resp.metadata.is_some());
    }
}
