import type { Store } from 'pinia';
import * as jsonpatch from 'fast-json-patch';
import type { BackendMessage, FrontendMessage, HelloMessage, PatchMessage } from '../types/sync';
import { wsManager } from './useWebSocket';

export function useStateSync(stores: Record<string, Store>) {
  console.log('[StateSync] Initializing with stores:', Object.keys(stores));

  // Load initial seqs from sessionStorage
  const getStoredSeqs = (): Record<string, number> => {
    const stored = sessionStorage.getItem('sync_seqs');
    return stored ? JSON.parse(stored) : {};
  };

  const saveStoredSeqs = (seqs: Record<string, number>) => {
    sessionStorage.setItem('sync_seqs', JSON.stringify(seqs));
  };

  const currentSeqs = getStoredSeqs();
  console.log('[StateSync] Loaded stored seqs:', currentSeqs);

  // Track when we're applying backend patches to prevent circular updates
  const isApplyingBackendPatch: Record<string, boolean> = {};
  // Track previous state for each store to compute diffs
  const previousStates: Record<string, any> = {};

  const sendJson = (msg: FrontendMessage) => {
    wsManager.send(msg);
  };

  const handleMessage = (msg: BackendMessage) => {
    console.log('[StateSync] Received WebSocket message:', msg);
    console.log('[StateSync] Parsed message type:', msg.type);

    if (msg.type === 'snapshot') {
      console.log('[StateSync] Processing snapshot for store:', msg.store);
      const store = stores[msg.store];
      if (store) {
        console.log('[StateSync] Applying snapshot state:', msg.state);
        isApplyingBackendPatch[msg.store] = true;
        store.$patch(msg.state);
        // Update previousState to stay in sync
        previousStates[msg.store] = JSON.parse(JSON.stringify(store.$state));
        isApplyingBackendPatch[msg.store] = false;
        currentSeqs[msg.store] = msg.seq;
        saveStoredSeqs(currentSeqs);
        console.log('[StateSync] Snapshot applied, new seq:', msg.seq);
      } else {
        console.warn('[StateSync] Store not found:', msg.store);
      }
    } else if (msg.type === 'patch') {
      console.log('[StateSync] Processing patch for store:', msg.store, 'seq:', msg.seq);
      const store = stores[msg.store];
      if (store) {
        // Check for sequence gaps
        const lastSeq = currentSeqs[msg.store] || 0;
        console.log('[StateSync] Current seq:', lastSeq, 'Received seq:', msg.seq);
        if (msg.seq !== lastSeq + 1) {
          console.warn(`[StateSync] Sequence gap detected for ${msg.store}: expected ${lastSeq + 1}, got ${msg.seq}. Requesting resync.`);
          sendHello();
          return;
        }

        console.log('[StateSync] Applying patch:', msg.patch);
        const currentState = JSON.parse(JSON.stringify(store.$state));
        console.log('[StateSync] Current state before patch:', currentState);
        const newState = jsonpatch.applyPatch(currentState, msg.patch, false, false).newDocument;
        console.log('[StateSync] New state after patch:', newState);
        isApplyingBackendPatch[msg.store] = true;
        store.$patch(newState);
        // Update previousState to stay in sync
        previousStates[msg.store] = JSON.parse(JSON.stringify(store.$state));
        isApplyingBackendPatch[msg.store] = false;
        currentSeqs[msg.store] = msg.seq;
        saveStoredSeqs(currentSeqs);
        console.log('[StateSync] Patch applied successfully, new seq:', msg.seq);
      } else {
        console.warn('[StateSync] Store not found:', msg.store);
      }
    } else {
      console.log('[StateSync] Ignoring message type:', msg.type);
    }
  };

  const sendHello = () => {
    const hello: HelloMessage = {
      type: 'hello',
      seqs: currentSeqs
    };
    console.log('[StateSync] Sending hello with seqs:', currentSeqs);
    sendJson(hello);
  };

  const sendPatch = (storeName: string, patch: any[]) => {
    const seq = currentSeqs[storeName] || 0;
    const msg: PatchMessage = {
      type: 'patch',
      store: storeName,
      seq: seq,
      patch: patch
    };
    sendJson(msg);

    // Optimistically increment sequence number after sending
    // This ensures the next patch we send has the correct base sequence
    currentSeqs[storeName] = seq + 1;
    saveStoredSeqs(currentSeqs);
    console.log(`[StateSync] Incremented seq for "${storeName}" to ${seq + 1} after sending patch`);
  };

  // Set up watchers for each store to send patches on changes
  const unwatchFunctions: (() => void)[] = [];

  // Register handler (connection will be initiated after authentication)
  wsManager.addHandler(handleMessage);

  // Send hello when authenticated
  const checkAuthAndSendHello = () => {
    if (wsManager.isAuthenticated.value && wsManager.isConnected.value) {
      console.log('[StateSync] Authenticated and connected, sending hello');
      sendHello();
    }
  };

  // Watch for authentication
  const unwatchAuth = wsManager.isAuthenticated.value
    ? (checkAuthAndSendHello(), () => {})
    : (() => {
        const interval = setInterval(() => {
          if (wsManager.isAuthenticated.value && wsManager.isConnected.value) {
            clearInterval(interval);
            checkAuthAndSendHello();
          }
        }, 100);
        return () => clearInterval(interval);
      })();

  unwatchFunctions.push(unwatchAuth);

  for (const [storeName, store] of Object.entries(stores)) {
    previousStates[storeName] = JSON.parse(JSON.stringify(store.$state));
    isApplyingBackendPatch[storeName] = false;

    const unwatch = store.$subscribe((mutation, state) => {
      // Skip if we're currently applying a patch from the backend
      if (isApplyingBackendPatch[storeName]) {
        console.log(`[StateSync] Skipping patch send for "${storeName}" - applying backend update`);
        return;
      }

      console.log(`[StateSync] Store "${storeName}" changed, mutation type:`, mutation.type);
      console.log('[StateSync] Previous state:', previousStates[storeName]);
      console.log('[StateSync] New state:', state);

      const currentState = JSON.parse(JSON.stringify(state));
      const patch = jsonpatch.compare(previousStates[storeName], currentState);

      if (patch.length > 0) {
        console.log(`[StateSync] Sending ${patch.length} patch operations for "${storeName}":`, patch);
        sendPatch(storeName, patch);
        previousStates[storeName] = currentState;
      } else {
        console.log(`[StateSync] No changes detected for "${storeName}"`);
      }
    }, { detached: true });

    unwatchFunctions.push(unwatch);
  }

  return {
    sendPatch,
    cleanup: () => {
      wsManager.removeHandler(handleMessage);
      unwatchFunctions.forEach(unwatch => unwatch());
    }
  };
}
