telegram-web-client/frontend/src/hooks/useWebSocket.ts

import { useEffect, useRef, useCallback } from 'react';
import { useChatStore } from '../stores/chatStore';
import type { Message } from '../types';

const WS_URL = import.meta.env.DEV
  ? 'ws://localhost:8000/ws'
  : `wss://${window.location.host}/ws`;

export function useWebSocket() {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | undefined>(undefined);

  const {
    setSessionId,
    setConnected,
    setAuthStep,
    setPhone,
    setPhoneCodeHash,
    setUser,
    addMessage,
    setMessages,
    setLoading,
    setError,
  } = useChatStore();

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

    const ws = new WebSocket(WS_URL);
    wsRef.current = ws;

    ws.onopen = () => {
      console.log('WebSocket connected');
      setConnected(true);
      setError(null);
    };

    ws.onclose = () => {
      console.log('WebSocket disconnected');
      setConnected(false);

      // Reconnect after 3 seconds
      reconnectTimeoutRef.current = window.setTimeout(() => {
        connect();
      }, 3000);
    };

    ws.onerror = (error) => {
      console.error('WebSocket error:', error);
      setError('Connection error');
    };

    ws.onmessage = (event) => {
      try {
        const message = JSON.parse(event.data);
        handleMessage(message);
      } catch (e) {
        console.error('Failed to parse message:', e);
      }
    };
  }, []);

  const handleMessage = useCallback(
    (message: { type: string; data: Record<string, unknown> }) => {
      const { type, data } = message;

      switch (type) {
        case 'connected':
          setSessionId(data.session_id as string);
          // Check auth status
          send({ type: 'check_auth', data: {} });
          break;

        case 'auth_status':
          if (data.authenticated) {
            setUser({
              id: data.user_id as number,
              username: data.username as string,
              firstName: data.first_name as string,
            });
            setAuthStep('authenticated');
            // Get chat history
            send({ type: 'get_history', data: {} });
          } else {
            setAuthStep('phone');
          }
          setLoading(false);
          break;

        case 'code_sent':
          setPhone(data.phone as string);
          setPhoneCodeHash(data.phone_code_hash as string);
          setAuthStep('code');
          setLoading(false);
          break;

        case 'need_password':
          setAuthStep('password');
          setLoading(false);
          break;

        case 'auth_success':
          setUser({
            id: data.user_id as number,
            username: data.username as string,
            firstName: data.first_name as string,
          });
          setAuthStep('authenticated');
          setLoading(false);
          // Get chat history
          send({ type: 'get_history', data: {} });
          break;

        case 'auth_error':
          setError(data.message as string);
          setLoading(false);
          break;

        case 'history':
          setMessages((data.messages as Message[]) || []);
          setLoading(false);
          break;

        case 'message_sent':
          addMessage(data as unknown as Message);
          break;

        case 'new_message':
          addMessage(data as unknown as Message);
          break;

        case 'error':
          setError(data.message as string);
          setLoading(false);
          break;

        default:
          console.log('Unknown message type:', type);
      }
    },
    []
  );

  const send = useCallback((message: { type: string; data: Record<string, unknown> }) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    }
  }, []);

  const sendCode = useCallback((phone: string) => {
    setLoading(true);
    send({ type: 'send_code', data: { phone } });
  }, [send]);

  const verifyCode = useCallback(
    (code: string) => {
      const { phone, phoneCodeHash } = useChatStore.getState();
      setLoading(true);
      send({
        type: 'verify_code',
        data: { phone, code, phone_code_hash: phoneCodeHash },
      });
    },
    [send]
  );

  const verifyPassword = useCallback(
    (password: string) => {
      setLoading(true);
      send({ type: 'verify_password', data: { password } });
    },
    [send]
  );

  const sendMessage = useCallback(
    (text: string) => {
      send({ type: 'send_message', data: { text } });
    },
    [send]
  );

  useEffect(() => {
    connect();

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

  return {
    sendCode,
    verifyCode,
    verifyPassword,
    sendMessage,
  };
}