import type { Schema, Node, NodeType } from 'prosemirror-model';
import type { Transaction } from 'prosemirror-state';
import type { EditorView } from 'prosemirror-view';

export const setTextAlign = (tr: Transaction, schema: Schema, alignment: string) => {
  const { selection, doc } = tr;
  if (!selection || !doc) return tr;

  const { from, to } = selection;
  const { nodes } = schema;

  const blockquote = nodes.blockquote;
  const listItem = nodes.list_item;
  const paragraph = nodes.paragraph;

  interface Task {
    node: Node;
    pos: number;
    nodeType: NodeType;
  }

  const tasks: Task[] = [];
  alignment = alignment || '';

  const allowedNodeTypes = new Set([blockquote, listItem, paragraph]);

  doc.nodesBetween(from, to, (node, pos) => {
    const nodeType = node.type;
    const align = node.attrs.align || '';
    if (align !== alignment && allowedNodeTypes.has(nodeType)) {
      tasks.push({
        node,
        pos,
        nodeType,
      });
    }
    return true;
  });

  if (!tasks.length) return tr;

  tasks.forEach((task) => {
    const { node, pos, nodeType } = task;
    let { attrs } = node;
    if (alignment) attrs = { ...attrs, align: alignment };
    else attrs = { ...attrs, align: null };
    tr = tr.setNodeMarkup(pos, nodeType, attrs, node.marks);
  });

  return tr;
};

export const alignmentCommand = (view: EditorView, alignment: string) => {
  const { state } = view;
  const { schema, selection } = state;
  const tr = setTextAlign(state.tr.setSelection(selection), schema, alignment);
  view.dispatch(tr);
};