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

import type { createStore } from "jotai";
import { getCellEditorView } from "@/core/cells/cells";
import { type CellId, HTMLCellId } from "@/core/cells/ids";
import { Logger } from "@/utils/Logger";

export function focusCellEditor(
  store: ReturnType<typeof createStore>,
  cellId: CellId,
): void {
  const editor = getCellEditorView(cellId);
  if (editor) {
    editor.focus();
  } else {
    Logger.warn(
      `[CellFocusManager] focusCellEditor: element not found: ${cellId}`,
    );
  }
}

export function focusCell(cellId: CellId): void {
  const element = document.getElementById(HTMLCellId.create(cellId));
  if (element) {
    tryFocus(element);
  } else {
    Logger.warn(`[CellFocusManager] focusCell: element not found: ${cellId}`);
  }
}

/**
 * Run a callback after two frames.
 * It is somewhat common/safer to run code after two frames to ensure the DOM is fully rendered.
 */
export function raf2(callback: () => void): void {
  requestAnimationFrame(() => {
    requestAnimationFrame(callback);
  });
}

/**
 * Checks if the cell is focused at the top level.
 */
export function isAnyCellFocused(): boolean {
  return (
    document.activeElement instanceof HTMLElement &&
    document.activeElement.classList.contains("marimo-cell")
  );
}

export function tryFocus(dom: HTMLElement) {
  try {
    dom.focus();
  } catch {
    Logger.warn("[CellFocusManager] element may not be focusable", dom);
  }
}