w-lab-app/src/hooks/chats/useChatSocket.ts

"use client";

import { useCallback, useEffect, useMemo, useRef, useState } from "react";

import {
  fetchChatMessages,
  fetchMyChatSettings,
  updateChatSettings,
} from "@/entities/chats/actions";
import type { IChatMessage } from "@/entities/chats/types";
import { useSingleCheckout } from "@/hooks/payment/useSingleCheckout";
import { useSocketEvent } from "@/hooks/socket/useSocketEvent";
import { useChatStore } from "@/providers/chat-store-provider";
import {
  ESocketStatus,
  useSocketEmit,
  useSocketEntity,
  useSocketStatus,
} from "@/services/socket";
import type {
  ICurrentBalance,
  IRefillModals,
  ISessionStarted,
} from "@/services/socket/events";

const PAGE_LIMIT = 50;

interface UseChatSocketOptions {
  initialMessages?: IChatMessage[];
  initialTotal?: number;
  onNewMessage?: (message: IChatMessage) => void;
}

export const useChatSocket = (
  chatId: string,
  options: UseChatSocketOptions = {}
) => {
  const socket = useSocketEntity();
  const status = useSocketStatus();
  const emit = useSocketEmit();

  const [messages, setMessages] = useState<IChatMessage[]>(
    () => options.initialMessages || []
  );
  const [page, setPage] = useState(1);
  const [totalCount, _setTotalCount] = useState<number | null>(
    options.initialTotal ?? null
  );
  const [isLoadingOlder, setIsLoadingOlder] = useState(false);
  const [balance, setBalance] = useState<ICurrentBalance | null>(null);
  const [session, setSession] = useState<ISessionStarted | null>(null);
  const [isLoadingSelfMessage, setIsLoadingSelfMessage] = useState(false);
  const [isSessionExpired, setIsSessionExpired] = useState(false);
  const [refillModals, setRefillModals] = useState<IRefillModals | null>(null);
  const { suggestions, setSuggestions } = useChatStore(state => state);

  const isLoadingAdvisorMessage = useMemo(() => {
    return (
      messages.length > 0 &&
      (messages[0].role !== "assistant" ||
        (Object.hasOwn(messages[0], "isLast") && !messages[0].isLast))
    );
  }, [messages]);

  const unreadMessagesCount = useMemo(() => {
    return messages.filter(m => !m.isRead && m.role === "assistant").length;
  }, [messages]);

  const joinChat = useCallback(
    () => emit("join_chat", { chatId }),
    [emit, chatId]
  );
  const leaveChat = useCallback(
    () => emit("leave_chat", { chatId }),
    [emit, chatId]
  );

  const send = useCallback(
    (text: string) => {
      const sendingMessage: IChatMessage = {
        id: `sending-message-${Date.now()}`,
        role: "user",
        text,
        createdDate: new Date().toISOString(),
        isRead: false,
        type: "text",
      };
      setMessages(prev => [sendingMessage, ...prev]);
      if (options.onNewMessage) {
        options.onNewMessage(sendingMessage);
      }

      setIsLoadingSelfMessage(true);
      emit("send_message", { chatId, message: text });
    },
    [options, emit, chatId]
  );

  const read = useCallback(
    (ids: string[]) => {
      emit("read_message", { messages: ids });
      // setMessages(prev =>
      //   prev.map(m => (ids.includes(m.id) ? { ...m, isRead: true } : m))
      // );
    },
    [emit]
  );
  const startSession = useCallback(
    () => emit("start_session", { chatId }),
    [emit, chatId]
  );
  const endSession = useCallback(
    () => emit("end_session", { chatId }),
    [emit, chatId]
  );
  const fetchBalance = useCallback(() => {
    emit("fetch_balance", { chatId });
  }, [emit, chatId]);

  const clearRefillModals = useCallback(() => {
    setRefillModals(null);
  }, []);

  // Auto top-up: silent flow (no UI prompt)
  const autoTopUpInProgressRef = useRef(false);
  // Timer for delayed unlock after error
  const autoTopUpUnlockTimerRef = useRef<NodeJS.Timeout | null>(null);

  // Auto top-up: use existing single checkout flow
  const { handleSingleCheckout, isLoading: isAutoTopUpLoading } =
    useSingleCheckout({
      onSuccess: fetchBalance,
      onError: async () => {
        // eslint-disable-next-line no-console
        console.error("Auto top-up payment failed - disabling auto top-up");

        // Disable auto top-up on payment failure
        try {
          // Fetch current settings to preserve topUpAmount
          const settingsResponse = await fetchMyChatSettings();
          if (settingsResponse.data?.settings) {
            await updateChatSettings({
              ...settingsResponse.data.settings,
              autoTopUp: false,
            });
            // eslint-disable-next-line no-console
            console.info("Auto top-up disabled successfully after payment failure");
          }
        } catch (error) {
          // eslint-disable-next-line no-console
          console.error("Failed to disable auto top-up:", error);
        }

        // Release the lock immediately after disabling
        autoTopUpInProgressRef.current = false;
        if (autoTopUpUnlockTimerRef.current) {
          clearTimeout(autoTopUpUnlockTimerRef.current);
          autoTopUpUnlockTimerRef.current = null;
        }
      },
      returnUrl: typeof window !== "undefined" ? window.location.href : "",
    });

  const balancePollId = useRef<NodeJS.Timeout | null>(null);
  // Avoid immediate leave_chat right after join in React 18 StrictMode (dev) double-invoke
  // We debounce leave so that a quick remount cancels it, but real navigation (or chat switch) proceeds.
  const leaveTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const lastChatIdRef = useRef(chatId);

  const startBalancePolling = useCallback(() => {
    if (balancePollId.current) return;
    balancePollId.current = setInterval(fetchBalance, 5_000);
  }, [fetchBalance]);

  const stopBalancePolling = useCallback(() => {
    if (balancePollId.current) {
      clearInterval(balancePollId.current);
      balancePollId.current = null;
    }
  }, []);

  const hasMoreOlderMessages =
    totalCount === null ? false : messages.length < totalCount;

  const loadOlder = useCallback(async () => {
    if (isLoadingOlder || !hasMoreOlderMessages) return;
    setIsLoadingOlder(true);
    try {
      const nextPage = page + 1;
      const { data } = await fetchChatMessages(chatId, {
        limit: PAGE_LIMIT,
        page: nextPage,
      });
      if (!data) return;
      const { messages: msgs } = data;
      setMessages(prev => {
        const ids = new Set(prev.map(m => m.id));
        return [...prev, ...msgs.filter(m => !ids.has(m.id))];
      });
      setPage(nextPage);
    } catch (e) {
      // eslint-disable-next-line no-console
      console.error("Failed to load older messages:", e);
    } finally {
      setIsLoadingOlder(false);
    }
  }, [isLoadingOlder, hasMoreOlderMessages, page, chatId]);

  useSocketEvent("receive_message", data => {
    if (!data?.length) return;

    if (data[0].role === "user") setIsLoadingSelfMessage(false);

    setMessages(prev => {
      const map = new Map<string, IChatMessage>();

      prev
        .filter(m => !m.id.startsWith("sending-message-"))
        .forEach(m => map.set(m.id, m));

      data.forEach(d => map.set(d.id, d));

      return Array.from(map.values()).sort(
        (a, b) =>
          new Date(b.createdDate).getTime() - new Date(a.createdDate).getTime()
      );
    });

    if (options.onNewMessage) {
      options.onNewMessage(data[0]);
    }
  });
  useSocketEvent("current_balance", b => {
    setBalance(b.data);
    // If auto top-up was in-flight, release the lock only after balance became positive
    if (autoTopUpInProgressRef.current && b?.data?.balance > 0) {
      autoTopUpInProgressRef.current = false;
      if (autoTopUpUnlockTimerRef.current) {
        clearTimeout(autoTopUpUnlockTimerRef.current);
        autoTopUpUnlockTimerRef.current = null;
      }
    }
  });
  useSocketEvent("balance_updated", b => {
    setBalance(prev => (prev ? { ...prev, balance: b.data.balance } : null));
    if (autoTopUpInProgressRef.current && b?.data?.balance > 0) {
      autoTopUpInProgressRef.current = false;
      if (autoTopUpUnlockTimerRef.current) {
        clearTimeout(autoTopUpUnlockTimerRef.current);
        autoTopUpUnlockTimerRef.current = null;
      }
    }
  });
  useSocketEvent("session_started", s => setSession(s.data));
  useSocketEvent("session_ended", () => setSession(null));
  useSocketEvent("show_refill_modals", r => setRefillModals(r.data));
  useSocketEvent("auto_topup_request", r => {
    if (!r?.data) return;
    if (isAutoTopUpLoading) return;
    // Prevent concurrent or rapid duplicate attempts
    if (autoTopUpInProgressRef.current) return;
    autoTopUpInProgressRef.current = true;
    // Trigger checkout silently
    handleSingleCheckout(r.data);
  });

  useEffect(() => {
    if (!session?.maxFinishedAt) return;

    const finishAt = new Date(session.maxFinishedAt).getTime();
    const now = Date.now();
    const delay = finishAt - now;

    if (delay <= 0) {
      setIsSessionExpired(true);
      return;
    }

    const id = setTimeout(() => setIsSessionExpired(true), delay);
    return () => clearTimeout(id);
  }, [session?.maxFinishedAt]);

  useEffect(() => {
    if (!socket || status !== ESocketStatus.CONNECTED) return;

    // If a leave was scheduled for the same chat (StrictMode remount), cancel it
    if (leaveTimeoutRef.current && lastChatIdRef.current === chatId) {
      clearTimeout(leaveTimeoutRef.current);
      leaveTimeoutRef.current = null;
    }

    joinChat();
    fetchBalance();
    lastChatIdRef.current = chatId;

    return () => {
      // Debounce leave to avoid leaving room between StrictMode unmount/mount cycle
      // For real chat switch (different chatId), we won't cancel this in the next effect
      leaveTimeoutRef.current = setTimeout(() => {
        leaveChat();
      }, 300);
    };
  }, [socket, status, joinChat, leaveChat, fetchBalance, chatId]);

  // Re-join chat on socket reconnects while staying on the chat page
  useEffect(() => {
    if (!socket) return;
    const handleConnect = () => {
      // Ensure current chat is tracked and we actually reconnect to it
      lastChatIdRef.current = chatId;
      joinChat();
      fetchBalance();
    };
    socket.on("connect", handleConnect);
    return () => {
      socket.off("connect", handleConnect);
    };
  }, [socket, chatId, joinChat, fetchBalance]);

  useEffect(() => {
    setSuggestions(messages[0]?.suggestions);
  }, [messages, setSuggestions]);

  useEffect(() => {
    if (session && status === ESocketStatus.CONNECTED) {
      startBalancePolling();
    } else {
      stopBalancePolling();
    }
    return () => {
      stopBalancePolling();
    };
  }, [session, status, startBalancePolling, stopBalancePolling]);

  useEffect(() => {
    if (!balance) return;

    const hasBalance = balance.balance > 0;
    const hasSession = !!session;

    if (hasBalance && !hasSession) startSession();
    if (!hasBalance && hasSession) endSession();
  }, [balance, session, startSession, endSession]);

  useEffect(() => {
    if (!session) return;

    return () => {
      endSession();
    };
  }, [session, endSession]);

  // Cleanup pending unlock timer on unmount
  useEffect(() => {
    return () => {
      if (autoTopUpUnlockTimerRef.current) {
        clearTimeout(autoTopUpUnlockTimerRef.current);
        autoTopUpUnlockTimerRef.current = null;
      }
    };
  }, []);

  const isAvailableChatting =
    !!balance?.balance && !!session && !isSessionExpired;

  return useMemo(
    () => ({
      chatId,
      messages,
      balance,
      session,
      refillModals,
      suggestions,
      unreadMessagesCount,

      send,
      read,
      startSession,
      endSession,
      clearRefillModals,

      isLoadingSelfMessage,
      isLoadingAdvisorMessage,
      isAvailableChatting,
      isConnected: status === ESocketStatus.CONNECTED,

      loadOlder,
      isLoadingOlder,
      hasMoreOlderMessages,
    }),
    [
      chatId,
      messages,
      balance,
      session,
      refillModals,
      suggestions,
      unreadMessagesCount,
      isLoadingSelfMessage,
      isLoadingAdvisorMessage,
      isAvailableChatting,
      status,
      loadOlder,
      isLoadingOlder,
      hasMoreOlderMessages,
      send,
      read,
      startSession,
      endSession,
      clearRefillModals,
    ]
  );
};