import { wsUrl } from '@/utils';
import { useEffect } from 'react';

export type Payload = {
  chatId: string;
  content: string;
  action: 'append' | 'done';
} | {
  chatId: string;
  title: string;
  action:'title_updated';
} | {
  chatId: string;
  url: string;
  action:'image';
}

export type OnMessage = (payload: Payload) => void;

type WebSocketEvent = {
  chatId: string;
  payload: Payload;
};

// TODO: Make this into class so we can mock for test.
const wsSingleton = new Map<string, WebSocket>();
const listeners = new Map<string, Set<OnMessage>>();
const pendingMessages = new Map<string, Array<{ content: string; resolve: () => void }>>();

function getOrCreateWebSocket(chatId: string): WebSocket {
  let ws = wsSingleton.get(chatId);
  
  if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
    return ws;
  }

  ws = new WebSocket(wsUrl(`/chats/${chatId}/ws`));

  ws.onopen = () => {
    console.log(`[WS] Connected for chat ${chatId}`);
    const queue = pendingMessages.get(chatId);
    if (queue) {
      queue.forEach(({ content }) => {
        ws!.send(JSON.stringify({ role: 'user', content }));
      });
      pendingMessages.delete(chatId);
    }
  };

  // ADD IT HERE:
  ws.onmessage = (event) => {
    const payload = JSON.parse(event.data);
    console.log('[WS] Received message:', payload);
    
    wsEventTarget.dispatchEvent(
      new CustomEvent('message', {
        detail: { chatId, payload },
      }),
    );
  };

  ws.onerror = (error) => {
    console.error(`[WS] Error for chat ${chatId}:`, error);
  };

  ws.onclose = () => {
    console.log(`[WS] Closed for chat ${chatId}`);
    wsSingleton.delete(chatId);
  };

  wsSingleton.set(chatId, ws);
  return ws;
}

// TODO: This could be not done rather cancel?
export function broadcastDone(chatId: string) {
  const set = listeners.get(chatId);
  // TODO: This does not update the histroy to be canceled on so there is decrepency but it should be fine for now.
  if (set) {
    set.forEach((cb) => cb({ chatId, content: '', action: 'done' }));
  }
}

export function useChatWebSocket(
  chatId: string | null,
  onMessage: OnMessage,
  onLoadingChange: (loading: boolean) => void,
) {

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

    // Only register listener to given chatId as chatId websocket has not been made yet. 
    let set = listeners.get(chatId);
    if (!set) {
      set = new Set();
      listeners.set(chatId, set);
    }
    set.add(onMessage);

    return () => {
      if (!chatId) return;
      const set = listeners.get(chatId);
      if (set) {
        set.delete(onMessage);
        if (set.size === 0) {
          listeners.delete(chatId);
          const ws = wsSingleton.get(chatId);
          if (ws && (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING)) {
            ws.close();
          }
        }
      }
    };
  }, [chatId, onMessage]);

  const sendMessage = (content: string) => {
    if (!chatId) return;

    const ws = getOrCreateWebSocket(chatId);

    onLoadingChange(true);

    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ role: 'user', content }));
    } else if (ws.readyState === WebSocket.CONNECTING) {
      // Queue message until open
      let queue = pendingMessages.get(chatId);
      if (!queue) {
        queue = [];
        pendingMessages.set(chatId, queue);
      }
      console.log('Queue: ', queue);
      if (queue.length > 1) return;
      queue.push({ content, resolve: () => {} });
    } else {
      console.warn('[WS] WebSocket not in usable state', ws.readyState);
      onLoadingChange(false);
    }
  };

  return { sendMessage };
}

const wsEventTarget = new EventTarget();

export function subscribeToChat(chatId: string, callback: (payload: Payload) => void) {
  const handler = (event: Event) => {
    const { chatId: eventChatId, payload } = (event as CustomEvent<WebSocketEvent>).detail;
    if (eventChatId === chatId) {
      callback(payload);
    }
  };
  
  wsEventTarget.addEventListener('message', handler);
  
  return () => wsEventTarget.removeEventListener('message', handler);
}

export function sendMessageChatId(content: string, chatId: string): void {
  if (!content.trim() || !chatId) return;
  
  const ws = getOrCreateWebSocket(chatId);
  
  if (ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify({ role: 'user', content }));
  } else if (ws.readyState === WebSocket.CONNECTING) {
    let queue = pendingMessages.get(chatId);
    if (!queue) {
      queue = [];
      pendingMessages.set(chatId, queue);
    }
    if (!queue.some(m => m.content === content)) {
      queue.push({ content, resolve: () => {} });
    }
  }
}