Draft js - Add entity on editorState change

Question

I am using Draft JS for its Mentions plugin.

What I would like to do is to give the user the possibility to insert mentions programmatically so that when he types (for example) "@*", the following word will be managed as an entity.

I've written this function that runs on editorState change to try and do exactly that:

export const replaceString = (editorState, setEditorState) => {
    const contentState = editorState.getCurrentContent();
    const selectionState = editorState.getSelection();
    const block = contentState.getBlockForKey(selectionState.getAnchorKey());
    
    const blockAsArray = block.getText().split(' ');
    const lastWord = blockAsArray.pop();
    if (lastWord.startsWith('@*') || lastWord.startsWith('#*') || lastWord.startsWith('$*') || lastWord.endsWith('*')) {

        const entity = contentState.createEntity(
            "MENTION",
            "MUTABLE",
            {nature: 'remedy', lastWord}
        );
        
        const entityKey = entity.getLastCreatedEntityKey();
        
        const newContentState = Modifier.replaceText(
            contentState,
            selectionState.merge({
                anchorOffset: selectionState.getEndOffset() - lastWord.length,
                focusOffset: selectionState.getEndOffset()
            }),
            lastWord,
            [getInlineStyle(lastWord[0]) || editorState.getCurrentInlineStyle()],
            entityKey
        );
        
        editorState = EditorState.push(
            editorState,
            newContentState
        );
    }
    
    setEditorState(editorState)
}

const getInlineStyle = (type) => {
    switch (type) {
        case '@':
            return "REMEDY";
        case '#':
            return "FAMILY";
        default:
            return;
    }
}

const defaultStyle = {
    padding: "5px 10px",
    color: "white",
    fontWeight: "bold",
    borderRadius: "10px",
    textDecoration: "none"
}
export const styleMap = {
    'REMEDY': {
        ...defaultStyle,
        backgroundColor: "green"
    },
    'FAMILY': {
        ...defaultStyle,
        backgroundColor: "blue"
    },
};

As you can see, when the lastword meets the if conditions, the following content will be treated as an entity and it works fine. The only problem is with styles.

I would like the inline style change to cover only the specific selectionState but this is what actually happens

The first word (and not the second) gets treated exactly as an entity when it comes to the editor state and that is perfect. But both words get the new inline style and that is the main problem.

Is there any problem in the replaceText function? Is there any other way I can style a mention?