import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
import { mergeRegister } from "@lexical/utils";
import { $getNodeByKey, $getSelection, COMMAND_PRIORITY_EDITOR, type LexicalNode, SELECTION_CHANGE_COMMAND } from "lexical";
import { useEffect, useState } from "react";
import { $isFocusableNode } from "../nodes/focusable_node";

export function SelectionPlugin() {
    const [editor] = useLexicalComposerContext();
    const [previousSelectedNodes, setPreviousSelectedNodes] = useState<string[]>([]);

    useEffect(() => mergeRegister(
        editor.registerCommand(SELECTION_CHANGE_COMMAND, () => {
            const selection = $getSelection();
            const nodes = selection?.getNodes() || [];

            previousSelectedNodes
                .map((key) => $getNodeByKey(key))
                .forEach(markTreeAsUnfocused);

            nodes.forEach(markTreeAsFocused);
            setPreviousSelectedNodes(nodes.map(n => n.getKey()));

            return false;
        }, COMMAND_PRIORITY_EDITOR),
    ));

    return null;
}


function markTreeAsFocused(node: LexicalNode | null) {
    if (!node) return;

    if ($isFocusableNode(node)) {
        node.setFocus(true);
    }

    markTreeAsFocused(node.getParent());
}

function markTreeAsUnfocused(node: LexicalNode | null) {
    if (!node) return;

    if ($isFocusableNode(node)) {
        node.setFocus(false);
    }

    markTreeAsUnfocused(node.getParent());
}