sw-fe/src/components/Chat/useChat.js

import { useState, useEffect, useRef, useCallback } from "react";
import {
  getOrCreateKeyPair,
  storePeerPublicKey,
  getPeerPublicKey,
  encryptMessage,
  decryptMessage,
  isE2EAvailable,
} from "./cryptoUtils";

const WS_BASE = import.meta.env.VITE_CHAT_WS_URL ||
  (window.location.protocol === "https:" ? "wss://" : "ws://") +
  window.location.host + "/chat/ws";

const API_BASE = import.meta.env.VITE_CHAT_API_URL || "/chat/api";

export function useChat({ username, enabled = true, onNewMessage }) {
  const [connected, setConnected] = useState(false);
  const [messages, setMessages] = useState({}); // { recipientUsername: Message[] }
  const [typing, setTyping] = useState({}); // { username: boolean }
  const [presence, setPresence] = useState({}); // { username: boolean }
  const [conversations, setConversations] = useState([]);
  const [e2eReady, setE2eReady] = useState(false);

  const wsRef = useRef(null);
  const reconnectTimer = useRef(null);
  const reconnectAttempts = useRef(0);
  const keyPairRef = useRef(null);
  const usernameRef = useRef(username);
  const enabledRef = useRef(enabled);
  const onNewMessageRef = useRef(onNewMessage);

  // Keep refs updated
  useEffect(() => {
    usernameRef.current = username;
  }, [username]);

  useEffect(() => {
    enabledRef.current = enabled;
  }, [enabled]);

  useEffect(() => {
    onNewMessageRef.current = onNewMessage;
  }, [onNewMessage]);

  // Initialize E2E encryption keys
  useEffect(() => {
    if (!isE2EAvailable()) {
      console.warn("[E2E] Web Crypto API not available");
      return;
    }

    getOrCreateKeyPair().then((kp) => {
      keyPairRef.current = kp;
      setE2eReady(true);
      console.log("[E2E] Key pair ready");
    }).catch((err) => {
      console.error("[E2E] Failed to initialize keys", err);
    });
  }, []);

  // Send public key to a peer for key exchange
  const sendPublicKey = useCallback((recipientUsername) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;
    if (!keyPairRef.current?.publicKeyExport) return;

    const msg = {
      type: "key_exchange",
      payload: {
        recipient_username: recipientUsername,
        public_key: keyPairRef.current.publicKeyExport,
      },
    };
    wsRef.current.send(JSON.stringify(msg));
  }, []);

  // Handle incoming WebSocket messages (uses ref to avoid dependency changes)
  const handleMessage = useCallback((data) => {
    const currentUsername = usernameRef.current;

    // Parse payload if it's a string (backwards compat), otherwise use as-is
    const parsePayload = (payload) => {
      if (typeof payload === "string") {
        try {
          return JSON.parse(payload);
        } catch {
          return payload;
        }
      }
      return payload;
    };

    switch (data.type) {
      case "message": {
        const msg = parsePayload(data.payload);
        // Case-insensitive comparison for username matching
        const isMySent = msg.sender_name?.toLowerCase() === currentUsername?.toLowerCase();
        const otherUser = isMySent ? msg.recipient_name : msg.sender_name;
        const isIncoming = !isMySent;

        console.log("[Chat] Received message from", msg.sender_name, "to", msg.recipient_name);

        setMessages((prev) => {
          const existing = prev[otherUser] || [];
          // Avoid duplicates (check by id or by temp-id prefix)
          if (existing.some((m) => m.id === msg.id || (m.id?.toString().startsWith("temp-") && m.content === msg.content && m.sender_name === msg.sender_name))) {
            // Replace temp message with real one
            return {
              ...prev,
              [otherUser]: existing.map((m) =>
                m.id?.toString().startsWith("temp-") && m.content === msg.content && m.sender_name === msg.sender_name
                  ? msg
                  : m
              ),
            };
          }
          return { ...prev, [otherUser]: [...existing, msg] };
        });

        // Notify parent component of new incoming message
        if (isIncoming && onNewMessageRef.current) {
          onNewMessageRef.current(msg);
        }
        break;
      }

      case "typing": {
        const payload = parsePayload(data.payload);
        setTyping((prev) => ({
          ...prev,
          [payload.username]: payload.is_typing,
        }));

        // Clear typing indicator after 3 seconds
        setTimeout(() => {
          setTyping((prev) => ({ ...prev, [payload.username]: false }));
        }, 3000);
        break;
      }

      case "presence": {
        const payload = parsePayload(data.payload);
        setPresence((prev) => ({
          ...prev,
          [payload.username]: payload.online,
        }));
        break;
      }

      default:
        console.log("[Chat] Unknown message type", data.type);
    }
  }, []); // No dependencies - uses refs

  // Connect to WebSocket (uses refs to avoid dependency changes)
  const connect = useCallback(() => {
    const currentUsername = usernameRef.current;
    const currentEnabled = enabledRef.current;

    if (!currentEnabled || !currentUsername) return;

    // Don't reconnect if already connected
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      console.log("[Chat] Already connected");
      return;
    }

    const wsUrl = `${WS_BASE}?username=${encodeURIComponent(currentUsername)}`;
    console.log("[Chat] Connecting to", wsUrl);

    try {
      const ws = new WebSocket(wsUrl);

      ws.onopen = () => {
        console.log("[Chat] WebSocket connected as", currentUsername);
        setConnected(true);
        reconnectAttempts.current = 0;
      };

      ws.onclose = (e) => {
        console.log("[Chat] WebSocket closed", e.code, e.reason);
        setConnected(false);
        wsRef.current = null;

        // Reconnect with exponential backoff (check ref values)
        if (enabledRef.current && reconnectAttempts.current < 10) {
          const delay = Math.min(1000 * Math.pow(2, reconnectAttempts.current), 30000);
          reconnectAttempts.current++;
          console.log("[Chat] Reconnecting in", delay, "ms (attempt", reconnectAttempts.current, ")");
          reconnectTimer.current = setTimeout(connect, delay);
        }
      };

      ws.onerror = (e) => {
        console.error("[Chat] WebSocket error", e);
      };

      ws.onmessage = (e) => {
        try {
          const data = JSON.parse(e.data);
          handleMessage(data);
        } catch (err) {
          console.error("[Chat] Parse error", err);
        }
      };

      wsRef.current = ws;
    } catch (err) {
      console.error("[Chat] Connection error", err);
    }
  }, [handleMessage]); // Only depends on stable handleMessage

  // Send a chat message
  const sendMessage = useCallback((recipientUsername, content) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      console.error("[Chat] WebSocket not connected");
      return false;
    }

    const currentUsername = usernameRef.current;

    // Optimistic update - add message to local state immediately
    const optimisticMsg = {
      id: `temp-${Date.now()}`,
      sender_name: currentUsername,
      recipient_name: recipientUsername,
      content: content,
      created_at: new Date().toISOString(),
    };

    setMessages((prev) => {
      const existing = prev[recipientUsername] || [];
      return { ...prev, [recipientUsername]: [...existing, optimisticMsg] };
    });

    const msg = {
      type: "message",
      payload: {
        recipient_username: recipientUsername,
        content: content,
      },
    };

    wsRef.current.send(JSON.stringify(msg));
    return true;
  }, []);

  // Send typing indicator
  const sendTyping = useCallback((recipientUsername, isTyping) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;

    const msg = {
      type: "typing",
      payload: {
        recipient_username: recipientUsername,
        is_typing: isTyping,
      },
    };

    wsRef.current.send(JSON.stringify(msg));
  }, []);

  // Mark messages as read
  const markAsRead = useCallback((senderUsername) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;

    const msg = {
      type: "read",
      payload: { sender_username: senderUsername },
    };

    wsRef.current.send(JSON.stringify(msg));
  }, []);

  // Fetch message history
  const fetchMessages = useCallback(async (otherUsername) => {
    if (!username || !otherUsername) return [];

    try {
      const res = await fetch(`${API_BASE}/messages?username=${encodeURIComponent(username)}&other=${encodeURIComponent(otherUsername)}`);
      const data = await res.json();

      if (data.ok && data.messages) {
        setMessages((prev) => ({ ...prev, [otherUsername]: data.messages }));
        return data.messages;
      }
    } catch (err) {
      console.error("[Chat] Fetch messages error", err);
    }
    return [];
  }, [username]);

  // Fetch conversations list
  const fetchConversations = useCallback(async () => {
    if (!username) return [];

    try {
      const res = await fetch(`${API_BASE}/conversations?username=${encodeURIComponent(username)}`);
      const data = await res.json();

      if (data.ok && data.conversations) {
        setConversations(data.conversations);

        // Update presence from conversations
        const newPresence = {};
        data.conversations.forEach((c) => {
          newPresence[c.username] = c.online;
        });
        setPresence((prev) => ({ ...prev, ...newPresence }));

        return data.conversations;
      }
    } catch (err) {
      console.error("[Chat] Fetch conversations error", err);
    }
    return [];
  }, [username]);

  // Connect on mount / when enabled or username changes
  useEffect(() => {
    if (enabled && username) {
      // Only connect if not already connected
      if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
        connect();
      }
    } else {
      // Disconnect if disabled or no username
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
    }
  }, [enabled, username, connect]);

  // Reconnect when page becomes visible (mobile app resume)
  useEffect(() => {
    if (!enabled || !username) return;

    const handleVisibilityChange = () => {
      if (document.visibilityState === 'visible') {
        console.log('[Chat] Page visible, checking connection...');
        // Check if WebSocket is disconnected or in bad state
        if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
          console.log('[Chat] Reconnecting after visibility change...');
          reconnectAttempts.current = 0; // Reset attempts
          connect();
        }
      }
    };

    document.addEventListener('visibilitychange', handleVisibilityChange);
    return () => {
      document.removeEventListener('visibilitychange', handleVisibilityChange);
    };
  }, [enabled, username, connect]);

  // Cleanup on unmount only
  useEffect(() => {
    return () => {
      if (reconnectTimer.current) {
        clearTimeout(reconnectTimer.current);
      }
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
    };
  }, []);

  return {
    connected,
    messages,
    typing,
    presence,
    conversations,
    sendMessage,
    sendTyping,
    markAsRead,
    fetchMessages,
    fetchConversations,
  };
}