from prompt_toolkit.document import Document as _Document
from prompt_toolkit.key_binding.key_processor import KeyPress as _KeyPress
from prompt_toolkit.keys import Keys as _Keys
from xonsh.events import events as _events

@_events.on_ptk_create
def _custom_keybindings(bindings, **kw):

    ...  # other, smaller stuff

    def del_first_words(n, event, space_at_start=False):
        orig_text = event.current_buffer.text
        *_, new_text = event.current_buffer.text.split(maxsplit=int(n))
        if space_at_start:
            new_text = f" {new_text}"
        if new_text != orig_text:
            pos = event.current_buffer.cursor_position - len(orig_text) + len(new_text)
            # Setting new document because using cursor_left sometimes is off by n
            event.current_buffer.document = _Document(new_text, pos)

    mods = {
        'f': ('f"{', '}"'),
        'u': ('sudo ', ''),
        'w': ('pkg which $(', ')'),
        'W': ('pkg which $(which ', ')'),
    }

    @bindings.add(_Keys.ControlO, _Keys.Any)
    def modify_input(event):
        match event.key_sequence[1]:
            case _KeyPress(key=n) if n.isdigit():
                del_first_words(n, event)
            case _KeyPress(key='o'):
                del_first_words(1, event)
            case _KeyPress(key=_Keys.ControlO):
                del_first_words(1, event, space_at_start=True)
                event.current_buffer.cursor_position = 0
            case _KeyPress(key=n) if n in mods:
                pre, post = mods[n]
                event.current_buffer.text = f'{pre}{event.current_buffer.text}{post}'
                event.current_buffer.cursor_right(len(pre))