import { type Edge, type Node, type Viewport } from '@xyflow/svelte';
import { getContext, setContext } from 'svelte';
import type { TinyflowNodeNormalizer } from '../utils/nodeInteraction';

const DEFAULT_VIEWPORT: Viewport = { x: 250, y: 100, zoom: 1 };
const TINYFLOW_STORE_CONTEXT = Symbol('tinyflow_store');

export const createStore = () => {
  let nodesInternal = $state.raw([] as Node[]);
  let edgesInternal = $state.raw([] as Edge[]);
  let viewport = $state.raw({ ...DEFAULT_VIEWPORT } as Viewport);
  let normalizeNode: TinyflowNodeNormalizer = (node) => node;
  const pendingEditFlushers = new Set<() => void>();

  const normalizeNodes = (nodes: Node[]) => nodes.map(normalizeNode);

  return {
    // nodes: nodesInternal,
    // edges: edgesInternal,
    // viewport,
    setNodeNormalizer: (normalizer: TinyflowNodeNormalizer) => {
      normalizeNode = normalizer;
      nodesInternal = normalizeNodes(nodesInternal);
    },
    init: (nodes: Node[], edges: Edge[], nextViewport?: Viewport | null) => {
      nodesInternal = normalizeNodes(nodes);
      edgesInternal = edges;
      viewport = nextViewport ? { ...nextViewport } : { ...DEFAULT_VIEWPORT };
    },

    getNodes: () => nodesInternal,
    setNodes: (nodes: Node[]) => {
      nodesInternal = normalizeNodes(nodes);
    },
    getEdges: () => edgesInternal,
    setEdges: (edges: Edge[]) => {
      edgesInternal = edges;
    },
    getViewport: () => viewport,
    setViewport: (v: Viewport) => {
      viewport = v;
    },
    registerPendingEditFlusher: (flusher: () => void) => {
      pendingEditFlushers.add(flusher);
      return () => pendingEditFlushers.delete(flusher);
    },
    flushPendingEdits: () => {
      [...pendingEditFlushers].forEach((flusher) => flusher());
    },

    getNode: (id: string) => nodesInternal.find((node) => node.id === id),
    addNode: (node: Node) => {
      nodesInternal = [...nodesInternal, normalizeNode(node)];
    },
    removeNode: (id: string) => {
      nodesInternal = nodesInternal.filter((node) => node.id !== id);
    },
    updateNode: (id: string, node: Node) => {
      nodesInternal = nodesInternal.map((n) =>
        n.id === id ? normalizeNode({ ...n, ...node }) : n,
      );
    },
    updateNodes: (update: (nodes: Node[]) => Node[]) => {
      nodesInternal = normalizeNodes(update(nodesInternal));
    },
    updateNodeData: (id: string, data: Node['data']) => {
      nodesInternal = nodesInternal.map((n) =>
        n.id === id ? { ...n, data: { ...n.data, ...data } } : n,
      );
    },
    selectNodeOnly: (id: string) => {
      nodesInternal = nodesInternal.map((n) =>
        n.id === id ? { ...n, selected: true } : { ...n, selected: false },
      );
    },
    getEdge: (id: string) => edgesInternal.find((edge) => edge.id === id),
    addEdge: (edge: Edge) => {
      edgesInternal = [...edgesInternal, edge];
    },
    removeEdge: (id: string) => {
      edgesInternal = edgesInternal.filter((edge) => edge.id !== id);
    },
    updateEdge: (id: string, edge: Edge) => {
      edgesInternal = edgesInternal.map((e) =>
        e.id === id ? { ...e, ...edge } : e,
      );
    },
    updateEdges: (update: (edges: Edge[]) => Edge[]) => {
      edgesInternal = update(edgesInternal);
    },
    updateEdgeData: (id: string, data: Edge['data']) => {
      edgesInternal = edgesInternal.map((e) =>
        e.id === id ? { ...e, data: { ...e.data, ...data } } : e,
      );
    },
  };
};

export type TinyflowStore = ReturnType<typeof createStore>;

export const createTinyflowStoreContext = () => {
  const store = createStore();
  setContext(TINYFLOW_STORE_CONTEXT, store);
  return store;
};

export const useTinyflowStore = () => {
  const store = getContext<TinyflowStore | undefined>(TINYFLOW_STORE_CONTEXT);
  if (!store) {
    throw new Error('Tinyflow store context is not initialized');
  }
  return store;
};