yaak-mountain-loop/apps/yaak-client/components/EnvironmentEditDialog.tsx

import type { Environment, Workspace } from "@yaakapp-internal/models";
import { duplicateModel, patchModel } from "@yaakapp-internal/models";
import type { TreeHandle, TreeNode, TreeProps } from "@yaakapp-internal/ui";
import { Banner, Icon, InlineCode, SplitLayout, Tree, IconTooltip } from "@yaakapp-internal/ui";
import { atom, useAtomValue } from "jotai";
import { atomFamily } from "jotai-family";
import {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useLayoutEffect,
  useRef,
  useState,
} from "react";
import { createSubEnvironment } from "../commands/createEnvironment";
import { activeWorkspaceAtom, activeWorkspaceIdAtom } from "../hooks/useActiveWorkspace";
import {
  environmentsBreakdownAtom,
  useEnvironmentsBreakdown,
} from "../hooks/useEnvironmentsBreakdown";
import { useHotKey } from "../hooks/useHotKey";
import { atomWithKVStorage } from "../lib/atoms/atomWithKVStorage";
import { deleteModelWithConfirm } from "../lib/deleteModelWithConfirm";
import { fireAndForget } from "../lib/fireAndForget";
import { jotaiStore } from "../lib/jotai";
import { isBaseEnvironment, isSubEnvironment } from "../lib/model_util";
import { resolvedModelName } from "../lib/resolvedModelName";
import { showColorPicker } from "../lib/showColorPicker";
import type { ContextMenuProps, DropdownItem } from "./core/Dropdown";
import { ContextMenu } from "./core/Dropdown";
import { IconButton } from "./core/IconButton";
import type { PairEditorHandle } from "./core/PairEditor";
import { EnvironmentColorIndicator } from "./EnvironmentColorIndicator";
import { EnvironmentEditor } from "./EnvironmentEditor";
import { EnvironmentSharableTooltip } from "./EnvironmentSharableTooltip";

const collapsedFamily = atomFamily((treeId: string) => {
  const key = ["env_collapsed", treeId ?? "n/a"];
  return atomWithKVStorage<Record<string, boolean>>(key, {});
});

interface Props {
  initialEnvironmentId: string | null;
  setRef?: (ref: PairEditorHandle | null) => void;
}

type TreeModel = Environment | Workspace;

/** Lets tree row slots create a sub-environment and select it in the editor. */
const CreateSubEnvironmentContext = createContext<() => Promise<void>>(async () => {});

export function EnvironmentEditDialog({ initialEnvironmentId, setRef }: Props) {
  const { allEnvironments, baseEnvironment, baseEnvironments } = useEnvironmentsBreakdown();
  const [selectedEnvironmentId, setSelectedEnvironmentId] = useState<string | null>(
    initialEnvironmentId ?? null,
  );

  const selectedEnvironment =
    selectedEnvironmentId != null
      ? allEnvironments.find((e) => e.id === selectedEnvironmentId)
      : baseEnvironment;

  return (
    <SplitLayout
      storageKey="env_editor"
      defaultRatio={0.75}
      layout="horizontal"
      className="gap-0"
      resizeHandleClassName="-translate-x-px"
      firstSlot={() => (
        <EnvironmentEditDialogSidebar
          selectedEnvironmentId={selectedEnvironment?.id ?? null}
          setSelectedEnvironmentId={setSelectedEnvironmentId}
        />
      )}
      secondSlot={() => (
        <div className="grid grid-rows-[auto_minmax(0,1fr)]">
          {baseEnvironments.length > 1 ? (
            <div className="p-3">
              <Banner color="notice">
                There are multiple base environments for this workspace. Please delete the
                environments you no longer need.
              </Banner>
            </div>
          ) : (
            <span />
          )}
          {selectedEnvironment == null ? (
            <div className="p-3 mt-10">
              <Banner color="danger">
                Failed to find selected environment <InlineCode>{selectedEnvironmentId}</InlineCode>
              </Banner>
            </div>
          ) : (
            <EnvironmentEditor
              key={selectedEnvironment.id}
              setRef={setRef}
              className="pl-4 pt-3"
              environment={selectedEnvironment}
            />
          )}
        </div>
      )}
    />
  );
}

const sharableTooltip = (
  <IconTooltip
    tabIndex={-1}
    icon="eye"
    iconSize="sm"
    content="This environment will be included in Directory Sync and data exports"
  />
);

function EnvironmentEditDialogSidebar({
  selectedEnvironmentId,
  setSelectedEnvironmentId,
}: {
  selectedEnvironmentId: string | null;
  setSelectedEnvironmentId: (id: string | null) => void;
}) {
  const activeWorkspaceId = useAtomValue(activeWorkspaceIdAtom) ?? "";
  const treeId = `environment.${activeWorkspaceId}.sidebar`;
  const treeRef = useRef<TreeHandle>(null);
  const { allEnvironments, baseEnvironment, baseEnvironments } = useEnvironmentsBreakdown();

  // oxlint-disable-next-line react-hooks/exhaustive-deps -- none
  useLayoutEffect(() => {
    if (selectedEnvironmentId == null) return;
    treeRef.current?.selectItem(selectedEnvironmentId);
    treeRef.current?.focus();
  }, []);

  // A just-created environment to select once it reaches the model store. Selecting it before
  // then would show a "failed to find" error in the editor, and the tree would drop the
  // selection of an item it doesn't have yet.
  const [pendingSelectionId, setPendingSelectionId] = useState<string | null>(null);
  const tree = useAtomValue(treeAtom);
  useEffect(() => {
    if (pendingSelectionId == null) return;
    if (!allEnvironments.some((e) => e.id === pendingSelectionId)) return;

    setPendingSelectionId(null);
    setSelectedEnvironmentId(pendingSelectionId);

    // The tree leaves out sub-environments while there are multiple base environments
    const node = tree == null ? null : findTreeNode(tree, pendingSelectionId);
    if (node == null) return;
    if (node.parent != null) {
      // Expand the parent so the new item is visible to be selected
      const parentId = node.parent.item.id;
      jotaiStore.set(collapsedFamily(treeId), (prev) => ({ ...prev, [parentId]: false }));
    }
    // Wait a frame so the tree's own pending selection fixups, which may have been scheduled
    // before the new item arrived, don't reset this selection
    requestAnimationFrame(() => treeRef.current?.selectItem(pendingSelectionId));
  }, [allEnvironments, pendingSelectionId, setSelectedEnvironmentId, tree, treeId]);

  const handleDeleteEnvironment = useCallback(
    async (environment: Environment) => {
      await deleteModelWithConfirm(environment);
      if (selectedEnvironmentId === environment.id) {
        setSelectedEnvironmentId(baseEnvironment?.id ?? null);
      }
    },
    [baseEnvironment?.id, selectedEnvironmentId, setSelectedEnvironmentId],
  );

  const handleCreateSubEnvironment = useCallback(async () => {
    const { baseEnvironment } = jotaiStore.get(environmentsBreakdownAtom);
    if (baseEnvironment == null) return;
    const id = await createSubEnvironment.mutateAsync(baseEnvironment);
    // Show the new environment in the editor without changing the workspace's active environment
    if (id != null) setPendingSelectionId(id);
  }, []);

  const treeHasFocus = useCallback(() => treeRef.current?.hasFocus() ?? false, []);

  const getSelectedTreeModels = useCallback(
    () => treeRef.current?.getSelectedItems() as TreeModel[] | undefined,
    [],
  );

  const handleRenameSelected = useCallback(() => {
    const items = getSelectedTreeModels();
    if (items?.length === 1 && items[0] != null) {
      treeRef.current?.renameItem(items[0].id);
    }
  }, [getSelectedTreeModels]);

  const handleDeleteSelected = useCallback(
    (items: TreeModel[]) => deleteModelWithConfirm(items),
    [],
  );

  const handleDuplicateSelected = useCallback(
    async (items: TreeModel[]) => {
      if (items.length === 1 && items[0]) {
        const newId = await duplicateModel(items[0]);
        setSelectedEnvironmentId(newId);
      } else {
        await Promise.all(items.map(duplicateModel));
      }
    },
    [setSelectedEnvironmentId],
  );

  useHotKey("sidebar.selected.rename", handleRenameSelected, {
    enable: treeHasFocus,
    allowDefault: true,
    priority: 100,
  });
  useHotKey(
    "sidebar.selected.delete",
    useCallback(() => {
      const items = getSelectedTreeModels();
      if (items) {
        fireAndForget(handleDeleteSelected(items));
      }
    }, [getSelectedTreeModels, handleDeleteSelected]),
    { enable: treeHasFocus, priority: 100 },
  );
  useHotKey(
    "sidebar.selected.duplicate",
    useCallback(async () => {
      const items = getSelectedTreeModels();
      if (items) await handleDuplicateSelected(items);
    }, [getSelectedTreeModels, handleDuplicateSelected]),
    { enable: treeHasFocus, priority: 100 },
  );

  const getContextMenu = useCallback(
    (items: TreeModel[]): ContextMenuProps["items"] => {
      const environment = items[0];
      const addEnvironmentItem: DropdownItem = {
        label: "Create Sub Environment",
        leftSlot: <Icon icon="plus" />,
        onSelect: handleCreateSubEnvironment,
      };
      // Sub-environments aren't shown while there are multiple base environments, same as the
      // plus button on the base environment row
      const canCreateSubEnvironment = baseEnvironments.length <= 1;

      if (environment == null || environment.model !== "environment") {
        return canCreateSubEnvironment ? [addEnvironmentItem] : [];
      }

      const singleEnvironment = items.length === 1;
      const canDeleteEnvironment =
        isSubEnvironment(environment) ||
        (isBaseEnvironment(environment) && baseEnvironments.length > 1);

      const menuItems: DropdownItem[] = [
        {
          label: "Rename",
          leftSlot: <Icon icon="pencil" />,
          hidden: isBaseEnvironment(environment) || !singleEnvironment,
          hotKeyAction: "sidebar.selected.rename",
          hotKeyLabelOnly: true,
          onSelect: () => {
            // Not sure why this is needed, but without it the
            // edit input blurs immediately after opening.
            requestAnimationFrame(() => handleRenameSelected());
          },
        },
        {
          label: "Duplicate",
          leftSlot: <Icon icon="copy" />,
          hidden: isBaseEnvironment(environment),
          hotKeyAction: "sidebar.selected.duplicate",
          hotKeyLabelOnly: true,
          onSelect: () => handleDuplicateSelected(items),
        },
        {
          label: environment.color ? "Change Color" : "Assign Color",
          leftSlot: <Icon icon="palette" />,
          hidden: isBaseEnvironment(environment) || !singleEnvironment,
          onSelect: async () => showColorPicker(environment),
        },
        {
          label: `Make ${environment.public ? "Private" : "Sharable"}`,
          leftSlot: <Icon icon={environment.public ? "eye_closed" : "eye"} />,
          rightSlot: <EnvironmentSharableTooltip />,
          hidden: items.length > 1,
          onSelect: async () => {
            await patchModel(environment, { public: !environment.public });
          },
        },
        {
          color: "danger",
          label: "Delete",
          hotKeyAction: "sidebar.selected.delete",
          hotKeyLabelOnly: true,
          hidden: !canDeleteEnvironment,
          leftSlot: <Icon icon="trash" />,
          onSelect: () => handleDeleteEnvironment(environment),
        },
      ];

      // Add sub environment to base environment
      if (isBaseEnvironment(environment) && singleEnvironment && canCreateSubEnvironment) {
        menuItems.push({ type: "separator" });
        menuItems.push(addEnvironmentItem);
      }

      return menuItems;
    },
    [
      baseEnvironments.length,
      handleCreateSubEnvironment,
      handleDeleteEnvironment,
      handleDuplicateSelected,
      handleRenameSelected,
    ],
  );

  const handleDragEnd = useCallback(async function handleDragEnd({
    items,
    children,
    insertAt,
  }: {
    items: TreeModel[];
    children: TreeModel[];
    insertAt: number;
  }) {
    const prev = children[insertAt - 1] as Exclude<TreeModel, Workspace>;
    const next = children[insertAt] as Exclude<TreeModel, Workspace>;

    const beforePriority = prev?.sortPriority ?? 0;
    const afterPriority = next?.sortPriority ?? 0;
    const shouldUpdateAll = afterPriority - beforePriority < 1;

    try {
      if (shouldUpdateAll) {
        // Add items to children at insertAt
        children.splice(insertAt, 0, ...items);
        await Promise.all(children.map((m, i) => patchModel(m, { sortPriority: i * 1000 })));
      } else {
        const range = afterPriority - beforePriority;
        const increment = range / (items.length + 2);
        await Promise.all(
          items.map((m, i) => {
            const sortPriority = beforePriority + (i + 1) * increment;
            // Spread item sortPriority out over before/after range
            return patchModel(m, { sortPriority });
          }),
        );
      }
    } catch (e) {
      console.error(e);
    }
  }, []);

  const handleActivate = useCallback(
    (item: TreeModel) => {
      setSelectedEnvironmentId(item.id);
    },
    [setSelectedEnvironmentId],
  );

  const renderContextMenuFn = useCallback<NonNullable<TreeProps<TreeModel>["renderContextMenu"]>>(
    ({ items, position, onClose }) =>
      // Right-clicking empty space has nothing to offer while there are multiple base environments
      items.length === 0 ? null : (
        <ContextMenu items={items as DropdownItem[]} triggerPosition={position} onClose={onClose} />
      ),
    [],
  );

  return (
    <aside className="x-theme-sidebar h-full w-full min-w-0 grid overflow-y-auto border-r border-border-subtle ">
      {tree != null && (
        <CreateSubEnvironmentContext.Provider value={handleCreateSubEnvironment}>
          <div className="pt-2">
            <Tree
              ref={treeRef}
              treeId={treeId}
              collapsedAtom={collapsedFamily(treeId)}
              className="px-2 pb-10"
              root={tree}
              getContextMenu={getContextMenu}
              renderContextMenu={renderContextMenuFn}
              onDragEnd={handleDragEnd}
              getItemKey={(i) => `${i.id}::${i.name}`}
              ItemLeftSlotInner={ItemLeftSlotInner}
              ItemRightSlot={ItemRightSlot}
              ItemInner={ItemInner}
              onActivate={handleActivate}
              getEditOptions={getEditOptions}
            />
          </div>
        </CreateSubEnvironmentContext.Provider>
      )}
    </aside>
  );
}

const treeAtom = atom<TreeNode<TreeModel> | null>((get) => {
  const activeWorkspace = get(activeWorkspaceAtom);
  const { baseEnvironment, baseEnvironments, subEnvironments } = get(environmentsBreakdownAtom);
  if (activeWorkspace == null || baseEnvironment == null) return null;

  const root: TreeNode<TreeModel> = {
    item: activeWorkspace,
    parent: null,
    children: [],
    depth: 0,
  };

  for (const item of baseEnvironments) {
    root.children?.push({
      item,
      parent: root,
      depth: 0,
      draggable: false,
    });
  }

  const parent = root.children?.[0];
  if (baseEnvironments.length <= 1 && parent != null) {
    parent.children = subEnvironments.map((item) => ({
      item,
      parent,
      depth: 1,
      localDrag: true,
    }));
  }

  return root;
});

function findTreeNode(node: TreeNode<TreeModel>, id: string): TreeNode<TreeModel> | null {
  if (node.item.id === id) return node;
  for (const child of node.children ?? []) {
    const found = findTreeNode(child, id);
    if (found != null) return found;
  }
  return null;
}

function ItemLeftSlotInner({ item }: { item: TreeModel }) {
  const { baseEnvironments } = useEnvironmentsBreakdown();
  return baseEnvironments.length > 1 ? (
    <Icon icon="alert_triangle" color="notice" />
  ) : (
    item.model === "environment" && item.color && <EnvironmentColorIndicator environment={item} />
  );
}

function ItemRightSlot({ item }: { item: TreeModel }) {
  const { baseEnvironments } = useEnvironmentsBreakdown();
  const handleCreateSubEnvironment = useContext(CreateSubEnvironmentContext);
  return (
    <>
      {item.model === "environment" && baseEnvironments.length <= 1 && isBaseEnvironment(item) && (
        <IconButton
          size="sm"
          color="custom"
          iconSize="sm"
          icon="plus_circle"
          className="opacity-50 hover:opacity-100"
          title="Add Sub-Environment"
          onClick={handleCreateSubEnvironment}
        />
      )}
    </>
  );
}

function ItemInner({ item }: { item: TreeModel }) {
  return (
    <div className="grid grid-cols-[auto_minmax(0,1fr)] w-full items-center">
      {item.model === "environment" && item.public ? (
        <div className="mr-2 flex items-center">{sharableTooltip}</div>
      ) : (
        <span aria-hidden />
      )}
      <div className="truncate min-w-0 text-left">{resolvedModelName(item)}</div>
    </div>
  );
}

function getEditOptions(item: TreeModel) {
  const options: ReturnType<NonNullable<TreeProps<TreeModel>["getEditOptions"]>> = {
    defaultValue: item.name,
    placeholder: "Name",
    async onChange(item, name) {
      await patchModel(item, { name });
    },
  };
  return options;
}