import {
  addEdge,
  applyEdgeChanges,
  applyNodeChanges,
  type Edge,
  type Node,
  type OnConnect,
  type OnEdgesChange,
  type OnNodesChange,
} from '@xyflow/react'
import {contain, filter, find, isEqual} from 'licia'
import {create} from 'zustand/react'

export const useFlowStore = create<{
  nodes: Node[],
  getNodes: () => Node[],
  onNodesChange: OnNodesChange,
  getNodeById: (id: string) => Node | undefined,
  addNode: (node: Node) => void,
  removeNode: (id: string) => void,
  setNodes: (nodes: Node[]) => void,
  setNode: (node: Node) => void,

  edges: Edge[],
  getEdges: () => Edge[],
  onEdgesChange: OnEdgesChange,
  removeEdge: (id: string) => void,
  removeEdges: (ids: string[]) => void,
  setEdges: (edges: Edge[]) => void,

  onConnect: OnConnect,
}>((set, get) => ({
  nodes: [],
  getNodes: () => get().nodes,
  onNodesChange: changes => {
    set({
      nodes: applyNodeChanges(changes, get().nodes),
    })
  },
  getNodeById: (id: string) => find(get().nodes, node => isEqual(node.id, id)),
  addNode: node => set({nodes: get().nodes.concat(node)}),
  removeNode: id => {
    set({
      nodes: filter(get().nodes, node => !isEqual(node.id, id)),
    })
  },
  setNodes: nodes => set({nodes}),
  setNode: node => {
    set({
      nodes: get().nodes.map(n => {
        if (isEqual(node.id, n.id)) {
          return node
        }
        return n
      }),
    })
  },

  edges: [],
  getEdges: () => get().edges,
  onEdgesChange: changes => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    })
  },
  removeEdge: id => {
    set({
      edges: filter(get().edges, edge => !isEqual(edge.id, id)),
    })
  },
  removeEdges: ids => {
    set({
      edges: filter(get().edges, edge => !contain(ids, edge.id)),
    })
  },
  setEdges: edges => set({edges}),

  onConnect: connection => {
    set({
      edges: addEdge(connection, get().edges),
    })
  },
}))