uchill/front_material/hooks/useNotifications.ts

'use client';

import { useCallback, useEffect, useRef, useState } from 'react';
import {
  getNotifications,
  getUnread,
  markAllAsRead as apiMarkAllAsRead,
  markAsRead,
  type Notification,
} from '@/api/notifications';

type WsMessage =
  | { type: 'connection_established' }
  | { type: 'unread_count'; count: number }
  | { type: 'new_notification'; notification: Notification; unread_count?: number }
  | { type: 'notification_read'; notification_id: number }
  | { type: 'all_notifications_read'; count: number }
  | { type: 'nav_badges_updated' }
  | { type: 'error'; message?: string }
  | { type: string; [k: string]: unknown };

function dedupeById(items: Notification[]): Notification[] {
  const seen = new Set<number>();
  return items.filter((n) => {
    if (seen.has(n.id)) return false;
    seen.add(n.id);
    return true;
  });
}

export function useNotifications(options?: { wsEnabled?: boolean; onNavBadgesUpdate?: () => void }) {
  const wsEnabled = options?.wsEnabled ?? true;
  const onNavBadgesUpdateRef = useRef(options?.onNavBadgesUpdate);
  onNavBadgesUpdateRef.current = options?.onNavBadgesUpdate;
  const PAGE_SIZE = 20;
  const [list, setList] = useState<Notification[]>([]);
  const [unreadCount, setUnreadCount] = useState(0);
  const [loading, setLoading] = useState(false);
  const [loadingMore, setLoadingMore] = useState(false);
  const [hasMore, setHasMore] = useState(true);
  const [nextPage, setNextPage] = useState(2);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  const fetchList = useCallback(async () => {
    setLoading(true);
    setHasMore(true);
    setNextPage(2);
    try {
      const res = await getNotifications({ page: 1, page_size: PAGE_SIZE });
      const results = res.results ?? [];
      setList(dedupeById(results));
      setHasMore(Boolean(res.next) || results.length >= PAGE_SIZE);
    } catch {
      setList([]);
      setHasMore(false);
    } finally {
      setLoading(false);
    }
  }, []);

  const fetchMore = useCallback(async () => {
    if (loadingMore || !hasMore) return;
    setLoadingMore(true);
    try {
      const res = await getNotifications({ page: nextPage, page_size: PAGE_SIZE });
      const results = res.results ?? [];
      setList((prev) => {
        const existingIds = new Set(prev.map((n) => n.id));
        const newItems = results.filter((n) => !existingIds.has(n.id));
        return newItems.length ? [...prev, ...newItems] : prev;
      });
      setHasMore(Boolean(res.next) || results.length >= PAGE_SIZE);
      setNextPage((p) => p + 1);
    } catch {
      setHasMore(false);
    } finally {
      setLoadingMore(false);
    }
  }, [loadingMore, hasMore, nextPage]);

  const fetchUnreadCount = useCallback(async () => {
    try {
      const { count } = await getUnread();
      setUnreadCount(count);
    } catch {
      setUnreadCount(0);
    }
  }, []);

  const markOneAsRead = useCallback(async (id: number) => {
    try {
      await markAsRead(id);
      setList((prev) =>
        prev.map((n) => (n.id === id ? { ...n, is_read: true } : n))
      );
      setUnreadCount((c) => Math.max(0, c - 1));
    } catch {
      // ignore
    }
  }, []);

  const markAllAsRead = useCallback(async () => {
    try {
      await apiMarkAllAsRead();
      setUnreadCount(0);
      setList((prev) => prev.map((n) => ({ ...n, is_read: true })));
    } catch {
      // ignore
    }
  }, []);

  useEffect(() => {
    fetchUnreadCount();
  }, [fetchUnreadCount]);

  useEffect(() => {
    if (!wsEnabled || typeof window === 'undefined') return;

    const token = localStorage.getItem('access_token');
    if (!token) return;

    let baseUrl = process.env.NEXT_PUBLIC_API_URL || window.location.origin;
    baseUrl = baseUrl.replace(/\/api\/?$/, '').replace(/\/$/, '');
    const wsProtocol = baseUrl.startsWith('https') ? 'wss:' : 'ws:';
    const wsHost = baseUrl.replace(/^https?:\/\//, '');
    const wsUrl = `${wsProtocol}//${wsHost}/ws/notifications/?token=${token}`;

    const connect = () => {
      if (wsRef.current?.readyState === WebSocket.OPEN) return;
      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onmessage = (ev) => {
        try {
          const data = JSON.parse(ev.data) as WsMessage;
          switch (data.type) {
            case 'unread_count':
              setUnreadCount((data as { count: number }).count);
              break;
            case 'new_notification': {
              const payload = data as { notification: Notification; unread_count?: number };
              if (payload.unread_count !== undefined) setUnreadCount(payload.unread_count);
              else setUnreadCount((c) => c + 1);
              setList((prev) => {
                const has = prev.some((n) => n.id === payload.notification.id);
                if (has) return prev;
                return [payload.notification, ...prev];
              });
              break;
            }
            case 'notification_read': {
              const id = (data as { notification_id: number }).notification_id;
              setList((prev) =>
                prev.map((n) => (n.id === id ? { ...n, is_read: true } : n))
              );
              setUnreadCount((c) => Math.max(0, c - 1));
              break;
            }
            case 'all_notifications_read':
              setUnreadCount(0);
              setList((prev) => prev.map((n) => ({ ...n, is_read: true })));
              break;
            case 'nav_badges_updated':
              onNavBadgesUpdateRef.current?.();
              break;
            default:
              break;
          }
        } catch {
          // ignore
        }
      };

      ws.onclose = () => {
        wsRef.current = null;
        reconnectTimeoutRef.current = setTimeout(connect, 3000);
      };

      ws.onerror = () => {};
    };

    connect();
    return () => {
      if (reconnectTimeoutRef.current) clearTimeout(reconnectTimeoutRef.current);
      wsRef.current?.close();
      wsRef.current = null;
    };
  }, [wsEnabled]);

  return {
    list,
    unreadCount,
    loading,
    loadingMore,
    hasMore,
    fetchList,
    fetchMore,
    fetchUnreadCount,
    markOneAsRead,
    markAllAsRead,
  };
}