import type { TreeDataNode } from "antd";
import React from "react";
import { create } from "zustand";
import type { NodeContext } from "../components/Graph";

export interface TreeStore {
  nodesFlatById: Map<React.Key, TreeDataNode>,
  parentIdByChildId: Map<React.Key, string>;
  rootNodes: TreeDataNode[];
  tree: TreeDataNode[];
  add: (childNode: NodeContext, parentNodeId: string | undefined) => void;
  remove: (nodeId: string) => void;
  reset: () => void;
}

export const useGraphLayersTreeStore = create<TreeStore>()((set) => ({
  nodesFlatById: new Map<React.Key, TreeDataNode>(),
  parentIdByChildId: new Map<React.Key, string>(),
  rootNodes: [],
  tree: [],
  add: (childNodeContext, parentNodeId) => set((state) => {
    const childNode = nodeContextToTreeNode(childNodeContext);
    if (parentNodeId) {
      const parentNode = state.nodesFlatById.get(parentNodeId);
      if (parentNode) {
        parentNode.children = parentNode.children ? [...parentNode.children, childNode] : [childNode];
        const parentIdByChildId = new Map(state.parentIdByChildId);
        parentIdByChildId.set(childNode.key, parentNodeId);
        const nodesFlatById = new Map(state.nodesFlatById);
        nodesFlatById.set(childNode.key, childNode);

        const newState = {
          nodesFlatById: nodesFlatById,
          parentIdByChildId: parentIdByChildId,
          rootNodes: [...state.rootNodes],
          tree: createTree([...state.rootNodes], nodesFlatById)
        }

        return newState;
      } else {
        throw Error(`There is no parent node with id: ${parentNodeId}`)
      }
    } else {
      const nodesFlatById = new Map(state.nodesFlatById);
      nodesFlatById.set(childNode.key, childNode);
      const newRootNodes = [...state.rootNodes, childNode];
      const newState = {
        nodesFlatById: nodesFlatById,
        parentIdByChildId: state.parentIdByChildId,
        rootNodes: newRootNodes,
        tree: createTree(newRootNodes, nodesFlatById)
      }
      return newState;
    }
  }),
  remove: (nodeId) => set((state) => {
    const node = state.nodesFlatById.get(nodeId);

    if (!node) {
      return state;
    }

    const nodesFlatById = new Map(state.nodesFlatById);
    nodesFlatById.delete(nodeId);
    const parentNodeId = state.parentIdByChildId.get(nodeId);

    if (parentNodeId) {
      const parentNode = state.nodesFlatById.get(parentNodeId);

      if (parentNode) {
        parentNode.children = parentNode.children
          ? [...parentNode.children?.filter(n => n.key !== nodeId)]
          : parentNode.children;
        const parentIdByChildId = new Map(state.parentIdByChildId);
        parentIdByChildId.delete(nodeId);

        return {
          rootNodes: [...state.rootNodes],
          nodesFlatById: nodesFlatById,
          parentIdByChildId: parentIdByChildId,
          tree: createTree([...state.rootNodes], nodesFlatById)
        }
      }
    }

    return {
      rootNodes: [...state.rootNodes.filter(n => n.key !== nodeId)],
      nodesFlatById: nodesFlatById,
      parentIdByChildId: state.parentIdByChildId,
      tree: createTree([...state.rootNodes], nodesFlatById)
    }
  }),
  reset: () => set({
    nodesFlatById: new Map<React.Key, TreeDataNode>(),
    parentIdByChildId: new Map<React.Key, string>(),
    rootNodes: [],
    tree: [],
  }),
}));

export function createTree(nodes: TreeDataNode[], nodesFlatById: Map<React.Key, TreeDataNode>): TreeDataNode[] {
  const n = [...nodes];
  const result = [];

  for (const node of n) {
    const stateNode = nodesFlatById.get(node.key);

    if (stateNode) {
      stateNode.children = createTree(stateNode?.children ?? [], nodesFlatById);
      result.push(stateNode);
    }
  }

  return result;
}

export function nodeContextToTreeNode(nodeContext: NodeContext): TreeDataNode {
  return {
    key: nodeContext.nodeId,
    title: nodeContext.nodeName,
  } as TreeDataNode;
}