'use client';

import { useCallback, useEffect, useRef, useState } from 'react';
import type { Message } from '@/api/chat';

type WsEvent =
  | { type: 'connection_established' }
  | { type: 'chat_message'; message: Message }
  | { type: 'error'; error?: string }
  | { type: string; [k: string]: unknown };

export function useChatWebSocket(options: {
  chatUuid: string | null;
  enabled?: boolean;
  onMessage?: (m: Message) => void;
}) {
  const { chatUuid, enabled = true, onMessage } = options;
  const [isConnected, setIsConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;

  const disconnect = useCallback(() => {
    if (wsRef.current) {
      wsRef.current.close();
      wsRef.current = null;
    }
    setIsConnected(false);
  }, []);

  const connect = useCallback(() => {
    if (!enabled || !chatUuid) return;

    const token = typeof window !== 'undefined' ? localStorage.getItem('access_token') : null;
    if (!token) return;

    let apiUrl = process.env.NEXT_PUBLIC_API_URL || window.location.origin;
    apiUrl = apiUrl.replace(/\/api\/?$/, '').replace(/\/api\//, '/');
    apiUrl = apiUrl.replace(/\/$/, '');

    const wsProtocol = apiUrl.startsWith('https') ? 'wss:' : 'ws:';
    const wsHost = apiUrl.replace(/^https?:\/\//, '');
    const wsUrl = `${wsProtocol}//${wsHost}/ws/chat/${chatUuid}/?token=${token}`;

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

    ws.onopen = () => setIsConnected(true);
    ws.onclose = () => setIsConnected(false);
    ws.onerror = () => setIsConnected(false);
    ws.onmessage = (ev) => {
      try {
        const data = JSON.parse(ev.data) as WsEvent;
        if (data.type === 'chat_message' && (data as any).message) {
          onMessageRef.current?.((data as any).message as Message);
        }
      } catch {
        // ignore
      }
    };
  }, [chatUuid, enabled]);

  useEffect(() => {
    disconnect();
    connect();
    return () => disconnect();
  }, [connect, disconnect]);

  return { isConnected };
}