src/socket/mod.rs

use serde::Deserialize;
use serde_json::Value as JsonValue;
use socketioxide::extract::{Data, SocketRef};
use socketioxide::SocketIo;
use sqlx::PgPool;
use tracing::info;
use uuid::Uuid;

use crate::services::push::PushService;
use crate::utils::jwt;

#[derive(Debug, Deserialize)]
struct ConvPayload {
    conversation_id: String,
}

#[derive(Debug, Deserialize)]
struct SendMsgPayload {
    conversation_id: String,
    text: Option<String>,
    message_type: Option<String>,
    file_name: Option<String>,
    thumbnail_url: Option<String>,
}

pub fn attach_socket(io: &SocketIo, pool: PgPool, jwt_secret: String) {
    let io_owned = io.clone();
    io.ns("/", move |socket: SocketRef, Data(auth): Data<JsonValue>| {
        let sid = socket.id;
        let io = io_owned.clone();

        let token = auth
            .get("token")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        let user_id = match jwt::user_id_from_token(token, &jwt_secret) {
            Ok(uid) => uid,
            Err(_) => {
                info!("Socket {sid} rejected: invalid token");
                let _ = socket.emit("error", &serde_json::json!({"message": "Invalid token"}));
                return;
            }
        };

        socket.join(user_id.to_string());
        info!("Socket {sid} authenticated as user {user_id}");

        socket.on("join_conversation", |socket: SocketRef, Data(data): Data<ConvPayload>| {
            let room = format!("conv_{}", data.conversation_id);
            socket.join(room);
        });

        socket.on("leave_conversation", |socket: SocketRef, Data(data): Data<ConvPayload>| {
            let room = format!("conv_{}", data.conversation_id);
            socket.leave(room);
        });

        let send_pool = pool.clone();
        let send_uid = user_id;
        let send_socket = socket.clone();
        let send_io = io.clone();
        socket.on("send_message", move |_socket: SocketRef, Data(data): Data<SendMsgPayload>| {
            let pool = send_pool.clone();
            let uid = send_uid;
            let s = send_socket.clone();
            let io = send_io.clone();
            async move {
                let conv_id = match Uuid::parse_str(&data.conversation_id) {
                    Ok(id) => id,
                    Err(_) => return,
                };
                let text = data.text.unwrap_or_default();
                let msg_type = data.message_type.unwrap_or_else(|| "text".into());
                let file_name = data.file_name;
                let thumbnail_url = data.thumbnail_url;

                match sqlx::query_as::<_, crate::models::Message>(
                    r#"INSERT INTO messages (conversation_id, sender_id, text, message_type, file_name, thumbnail_url)
                       VALUES ($1, $2, $3, $4, $5, $6) RETURNING *"#,
                )
                .bind(conv_id)
                .bind(uid)
                .bind(&text)
                .bind(&msg_type)
                .bind(&file_name)
                .bind(&thumbnail_url)
                .fetch_one(&pool)
                .await
                {
                    Ok(message) => {
                        let _ = sqlx::query(
                            r#"UPDATE conversations SET
                                   last_message_text = $1,
                                   last_message_sender_id = $2,
                                   last_message_at = NOW()
                               WHERE id = $3"#,
                        )
                        .bind(&text)
                        .bind(uid)
                        .bind(conv_id)
                        .execute(&pool)
                        .await;

                        // Update unread_counts + notify other participants
                        let others = sqlx::query_scalar::<_, Uuid>(
                            "SELECT user_id FROM conversation_participants WHERE conversation_id = $1 AND user_id != $2"
                        )
                        .bind(conv_id)
                        .bind(uid)
                        .fetch_all(&pool)
                        .await
                        .unwrap_or_default();

                        for other_uid in &others {
                            let _ = sqlx::query(
                                r#"UPDATE conversations SET unread_counts =
                                   COALESCE(unread_counts, '{}'::jsonb) ||
                                   jsonb_build_object($1::text,
                                       COALESCE((unread_counts->>$1)::int, 0) + 1)
                                   WHERE id = $2"#,
                            )
                            .bind(other_uid.to_string())
                            .bind(conv_id)
                            .execute(&pool)
                            .await;
                        }

                        // Create in-app notification + push for each other participant
                        if !others.is_empty() {
                            let sender_name = sqlx::query_scalar::<_, String>(
                                "SELECT COALESCE(name, email) FROM users WHERE id = $1"
                            )
                            .bind(uid)
                            .fetch_optional(&pool)
                            .await
                            .unwrap_or(None)
                            .unwrap_or_else(|| "Someone".into());

                            let preview = if text.len() > 100 {
                                format!("{}...", &text[..100])
                            } else {
                                text.clone()
                            };

                            for other_uid in &others {
                                let _ = sqlx::query(
                                    r#"INSERT INTO notifications
                                       (recipient_id, sender_id, title, message, notification_type, data)
                                       VALUES ($1, $2, $3, $4, 'chat', $5)"#,
                                )
                                .bind(other_uid)
                                .bind(uid)
                                .bind(format!("New message from {}", sender_name))
                                .bind(&preview)
                                .bind(serde_json::json!({
                                    "conversationId": conv_id.to_string(),
                                    "messageId": message.id.to_string(),
                                }))
                                .execute(&pool)
                                .await;

                                let _ = io
                                    .to(other_uid.to_string())
                                    .emit("new_notification", &serde_json::json!({
                                        "conversationId": conv_id.to_string(),
                                        "senderName": &sender_name,
                                        "preview": &preview,
                                    }));

                                // Send push notification
                                let push_token = sqlx::query_scalar::<_, Option<String>>(
                                    "SELECT expo_push_token FROM users WHERE id = $1"
                                )
                                .bind(other_uid)
                                .fetch_optional(&pool)
                                .await
                                .ok()
                                .flatten()
                                .flatten();

                                if let Some(ref token) = push_token {
                                    if !token.is_empty() {
                                        let push = PushService::new();
                                        let _ = push.send_push(
                                            token,
                                            &format!("New message from {}", sender_name),
                                            &preview,
                                            Some(serde_json::json!({
                                                "conversationId": conv_id.to_string(),
                                                "senderName": &sender_name,
                                                "type": "chat",
                                            })),
                                        ).await;
                                    }
                                }
                            }
                        }

                        // Emit new_message to conversation room
                        let room = format!("conv_{}", conv_id);
                        let _ = s
                            .to(room)
                            .emit(
                                "new_message",
                                &serde_json::json!({
                                    "message": message,
                                    "conversationId": conv_id.to_string(),
                                }),
                            );

                        // Emit conversation_updated to each other participant's private room
                        for other_uid in &others {
                            let _ = io
                                .to(other_uid.to_string())
                                .emit(
                                    "conversation_updated",
                                    &serde_json::json!({
                                        "conversationId": conv_id.to_string(),
                                    }),
                                );
                        }
                    }
                    Err(e) => {
                        let _ = s
                            .emit("error", &serde_json::json!({"message": format!("{e}")}));
                    }
                }
            }
        });

        let mark_pool = pool.clone();
        let mark_uid = user_id;
        socket.on("mark_read", move |_socket: SocketRef, Data(data): Data<ConvPayload>| {
            let pool = mark_pool.clone();
            let uid = mark_uid;
            async move {
                if let Ok(conv_id) = Uuid::parse_str(&data.conversation_id) {
                    let _ = sqlx::query(
                        r#"UPDATE messages SET is_read = true
                           WHERE conversation_id = $1 AND sender_id != $2 AND is_read = false"#,
                    )
                    .bind(conv_id)
                    .bind(uid)
                    .execute(&pool)
                    .await;

                    let _ = sqlx::query(
                        r#"UPDATE conversations SET unread_counts =
                           unread_counts || jsonb_build_object($2::text, 0)
                           WHERE id = $1"#,
                    )
                    .bind(conv_id)
                    .bind(uid.to_string())
                    .execute(&pool)
                    .await;
                }
            }
        });

        socket.on("typing", |socket: SocketRef, Data(data): Data<ConvPayload>| {
            let room = format!("conv_{}", data.conversation_id);
            let _ = socket
                .to(room)
                .emit(
                    "user_typing",
                    &serde_json::json!({"conversation_id": data.conversation_id}),
                );
        });

        socket.on("stop_typing", |socket: SocketRef, Data(data): Data<ConvPayload>| {
            let room = format!("conv_{}", data.conversation_id);
            let _ = socket
                .to(room)
                .emit(
                    "user_stop_typing",
                    &serde_json::json!({"conversation_id": data.conversation_id}),
                );
        });

        socket.on("disconnect", move |_socket: SocketRef| {
            info!("Socket {sid} disconnected");
        });
    });
}