client/src/hooks/useWebSocket.ts

import { useEffect, useRef, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';
import { useAuthStore } from '../store/authStore';
import { useCollaborationStore, Collaborator } from '../store/collaborationStore';
import { useEditorStore } from '../store/editorStore';
import { useWsStore } from '../store/wsStore';
import {
  BLOCK_STATE_SYNC,
  BLOCK_STATE_SYNCED,
  REQUEST_BLOCK_STATE,
  COLLABORATOR_POINTER_MOVE,
  COLLABORATOR_BLOCK_DRAG_START,
  COLLABORATOR_BLOCK_DRAG_MOVE,
  COLLABORATOR_BLOCK_DRAG_END,
  COLLABORATOR_BLOCK_CONNECT,
} from '@shared/wsEvents';

let globalSocket: Socket | null = null;

// ─── Generic hook to subscribe to a WS event ───
export function useWsEvent(event: string, handler: (...args: any[]) => void) {
  const socket = useWsStore((s) => s.socket);
  const handlerRef = useRef(handler);
  handlerRef.current = handler;

  useEffect(() => {
    if (!socket) return;
    const wrapper = (...args: any[]) => handlerRef.current(...args);
    socket.on(event, wrapper);
    return () => {
      socket.off(event, wrapper);
    };
  }, [socket, event]);
}

// ─── Socket lifecycle: connect on token, reconnect with stored projectId ───
export function useWebSocket() {
  const token = useAuthStore((s) => s.token);
  const { setConnected: setCollabConnected } = useCollaborationStore();
  const { setConnected, setSocket, pushPending, shiftPending } = useWsStore();

  useEffect(() => {
    if (!token) {
      if (globalSocket) {
        globalSocket.disconnect();
        globalSocket = null;
        setSocket(null);
        setConnected(false);
        setCollabConnected(false);
      }
      return;
    }

    if (globalSocket?.connected) return;

    const s = io({
      auth: { token },
      transports: ['websocket', 'polling'],
    });

    globalSocket = s;
    setSocket(s);

    s.on('connect', () => {
      setConnected(true);
      setCollabConnected(true);

      // Re-join project if we have one stored
      const pid = useWsStore.getState().currentProjectId;
      if (pid) {
        s.emit('join_project', { projectId: pid });
      }

      // Replay pending queue
      let pending = shiftPending();
      while (pending) {
        s.emit(pending.action, pending.data);
        pending = shiftPending();
      }
    });

    s.on('disconnect', () => {
      setConnected(false);
      setCollabConnected(false);
    });

    s.on('collaborator_list', ({ collaborators }: { collaborators: Collaborator[] }) => {
      useCollaborationStore.getState().setCollaborators(collaborators);
    });

    s.on('collaborator_joined', (collaborator: Collaborator) => {
      useCollaborationStore.getState().addCollaborator(collaborator);
    });

    s.on('collaborator_left', ({ userId }: { userId: string }) => {
      useCollaborationStore.getState().removeCollaborator(userId);
    });

    s.on('collaborator_active_file', ({ userId, fileId }: { userId: string; fileId: string | null }) => {
      useCollaborationStore.getState().updateActiveFile(userId, fileId);
    });

    s.on('collaborator_cursor', ({ userId, fileId, cursor }: { userId: string; fileId: string; cursor: { line: number; ch: number } | null }) => {
      useCollaborationStore.getState().updateCursor(userId, fileId, cursor);
    });

    s.on('collaborator_block_selected', ({ userId, blockId }: { userId: string; blockId: string | null }) => {
      useCollaborationStore.getState().updateBlockSelection(userId, blockId);
    });

    s.on('collaborator_element_hover', ({ userId, elementId }: { userId: string; elementId: string | null }) => {
      useCollaborationStore.getState().updateElementHover(userId, elementId);
    });

    s.on('block_state_synced', (data: { fileId: string; xml: string; updatedBy: string; initialBlocksXml?: Record<string, { xml: string; updatedBy: string; updatedAt: number }> }) => {
      // Handle initial block state for new joiners (sent on join_project).
      // server sends initialBlocksXml as { fileId: { xml, updatedBy, updatedAt } }
      if (data.initialBlocksXml) {
        console.log('[WS] Received initialBlocksXml with', Object.keys(data.initialBlocksXml).length, 'files');
        const store = useEditorStore.getState();
        for (const [fid, entry] of Object.entries(data.initialBlocksXml)) {
          console.log('[WS]   initialBlocksXml file:', fid, 'xmlLength:', entry.xml?.length, 'updatedBy:', entry.updatedBy);
          if (entry.xml) store.setBlocksXml(fid, entry.xml);
        }
        if (store.activeFileId && data.initialBlocksXml[store.activeFileId]?.xml) {
          console.log('[WS]   setting visual elements for active file:', store.activeFileId);
          store.setVisualElements(store.elementRegistry[store.activeFileId] || []);
        }
        return;
      }
      // Real-time sync update
      // Note: socket.to() already excludes the sender, so no self-update check needed.
      // The updatedBy filter was removed because shared auth tokens can cause
      // different users to appear as the same user ID, incorrectly discarding syncs.
      console.log('[WS] block_state_synced: updating blocksXml for file:', data.fileId, 'xml length:', data.xml?.length);
      useEditorStore.getState().setBlocksXml(data.fileId, data.xml);
      const store = useEditorStore.getState();
      if (store.activeFileId === data.fileId) {
        console.log('[WS] block_state_synced: file is active, setting visual elements');
        store.setVisualElements(store.elementRegistry[data.fileId] || []);
      } else {
        console.log('[WS] block_state_synced: file is NOT active (activeFileId:', store.activeFileId, ')');
      }
    });

    s.on('collaborator_pointer_move', ({ userId, username, x, y }: { userId: string; username: string; x: number; y: number }) => {
      useCollaborationStore.getState().setRemotePointer(userId, { x, y, username });
    });

    s.on('collaborator_block_drag_start', ({ userId, blockId, x, y }: { userId: string; blockId: string; x: number; y: number }) => {
      useCollaborationStore.getState().setRemoteBlockDrag(userId, { blockId, x, y });
    });

    s.on('collaborator_block_drag_move', ({ userId, blockId, x, y }: { userId: string; blockId: string; x: number; y: number }) => {
      useCollaborationStore.getState().updateRemoteBlockDrag(userId, { x, y });
    });

    s.on('collaborator_block_drag_end', ({ userId }: { userId: string }) => {
      useCollaborationStore.getState().clearRemoteBlockDrag(userId);
    });

    s.on('collaborator_block_connect', ({ userId, childId, parentId, inputName }: { userId?: string; childId: string; parentId: string; inputName?: string }) => {
      // Clear any remaining drag state for this user (connect signals drag end + snap).
      // Prevents stale position updates from disconnecting the block after it was connected.
      if (userId) useCollaborationStore.getState().clearRemoteBlockDrag(userId);
      useCollaborationStore.getState().pushPendingBlockConnect({ childId, parentId, inputName });
    });

    s.on('collaborator_block_change', ({ blockId, name, value }: { blockId: string; name: string; value: any }) => {
      useCollaborationStore.getState().pushPendingBlockChange({ blockId, name, value });
    });

    s.on('collaborator_block_create_event', ({ blockId, fileId, xml }: { blockId: string; fileId: string; xml: string }) => {
      const store = useEditorStore.getState();
      console.log('[WS] collaborator_block_create_event received:', { blockId, fileId, activeFileId: store.activeFileId, xmlLength: xml?.length });
      // Only apply to the currently active file; otherwise the block would appear
      // on the wrong file's workspace. The full XML sync (block_state_synced) will
      // handle it when the collaborator switches to that file.
      if (store.activeFileId === fileId) {
        console.log('[WS] pushing pending block create:', blockId);
        useCollaborationStore.getState().pushPendingBlockCreate({ blockId, xml });
      } else {
        console.log('[WS] skipping block create - fileId mismatch:', { eventFileId: fileId, activeFileId: store.activeFileId });
      }
    });

    s.on('file_updated', ({ fileId, content }: { fileId: string; content: string }) => {
      useCollaborationStore.getState().updateFile(fileId, content);
    });

    s.on('file_added', ({ file, parentId }: { file: any; parentId?: string }) => {
      useEditorStore.getState().addFile(file, parentId);
    });

    s.on('file_renamed', ({ fileId, name }: { fileId: string; name: string }) => {
      useEditorStore.getState().updateFile(fileId, { name });
    });

    s.on('file_deleted', ({ fileId }: { fileId: string }) => {
      useEditorStore.getState().removeFile(fileId);
    });

    s.on('visual_element_added', ({ element, parentId }: { element: any; parentId?: string }) => {
      useEditorStore.getState().addVisualElement(element, parentId);
    });

    s.on('visual_element_updated', ({ elementId, updates }: { elementId: string; updates: any }) => {
      useEditorStore.getState().updateVisualElement(elementId, updates);
    });

    s.on('visual_element_removed', ({ elementId }: { elementId: string }) => {
      useEditorStore.getState().removeVisualElement(elementId);
    });

    s.on('visual_elements_reordered', ({ elements }: { elements: any[] }) => {
      const store = useEditorStore.getState();
      store.setVisualElements(elements);
      if (store.activeFileId) {
        store.syncVisualElementsToRegistry(store.activeFileId);
      }
    });

    s.on('project_saved', ({ updatedAt, savedBy }: { updatedAt: number; savedBy: string }) => {
      useCollaborationStore.getState().setLastSaved(updatedAt);
      useCollaborationStore.getState().setSavedBy(savedBy);
    });

    return () => {
      const pid = useWsStore.getState().currentProjectId;
      if (pid) {
        s.emit('leave_project', { projectId: pid });
      }
      s.disconnect();
      if (globalSocket === s) globalSocket = null;
      setSocket(null);
      setConnected(false);
      setCollabConnected(false);
    };
  }, [token, setCollabConnected, setConnected, setSocket, pushPending, shiftPending]);
}

function usePageVisibility() {
  const visibleRef = useRef(!document.hidden);
  useEffect(() => {
    const handler = () => { visibleRef.current = !document.hidden; };
    document.addEventListener('visibilitychange', handler);
    return () => document.removeEventListener('visibilitychange', handler);
  }, []);
  return visibleRef;
}

// ─── Project-scoped socket: join/leave room + emit functions ───
export function useProjectSocket(projectId: string | undefined) {
  const setCurrentProjectId = useWsStore((s) => s.setCurrentProjectId);
  const pageVisible = usePageVisibility();

  // Join/leave project room when projectId changes
  useEffect(() => {
    if (!projectId) {
      setCurrentProjectId(null);
      return;
    }

    setCurrentProjectId(projectId);
    const s = globalSocket;
    if (s?.connected) {
      s.emit('join_project', { projectId });
    }

    return () => {
      setCurrentProjectId(null);
      if (s?.connected) {
        s.emit('leave_project', { projectId });
      }
    };
  }, [projectId, setCurrentProjectId]);

  // Leave/rejoin project room when tab visibility changes
  useEffect(() => {
    if (!projectId) return;
    const handler = () => {
      const s = globalSocket;
      if (!s?.connected) return;
      if (document.hidden) {
        s.emit('leave_project', { projectId });
      } else {
        s.emit('join_project', { projectId });
      }
    };
    document.addEventListener('visibilitychange', handler);
    return () => document.removeEventListener('visibilitychange', handler);
  }, [projectId]);

  // ─── Generic emit for project actions ───
  const emitProjectAction = useCallback(<T>(action: string, data: any): Promise<T> => {
    return new Promise((resolve, reject) => {
      const s = globalSocket;
      if (!s?.connected || !projectId) {
        reject(new Error('Not connected'));
        return;
      }
      s.emit(action, { projectId, ...data }, (response: any) => {
        if (response?.error) reject(new Error(response.error));
        else resolve(response as T);
      });
    });
  }, [projectId]);

  const emitFileChanged = useCallback((fileId: string, content: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('file_changed', { projectId, fileId, content });
    }
  }, [projectId]);

  const emitActiveFileChanged = useCallback((fileId: string | null) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('active_file_changed', { projectId, fileId });
    }
  }, [projectId]);

  const emitCursorMove = useCallback((fileId: string, cursor: { line: number; ch: number }) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('cursor_move', { projectId, fileId, cursor });
    }
  }, [projectId]);

  const emitBlockSelection = useCallback((blockId: string | null) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_selection_changed', { projectId, blockId });
    }
  }, [projectId]);

  const emitElementHover = useCallback((elementId: string | null) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('visual_element_hover', { projectId, elementId });
    }
  }, [projectId]);

  const emitFileAdded = useCallback((file: any, parentId?: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('file_added', { projectId, file, parentId });
    }
  }, [projectId]);

  const emitFileRenamed = useCallback((fileId: string, name: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('file_renamed', { projectId, fileId, name });
    }
  }, [projectId]);

  const emitFileDeleted = useCallback((fileId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('file_deleted', { projectId, fileId });
    }
  }, [projectId]);

  const emitVisualElementAdded = useCallback((element: any, parentId?: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('visual_element_added', { projectId, element, parentId });
    }
  }, [projectId]);

  const emitVisualElementUpdated = useCallback((elementId: string, updates: any) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('visual_element_updated', { projectId, elementId, updates });
    }
  }, [projectId]);

  const emitVisualElementRemoved = useCallback((elementId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('visual_element_removed', { projectId, elementId });
    }
  }, [projectId]);

  const emitVisualElementsReordered = useCallback((elements: any[]) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('visual_elements_reordered', { projectId, elements });
    }
  }, [projectId]);

  const emitSave = useCallback((data: any): Promise<any> => {
    return new Promise((resolve, reject) => {
      const s = globalSocket;
      if (!s?.connected || !projectId) {
        reject(new Error('Not connected'));
        return;
      }
      s.emit('save_project', { projectId, data }, (response: any) => {
        if (response?.error) reject(new Error(response.error));
        else resolve(response);
      });
    });
  }, [projectId]);

  const lastPointerEmit = useRef(0);
  const emitPointerMove = useCallback((x: number, y: number) => {
    const now = Date.now();
    if (now - lastPointerEmit.current < 100) return;
    lastPointerEmit.current = now;
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('pointer_move', { projectId, x, y });
    }
  }, [projectId]);

  const emitBlockDragStart = useCallback((blockId: string, x: number, y: number) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_drag_start', { projectId, blockId, x, y });
    }
  }, [projectId]);

  const lastBlockDragEmit = useRef(0);
  const emitBlockDragMove = useCallback((blockId: string, x: number, y: number) => {
    const now = Date.now();
    if (now - lastBlockDragEmit.current < 50) return;
    lastBlockDragEmit.current = now;
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_drag_move', { projectId, blockId, x, y });
    }
  }, [projectId]);

  const emitBlockDragEnd = useCallback((blockId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_drag_end', { projectId, blockId });
    }
  }, [projectId]);

  const emitBlockStateSync = useCallback((fileId: string, xml: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      console.log('[emitBlockStateSync] sending:', { projectId, fileId, xmlLength: xml?.length, socketConnected: s?.connected });
      s.emit('block_state_sync', { projectId, fileId, xml });
    } else {
      console.warn('[emitBlockStateSync] FAILED: socket not connected or no projectId', { connected: s?.connected, projectId });
    }
  }, [projectId]);

  const emitBlockConnect = useCallback((childId: string, parentId: string, inputName?: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_connect', { projectId, childId, parentId, inputName });
    }
  }, [projectId]);

  const emitBlockChange = useCallback((blockId: string, name: string, value: any, element: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_change', { projectId, blockId, name, value, element });
    }
  }, [projectId]);

  const emitBlockCreateEvent = useCallback((blockId: string, fileId: string, xml: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('block_create_event', { projectId, fileId, blockId, xml });
    }
  }, [projectId]);

  const emitTextFieldFocus = useCallback((fieldId: string, fileId: string, cursorPosition: number) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('text_field_focus', { projectId, fieldId, fileId, cursorPosition });
    }
  }, [projectId]);

  const lastTextFieldCursorEmit = useRef(0);
  const emitTextFieldCursor = useCallback((fieldId: string, cursorPosition: number) => {
    const now = Date.now();
    if (now - lastTextFieldCursorEmit.current < 80) return;
    lastTextFieldCursorEmit.current = now;
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('text_field_cursor', { projectId, fieldId, cursorPosition });
    }
  }, [projectId]);

  const emitTextFieldBlur = useCallback((fieldId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('text_field_blur', { projectId, fieldId });
    }
  }, [projectId]);

  const emitCssRuleAdded = useCallback((rule: any) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_rule_added', { projectId, rule });
    }
  }, [projectId]);

  const emitCssRuleUpdated = useCallback((ruleId: string, updates: any) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_rule_updated', { projectId, ruleId, updates });
    }
  }, [projectId]);

  const emitCssRuleRemoved = useCallback((ruleId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_rule_removed', { projectId, ruleId });
    }
  }, [projectId]);

  const emitCssRuleReordered = useCallback((ruleIds: string[]) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_rule_reordered', { projectId, ruleIds });
    }
  }, [projectId]);

  const emitCssPropAdded = useCallback((ruleId: string, property: any) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_prop_added', { projectId, ruleId, property });
    }
  }, [projectId]);

  const emitCssPropUpdated = useCallback((ruleId: string, propertyId: string, updates: any) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_prop_updated', { projectId, ruleId, propertyId, updates });
    }
  }, [projectId]);

  const emitCssPropRemoved = useCallback((ruleId: string, propertyId: string) => {
    const s = globalSocket;
    if (s?.connected && projectId) {
      s.emit('css_prop_removed', { projectId, ruleId, propertyId });
    }
  }, [projectId]);

  return {
    emitProjectAction,
    emitFileChanged,
    emitActiveFileChanged,
    emitCursorMove,
    emitBlockSelection,
    emitElementHover,
    emitFileAdded,
    emitFileRenamed,
    emitFileDeleted,
    emitVisualElementAdded,
    emitVisualElementUpdated,
    emitVisualElementRemoved,
    emitVisualElementsReordered,
    emitSave,
    emitPointerMove,
    emitBlockDragStart,
    emitBlockDragMove,
    emitBlockDragEnd,
    emitBlockStateSync,
    emitBlockConnect,
    emitBlockChange,
    emitBlockCreateEvent,
    emitTextFieldFocus,
    emitTextFieldCursor,
    emitTextFieldBlur,
    emitCssRuleAdded,
    emitCssRuleUpdated,
    emitCssRuleRemoved,
    emitCssRuleReordered,
    emitCssPropAdded,
    emitCssPropUpdated,
    emitCssPropRemoved,
  };
}

export function getSocket(): Socket | null {
  return globalSocket;
}