mosaicmap/app/ws-context.tsx

'use client'

import React, { createContext, useContext, useRef, useState, ReactNode, useEffect } from 'react'
import { Map } from 'maplibre-gl';
import { gql, useSubscription } from '@apollo/client';
import { subscribeToConnectionState, reconnectWebSocket, WSConnectionState } from '@/lib/apollo-client';
import * as Sentry from '@sentry/nextjs';

// 定义WSContext的类型
interface WSContextType {
    wsStatus: WsStatus
    connectionState: WSConnectionState
    data: any
    loading: boolean
    error: any
    restart: () => void
    forceReconnect: () => void
    isOnline: boolean
}

enum WsStatus {
    CONNECTING = 'connecting',
    CONNECTED = 'connected',
    DISCONNECTED = 'disconnected',
    RECONNECTING = 'reconnecting'
}

// 创建Context
const WSContext = createContext<WSContextType | undefined>(undefined)

// Provider组件的Props类型
interface MapProviderProps {
    children: ReactNode
}

const SUBSCRIPTION_QUERY = gql`
  subscription {
    statusUpdates {
      id
      message
      status
      timestamp
      newestDt
    }
  }
`


// Provider组件
export function WSProvider({ children }: MapProviderProps) {
    const [wsStatus, setWsStatus] = useState<WsStatus>(WsStatus.DISCONNECTED)
    const [connectionState, setConnectionState] = useState<WSConnectionState>({
        status: 'disconnected',
        reconnectAttempts: 0
    });
    const [isOnline, setIsOnline] = useState(typeof window !== 'undefined' ? navigator.onLine : true);

    const { data, loading, error, restart } = useSubscription(SUBSCRIPTION_QUERY, {
        errorPolicy: 'all',
        onError: (error) => {
            Sentry.captureException(error);
        }
    });

    // 监听网络状态变化
    useEffect(() => {
        if (typeof window === 'undefined') return;

        const handleOnline = () => {
            setIsOnline(true);
            if (wsStatus === WsStatus.DISCONNECTED) {
                restart();
            }
        };

        const handleOffline = () => {
            setIsOnline(false);
            setWsStatus(WsStatus.DISCONNECTED);
        };

        window.addEventListener('online', handleOnline);
        window.addEventListener('offline', handleOffline);

        return () => {
            window.removeEventListener('online', handleOnline);
            window.removeEventListener('offline', handleOffline);
        };
    }, [wsStatus, restart]);

    // 监听 WebSocket 连接状态
    useEffect(() => {
        const unsubscribe = subscribeToConnectionState((state) => {
            setConnectionState(state);

            switch (state.status) {
                case 'connecting':
                    setWsStatus(WsStatus.CONNECTING);
                    break;
                case 'connected':
                    setWsStatus(WsStatus.CONNECTED);
                    break;
                case 'disconnected':
                    setWsStatus(WsStatus.DISCONNECTED);
                    break;
                case 'reconnecting':
                    setWsStatus(WsStatus.RECONNECTING);
                    break;
            }
        });

        return unsubscribe;
    }, []);

    // 监听订阅状态变化
    useEffect(() => {
        if (loading) {
            // 只在第一次加载时设置为 connecting，避免覆盖 WebSocket 的实际状态
            if (connectionState.status === 'disconnected') {
                setWsStatus(WsStatus.CONNECTING);
            }
        } else if (error && !loading) {
            console.error('WebSocket subscription error:', error);
            if (isOnline && connectionState.status !== 'reconnecting') {
                setTimeout(() => {
                    restart();
                }, 5000);
            }
        } else if (!loading && !error && data) {
            // 确保在有数据时状态是 connected
            if (wsStatus !== WsStatus.CONNECTED) {
                setWsStatus(WsStatus.CONNECTED);
            }
        }
    }, [loading, error, data, isOnline, restart, wsStatus, connectionState.status]);

    const forceReconnect = () => {
        Sentry.startSpan({
            op: 'websocket.manual-reconnect',
            name: 'Manual WebSocket Reconnection'
        }, () => {
            reconnectWebSocket();
            restart();
        });
    };

    const value: WSContextType = {
        wsStatus,
        connectionState,
        data,
        loading,
        error,
        restart,
        forceReconnect,
        isOnline
    }

    return (
        <WSContext.Provider value={value}>
            {children}
        </WSContext.Provider>
    )
}

// 自定义Hook用于使用MapContext
export function useWS() {
    const context = useContext(WSContext)

    if (context === undefined) {
        throw new Error('useWS must be used within a WSProvider')
    }

    return context
}