/**
 * AeThex Visual Scripting - State Store
 * Zustand store for managing visual script state
 */

import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import {
  Node,
  Edge,
  Connection,
  addEdge,
  applyNodeChanges,
  applyEdgeChanges,
  NodeChange,
  EdgeChange,
} from 'reactflow';
import { NodeData } from '../lib/visual-scripting/code-generator';

interface VisualScriptState {
  // Current script data
  nodes: Node<NodeData>[];
  edges: Edge[];

  // UI state
  selectedNodes: string[];
  selectedEdges: string[];

  // Generated code
  generatedCode: string;

  // History for undo/redo
  history: { nodes: Node<NodeData>[]; edges: Edge[] }[];
  historyIndex: number;

  // Actions
  setNodes: (nodes: Node<NodeData>[]) => void;
  setEdges: (edges: Edge[]) => void;
  onNodesChange: (changes: NodeChange[]) => void;
  onEdgesChange: (changes: EdgeChange[]) => void;
  onConnect: (connection: Connection) => void;
  addNode: (node: Node<NodeData>) => void;
  removeNode: (nodeId: string) => void;
  updateNodeData: (nodeId: string, data: Partial<NodeData>) => void;
  updateNodeValue: (nodeId: string, key: string, value: any) => void;
  setSelectedNodes: (nodeIds: string[]) => void;
  setSelectedEdges: (edgeIds: string[]) => void;
  setGeneratedCode: (code: string) => void;
  clearScript: () => void;
  saveToHistory: () => void;
  undo: () => void;
  redo: () => void;
  loadScript: (nodes: Node<NodeData>[], edges: Edge[]) => void;
}

export const useVisualScriptStore = create<VisualScriptState>()(
  persist(
    (set, get) => ({
      nodes: [],
      edges: [],
      selectedNodes: [],
      selectedEdges: [],
      generatedCode: '',
      history: [],
      historyIndex: -1,

      setNodes: (nodes) => set({ nodes }),

      setEdges: (edges) => set({ edges }),

      onNodesChange: (changes) => {
        set({
          nodes: applyNodeChanges(changes, get().nodes) as Node<NodeData>[],
        });
      },

      onEdgesChange: (changes) => {
        set({
          edges: applyEdgeChanges(changes, get().edges),
        });
      },

      onConnect: (connection) => {
        // Validate connection types could go here
        set({
          edges: addEdge(
            {
              ...connection,
              id: `edge-${Date.now()}`,
              type: 'smoothstep',
              animated: connection.sourceHandle === 'flow',
            },
            get().edges
          ),
        });
        get().saveToHistory();
      },

      addNode: (node) => {
        set({ nodes: [...get().nodes, node] });
        get().saveToHistory();
      },

      removeNode: (nodeId) => {
        set({
          nodes: get().nodes.filter((n) => n.id !== nodeId),
          edges: get().edges.filter(
            (e) => e.source !== nodeId && e.target !== nodeId
          ),
        });
        get().saveToHistory();
      },

      updateNodeData: (nodeId, data) => {
        set({
          nodes: get().nodes.map((node) =>
            node.id === nodeId
              ? { ...node, data: { ...node.data, ...data } }
              : node
          ),
        });
      },

      updateNodeValue: (nodeId, key, value) => {
        set({
          nodes: get().nodes.map((node) =>
            node.id === nodeId
              ? {
                  ...node,
                  data: {
                    ...node.data,
                    values: {
                      ...node.data.values,
                      [key]: value,
                    },
                  },
                }
              : node
          ),
        });
      },

      setSelectedNodes: (nodeIds) => set({ selectedNodes: nodeIds }),

      setSelectedEdges: (edgeIds) => set({ selectedEdges: edgeIds }),

      setGeneratedCode: (code) => set({ generatedCode: code }),

      clearScript: () => {
        set({
          nodes: [],
          edges: [],
          selectedNodes: [],
          selectedEdges: [],
          generatedCode: '',
        });
        get().saveToHistory();
      },

      saveToHistory: () => {
        const { nodes, edges, history, historyIndex } = get();
        const newHistory = history.slice(0, historyIndex + 1);
        newHistory.push({ nodes: [...nodes], edges: [...edges] });

        // Keep only last 50 states
        if (newHistory.length > 50) {
          newHistory.shift();
        }

        set({
          history: newHistory,
          historyIndex: newHistory.length - 1,
        });
      },

      undo: () => {
        const { history, historyIndex } = get();
        if (historyIndex > 0) {
          const prevState = history[historyIndex - 1];
          set({
            nodes: prevState.nodes,
            edges: prevState.edges,
            historyIndex: historyIndex - 1,
          });
        }
      },

      redo: () => {
        const { history, historyIndex } = get();
        if (historyIndex < history.length - 1) {
          const nextState = history[historyIndex + 1];
          set({
            nodes: nextState.nodes,
            edges: nextState.edges,
            historyIndex: historyIndex + 1,
          });
        }
      },

      loadScript: (nodes, edges) => {
        set({ nodes, edges });
        get().saveToHistory();
      },
    }),
    {
      name: 'aethex-visual-script',
      partialize: (state) => ({
        nodes: state.nodes,
        edges: state.edges,
      }),
    }
  )
);
