/* Copyright 2026 Marimo. All rights reserved. */

import { useAtomValue } from "jotai";
import type { GridLayout } from "@/components/editor/renderers/grid-layout/types";
import { cellRendererPlugins } from "@/components/editor/renderers/plugins";
import type { LayoutType } from "@/components/editor/renderers/types";
import { createReducerAndAtoms } from "@/utils/createReducer";
import { Logger } from "@/utils/Logger";
import { getNotebook } from "../cells/cells";
import { notebookCells } from "../cells/utils";
import { store } from "../state/jotai";

export type LayoutData = GridLayout | undefined;

export interface LayoutState {
  selectedLayout: LayoutType;
  layoutData: Partial<Record<LayoutType, LayoutData>>;
}

export function initialLayoutState(): LayoutState {
  return {
    selectedLayout: "vertical",
    layoutData: {},
  };
}

const { valueAtom: layoutStateAtom, useActions } = createReducerAndAtoms(
  initialLayoutState,
  {
    setLayoutView: (state, payload: LayoutType) => {
      return {
        ...state,
        selectedLayout: payload,
      };
    },
    setLayoutData: (
      state,
      payload: { layoutView: LayoutType; data: LayoutData },
    ) => {
      return {
        ...state,
        selectedLayout: payload.layoutView,
        layoutData: {
          ...state.layoutData,
          [payload.layoutView]: payload.data,
        },
      };
    },
    setCurrentLayoutData: (state, payload: LayoutData) => {
      return {
        ...state,
        layoutData: {
          ...state.layoutData,
          [state.selectedLayout]: payload,
        },
      };
    },
  },
);

export { layoutStateAtom };

export const useLayoutState = () => {
  return useAtomValue(layoutStateAtom);
};

export const useLayoutActions = () => {
  return useActions();
};

/**
 * Get the serialized layout data, to be used when saving.
 */
export function getSerializedLayout() {
  const notebook = getNotebook();
  const { layoutData, selectedLayout } = store.get(layoutStateAtom);

  // Vertical layout has no data, as it is the default.
  if (selectedLayout === "vertical") {
    return null;
  }

  if (layoutData === undefined) {
    return null;
  }

  const plugin = cellRendererPlugins.find(
    (plugin) => plugin.type === selectedLayout,
  );
  if (plugin === undefined) {
    Logger.error(`Unknown layout type: ${selectedLayout}`);
    return null;
  }
  const cells = notebookCells(notebook);
  // Fall back to the plugin's initial layout when the user has not yet
  // interacted with this layout — otherwise serializers that expect a
  // structured layout object crash on `undefined`.
  const data = layoutData[selectedLayout] ?? plugin.getInitialLayout(cells);
  return {
    type: selectedLayout,
    data: plugin.serializeLayout(data, cells),
  };
}