Chat-App-Frontend/src/context/SocketContext.tsx

"use client";

import {
  createContext,
  useContext,
  useEffect,
  useState,
  useRef,
  useCallback,
  type ReactNode,
} from "react";
import { io, Socket } from "socket.io-client";
import { useAuth } from "./AuthContext";
import { API_BASE_URL } from "@/lib/api";

// Shape of a chat message coming from the server
export interface ChatMessage {
  _id?: string;
  roomId: string;
  senderId: string;
  username: string;
  text: string;
  createdAt?: string;
}

// System messages related to a room (join/leave/info)
export interface RoomNotice {
  roomId: string;
  message: string;
}

// Everything that the socket context will provide to the app
interface SocketContextValue {
  socket: Socket | null;
  isConnected: boolean;
  mySocketId: string | null;

  // join / leave / send chat messages
  joinRoom: (roomId: string) => void;
  leaveRoom: (roomId: string) => void;
  sendMessage: (roomId: string, text: string) => void;

  // chat state stored globally
  messages: ChatMessage[];
  notices: RoomNotice[];
  currentRoom: string | null;
  setCurrentRoom: (roomId: string | null) => void;
  clearMessages: () => void;
}

//context
// Create a global socket context
const SocketContext = createContext<SocketContextValue | undefined>(undefined);

/*  Provider   */

export function SocketProvider({ children }: { children: ReactNode }) {
  const { user } = useAuth();

  // Keep socket instance in a ref so it persists without rerenders
  const socketRef = useRef<Socket | null>(null);

  // Connection status
  const [isConnected, setIsConnected] = useState(false);

  // Store current socket id assigned by server
  const [mySocketId, setMySocketId] = useState<string | null>(null);

  // All chat messages received
  const [messages, setMessages] = useState<ChatMessage[]>([]);

  // System notices (join/leave/info messages)
  const [notices, setNotices] = useState<RoomNotice[]>([]);

  // Currently active chat room
  const [currentRoom, setCurrentRoom] = useState<string | null>(null);

  // Connect to socket when user logs in, disconnect when logged out
  useEffect(() => {
    if (!user?.token) {
      // If user logs out, close socket connection
      if (socketRef.current) {
        socketRef.current.disconnect();
        socketRef.current = null;
        setIsConnected(false);
        setMySocketId(null);
      }
      return;
    }

    // If already connected, don't reconnect again
    if (socketRef.current?.connected) return;

    // Create socket connection with auth token
    const socket = io(`${API_BASE_URL}/chat`, {
      auth: { token: `Bearer ${user.token}` },
      transports: ["websocket", "polling"],
    });

    socketRef.current = socket;

    /* ---------------- Server Events ---------------- */

    // When connection is successfully established
    socket.on("connect", () => {
      setIsConnected(true);
    });

    // When server confirms who you are
    socket.on(
      "connected",
      (data: { YourId: string; username: string; message: string }) => {
        setMySocketId(data.YourId);
        console.log("[Socket] Connected:", data.message);
      },
    );

    // When socket disconnects
    socket.on("disconnect", () => {
      setIsConnected(false);
      setMySocketId(null);
      console.log("[Socket] Disconnected");
    });

    // If authentication fails
    socket.on("authError", (data: { mesage?: string }) => {
      console.error("[Socket] Auth error:", data.mesage);
      socket.disconnect();
    });

    // General server messages
    socket.on("serverNotice", (message: string) => {
      console.log("[Socket] Server notice:", message);
    });

    //- Room Events

    // When you join a room
    socket.on("roomJoined", (data: { roomId: string; message: string }) => {
      console.log("[Socket] Room joined:", data.message);
    });

    // When you leave a room
    socket.on("roomLeft", (data: { roomId: string }) => {
      console.log("[Socket] Room left:", data.roomId);
    });

    // System messages inside a room
    socket.on("roomNotice", (data: RoomNotice) => {
      setNotices((prev) => [...prev, data]);
    });

    // New chat message received
    socket.on("roomMessage", (data: ChatMessage) => {
      setMessages((prev) => [...prev, data]);
    });

    // Room-related errors
    socket.on("roomError", (data: { message: string }) => {
      console.error("[Socket] Room error:", data.message);
    });

    // Cleanup socket when component unmounts or token changes
    return () => {
      socket.disconnect();
      socketRef.current = null;
      setIsConnected(false);
      setMySocketId(null);
    };
  }, [user?.token]);

  // Actions you can call from UI

  // Join a chat room
  const joinRoom = useCallback((roomId: string) => {
    socketRef.current?.emit("joinRoom", { roomId });
  }, []);

  // Leave a chat room
  const leaveRoom = useCallback((roomId: string) => {
    socketRef.current?.emit("leaveRoom", { roomId });
  }, []);

  // Send a message to a room
  const sendMessage = useCallback((roomId: string, text: string) => {
    socketRef.current?.emit("roomMessage", { roomId, text });
  }, []);

  // Clear all messages and notices (reset chat state)
  const clearMessages = useCallback(() => {
    setMessages([]);
    setNotices([]);
  }, []);

  return (
    <SocketContext.Provider
      value={{
        socket: socketRef.current,
        isConnected,
        mySocketId,
        joinRoom,
        leaveRoom,
        sendMessage,
        messages,
        notices,
        currentRoom,
        setCurrentRoom,
        clearMessages,
      }}
    >
      {children}
    </SocketContext.Provider>
  );
}

// Hook to use socket context

// Custom hook to access socket anywhere in app
export function useSocket() {
  const ctx = useContext(SocketContext);
  if (!ctx) throw new Error("useSocket must be used inside <SocketProvider>");
  return ctx;
}