import { Effect, Option } from "effect";
import type * as vscode from "vscode";

import { LanguageId } from "../constants.ts";
import { VariablesService } from "../panel/variables/VariablesService.ts";
import type { MarimoNotebookDocument } from "../schemas/MarimoNotebookDocument.ts";
import type { NotebookCellId } from "../schemas/MarimoNotebookDocument.ts";
import { getTopologicalCellIds } from "./getTopologicalCellIds.ts";

/**
 * Get raw notebook cells in topological order, based on variable dependencies.
 *
 * Only includes Python cells (mo-python language ID) since the LSP server
 * only understands Python. Non-Python cells (SQL, markdown) are filtered out.
 *
 * Cells with no dependencies or stableId are placed at the end.
 */
export function getTopologicalCells(
  doc: MarimoNotebookDocument,
): Effect.Effect<Array<vscode.NotebookCell>, never, VariablesService> {
  return Effect.gen(function* () {
    const variablesService = yield* VariablesService;

    // Filter to only Python cells - LSP server only understands Python
    const cells = doc
      .getCells()
      .filter(
        (cell) =>
          cell.rawNotebookCell.document.languageId === LanguageId.Python ||
          cell.rawNotebookCell.document.languageId === "python",
      );

    if (cells.length === 0) {
      // Don't need to do anything for no cells
      return [];
    }

    const variables = yield* variablesService.getVariables(doc.id);

    // No variables yet - fall back to document order
    if (Option.isNone(variables)) {
      return cells.map((cell) => cell.rawNotebookCell);
    }

    const cellMap = new Map<NotebookCellId, vscode.NotebookCell>();
    const cellsWithoutIds: Array<vscode.NotebookCell> = [];

    for (const cell of cells) {
      if (Option.isNone(cell.id)) {
        cellsWithoutIds.push(cell.rawNotebookCell);
      } else {
        cellMap.set(cell.id.value, cell.rawNotebookCell);
      }
    }

    const sortedIds = getTopologicalCellIds(
      [...cellMap.keys()],
      variables.value,
    );

    // `getTopologicalCellIds` does not strictly guarantee a permutation of
    // its input (cyclic graphs can drop cells; stale variables can introduce
    // unknown IDs), so filter out any lookups that miss the map instead of
    // assuming presence.
    const sortedCells = sortedIds
      .map((id) => cellMap.get(id))
      .filter((cell): cell is vscode.NotebookCell => cell !== undefined);

    return [...sortedCells, ...cellsWithoutIds];
  });
}
