janichat/backend/src/ws.ts

import { FastifyInstance } from 'fastify';
import { pool } from './db.js';
import webpush from 'web-push';

// uid -> Set of WebSocket-like connections
export const connections = new Map<string, Set<any>>();

function broadcast(uids: string[], data: object) {
  const msg = JSON.stringify(data);
  for (const uid of uids) {
    const sockets = connections.get(uid);
    if (sockets) {
      sockets.forEach((ws: any) => {
        try {
          if (ws.readyState === 1) ws.send(msg);
        } catch {}
      });
    }
  }
}

async function getChatMemberIds(chatId: string): Promise<string[]> {
  const { rows } = await pool.query(
    'SELECT user_id FROM chat_members WHERE chat_id = $1', [chatId]
  );
  return rows.map((r: any) => r.user_id);
}

async function sendPushToOfflineUsers(uids: string[], payload: object) {
  const online = new Set(connections.keys());
  const offline = uids.filter(id => !online.has(id));
  if (offline.length === 0) return;

  const { rows } = await pool.query(
    'SELECT endpoint, auth_key, p256dh FROM push_subscriptions WHERE user_id = ANY($1)',
    [offline]
  );

  const body = JSON.stringify(payload);
  for (const sub of rows) {
    try {
      await webpush.sendNotification(
        { endpoint: sub.endpoint, keys: { auth: sub.auth_key, p256dh: sub.p256dh } },
        body
      );
    } catch (e: any) {
      if (e.statusCode === 410) {
        await pool.query('DELETE FROM push_subscriptions WHERE endpoint = $1', [sub.endpoint]);
      }
    }
  }
}

export function setupWebSocket(app: FastifyInstance) {
  (app as any).get('/ws', { websocket: true }, (connection: any, req: any) => {
    // @fastify/websocket v8 passes a Duplex stream — actual ws.WebSocket is at connection.socket
    let uid: string;
    try {
      const token = req.query?.token;
      const decoded = app.jwt.verify<{ id: string }>(token);
      uid = decoded.id;
    } catch {
      connection.socket.close(1008, 'Unauthorized');
      return;
    }

    const ws = connection.socket; // actual ws.WebSocket

    // Register connection (store ws.WebSocket for broadcast)
    if (!connections.has(uid)) connections.set(uid, new Set());
    connections.get(uid)!.add(ws);

    // Update last_seen + notify online (fire-and-forget)
    pool.query('UPDATE users SET last_seen = NOW() WHERE id = $1', [uid]).catch(() => {});
    pool.query('SELECT DISTINCT chat_id FROM chat_members WHERE user_id = $1', [uid]).then(async ({ rows: memberChats }) => {
      for (const row of memberChats) {
        const ids = await getChatMemberIds(row.chat_id);
        broadcast(ids.filter(id => id !== uid), { type: 'user_online', payload: { userId: uid, online: true } });
      }
    }).catch(() => {});

    ws.on('message', async (raw: any, isBinary: boolean) => {
      try {
        const str = isBinary ? raw.toString('utf8') : raw.toString();
        console.log('[WS] received:', str.substring(0, 200));
        const { type, payload } = JSON.parse(str);

        if (type === 'typing') {
          const members = await getChatMemberIds(payload.chatId);
          broadcast(members.filter(id => id !== uid), {
            type: 'typing',
            payload: { chatId: payload.chatId, userId: uid, typing: payload.typing }
          });
        }

        if (type === 'read_messages') {
          await pool.query(
            `UPDATE chat_members SET last_read_at = NOW() WHERE chat_id = $1 AND user_id = $2`,
            [payload.chatId, uid]
          );
          const members = await getChatMemberIds(payload.chatId);
          broadcast(members.filter(id => id !== uid), {
            type: 'messages_read',
            payload: { chatId: payload.chatId, userId: uid }
          });
        }

        if (type === 'send_message') {
          const { rows: [member] } = await pool.query(
            'SELECT role, can_send_messages FROM chat_members WHERE chat_id = $1 AND user_id = $2',
            [payload.chatId, uid]
          );
          if (!member) return;

          const { rows: [chat] } = await pool.query('SELECT type FROM chats WHERE id = $1', [payload.chatId]);
          if (chat.type === 'channel' && !['owner', 'admin'].includes(member.role)) return;
          if (member.can_send_messages === false && member.role === 'member') return;

          const { rows: [msg] } = await pool.query(
            `INSERT INTO messages (chat_id, sender_id, content, type, reply_to_id)
             VALUES ($1, $2, $3, $4, $5)
             RETURNING *`,
            [payload.chatId, uid, payload.content, payload.type || 'text', payload.replyToId || null]
          );

          const { rows: [sender] } = await pool.query(
            'SELECT id, username, display_name, avatar_color FROM users WHERE id = $1', [uid]
          );

          const fullMsg = {
            id: msg.id, chatId: msg.chat_id, content: msg.content, type: msg.type,
            isDeleted: false, isEdited: false, createdAt: msg.created_at,
            replyTo: null, attachments: [],
            sender: { id: sender.id, username: sender.username, displayName: sender.display_name, avatarColor: sender.avatar_color }
          };

          const members = await getChatMemberIds(payload.chatId);
          broadcast(members, { type: 'new_message', payload: fullMsg });

          await sendPushToOfflineUsers(members, {
            title: sender.display_name,
            body: payload.content.substring(0, 100),
            chatId: payload.chatId
          });

          await pool.query(
            'UPDATE chat_members SET last_read_at = NOW() WHERE chat_id = $1 AND user_id = $2',
            [payload.chatId, uid]
          );
        }

        if (type === 'delete_message') {
          const { rows: [msg] } = await pool.query(
            'SELECT sender_id, chat_id FROM messages WHERE id = $1', [payload.messageId]
          );
          if (!msg) return;

          const { rows: [member] } = await pool.query(
            'SELECT role FROM chat_members WHERE chat_id = $1 AND user_id = $2', [msg.chat_id, uid]
          );
          if (msg.sender_id !== uid && !['owner', 'admin'].includes(member?.role)) return;

          await pool.query('UPDATE messages SET is_deleted = TRUE WHERE id = $1', [payload.messageId]);
          const members = await getChatMemberIds(msg.chat_id);
          broadcast(members, { type: 'message_deleted', payload: { messageId: payload.messageId, chatId: msg.chat_id } });
        }

        if (type === 'edit_message') {
          const { rows: [msg] } = await pool.query(
            'SELECT sender_id, chat_id FROM messages WHERE id = $1', [payload.messageId]
          );
          if (!msg || msg.sender_id !== uid) return;

          const { rows: [updated] } = await pool.query(
            `UPDATE messages SET content = $1, is_edited = TRUE, edited_at = NOW() WHERE id = $2
             RETURNING id, chat_id, sender_id, content, type, is_edited, edited_at, created_at`,
            [payload.content, payload.messageId]
          );

          const { rows: [sender] } = await pool.query(
            'SELECT id, username, display_name, avatar_color FROM users WHERE id = $1', [uid]
          );

          const fullMsg = {
            id: updated.id, chatId: updated.chat_id, content: updated.content, type: updated.type,
            isDeleted: false, isEdited: updated.is_edited, editedAt: updated.edited_at,
            createdAt: updated.created_at, replyTo: null, attachments: [],
            sender: { id: sender.id, username: sender.username, displayName: sender.display_name, avatarColor: sender.avatar_color }
          };

          const members = await getChatMemberIds(msg.chat_id);
          broadcast(members, { type: 'message_edited', payload: fullMsg });
        }

      } catch (e) {
        console.error('WS message error:', e);
      }
    });

    ws.on('close', async () => {
      const set = connections.get(uid);
      if (set) {
        set.delete(ws);
        if (set.size === 0) {
          connections.delete(uid);
          await pool.query('UPDATE users SET last_seen = NOW() WHERE id = $1', [uid]);

          const { rows: memberChats } = await pool.query(
            'SELECT DISTINCT chat_id FROM chat_members WHERE user_id = $1', [uid]
          );
          for (const row of memberChats) {
            const ids = await getChatMemberIds(row.chat_id);
            broadcast(ids, { type: 'user_online', payload: { uid: uid, online: false } });
          }
        }
      }
    });
  });
}