import React, { useEffect } from 'react';
import type { Decorator } from '@storybook/react';
import { useStore, Node, Edge, Component } from '../../src/store/store';

// Default mock state for stories
interface MockState {
  nodes?: Node[];
  edges?: Edge[];
  selectedNode?: Node | null;
  selectedEdge?: Edge | null;
  visibleStartNodeIds?: string[] | null;
  components?: Component[];
}

interface StoreDecoratorOptions {
  initialState?: MockState;
}

/**
 * Decorator that provides a mock Zustand store for components
 * Usage in stories:
 *
 * export const MyStory: Story = {
 *   decorators: [withStore({ initialState: { nodes: [...] } })],
 * };
 */
export const withStore = (options: StoreDecoratorOptions = {}): Decorator => {
  return (Story) => {
    // Reset store to initial state before each story
    useEffect(() => {
      const store = useStore.getState();

      // Apply mock state
      if (options.initialState?.nodes) {
        store.setNodes(options.initialState.nodes);
      }
      if (options.initialState?.edges) {
        store.setEdges(options.initialState.edges);
      }
      if (options.initialState?.selectedNode !== undefined) {
        store.setSelectedNode(options.initialState.selectedNode);
      }
      if (options.initialState?.selectedEdge !== undefined) {
        store.setSelectedEdge(options.initialState.selectedEdge);
      }
      if (options.initialState?.visibleStartNodeIds !== undefined) {
        store.setVisibleStartNodeIds(options.initialState.visibleStartNodeIds);
      }
      if (options.initialState?.components) {
        store.setComponents(options.initialState.components);
      }
    }, []);

    return <Story />;
  };
};

export default withStore;
