babi/babi.py

import argparse
import collections
import contextlib
import curses
import enum
import functools
import hashlib
import io
import os
import re
import signal
import sys
from typing import Any
from typing import Callable
from typing import cast
from typing import Dict
from typing import FrozenSet
from typing import Generator
from typing import IO
from typing import Iterator
from typing import List
from typing import Match
from typing import NamedTuple
from typing import Optional
from typing import Pattern
from typing import Tuple
from typing import TYPE_CHECKING
from typing import TypeVar
from typing import Union

if TYPE_CHECKING:
    from typing import Protocol  # python3.8+
else:
    Protocol = object

VERSION_STR = 'babi v0'
TCallable = TypeVar('TCallable', bound=Callable[..., Any])
EditResult = enum.Enum('EditResult', 'EXIT NEXT PREV')
PromptResult = enum.Enum('PromptResult', 'CANCELLED')


def _line_x(x: int, width: int) -> int:
    margin = min(width - 3, 6)
    if x + 1 < width:
        return 0
    elif width == 1:
        return x
    else:
        return (
            width - margin - 2 +
            (x + 1 - width) //
            (width - margin - 2) *
            (width - margin - 2)
        )


def _scrolled_line(s: str, x: int, width: int) -> str:
    line_x = _line_x(x, width)
    if line_x:
        s = f'«{s[line_x + 1:]}'
        if line_x and len(s) > width:
            return f'{s[:width - 1]}»'
        else:
            return s.ljust(width)
    elif len(s) > width:
        return f'{s[:width - 1]}»'
    else:
        return s.ljust(width)


class MutableSequenceNoSlice(Protocol):
    def __len__(self) -> int: ...
    def __getitem__(self, idx: int) -> str: ...
    def __setitem__(self, idx: int, val: str) -> None: ...
    def __delitem__(self, idx: int) -> None: ...
    def insert(self, idx: int, val: str) -> None: ...

    def __iter__(self) -> Iterator[str]:
        for i in range(len(self)):
            yield self[i]

    def append(self, val: str) -> None:
        self.insert(len(self), val)

    def pop(self, idx: int = -1) -> str:
        victim = self[idx]
        del self[idx]
        return victim


def _del(lst: MutableSequenceNoSlice, *, idx: int) -> None:
    del lst[idx]


def _set(lst: MutableSequenceNoSlice, *, idx: int, val: str) -> None:
    lst[idx] = val


def _ins(lst: MutableSequenceNoSlice, *, idx: int, val: str) -> None:
    lst.insert(idx, val)


class ListSpy(MutableSequenceNoSlice):
    def __init__(self, lst: MutableSequenceNoSlice) -> None:
        self._lst = lst
        self._undo: List[Callable[[MutableSequenceNoSlice], None]] = []

    def __repr__(self) -> str:
        return f'{type(self).__name__}({self._lst})'

    def __len__(self) -> int:
        return len(self._lst)

    def __getitem__(self, idx: int) -> str:
        return self._lst[idx]

    def __setitem__(self, idx: int, val: str) -> None:
        self._undo.append(functools.partial(_set, idx=idx, val=self._lst[idx]))
        self._lst[idx] = val

    def __delitem__(self, idx: int) -> None:
        if idx < 0:
            idx %= len(self)
        self._undo.append(functools.partial(_ins, idx=idx, val=self._lst[idx]))
        del self._lst[idx]

    def insert(self, idx: int, val: str) -> None:
        if idx < 0:
            idx %= len(self)
        self._undo.append(functools.partial(_del, idx=idx))
        self._lst.insert(idx, val)

    def undo(self, lst: MutableSequenceNoSlice) -> None:
        for fn in reversed(self._undo):
            fn(lst)

    @property
    def has_modifications(self) -> bool:
        return bool(self._undo)


class Key(NamedTuple):
    wch: Union[int, str]
    keyname: bytes


class Margin(NamedTuple):
    header: bool
    footer: bool

    @property
    def body_lines(self) -> int:
        return curses.LINES - self.header - self.footer

    @property
    def page_size(self) -> int:
        if self.body_lines <= 2:
            return 1
        else:
            return self.body_lines - 2

    @classmethod
    def from_current_screen(cls) -> 'Margin':
        if curses.LINES == 1:
            return cls(header=False, footer=False)
        elif curses.LINES == 2:
            return cls(header=False, footer=True)
        else:
            return cls(header=True, footer=True)


def _get_color_pair_mapping() -> Dict[Tuple[int, int], int]:
    ret = {}
    i = 0
    for bg in range(-1, 16):
        for fg in range(bg, 16):
            ret[(fg, bg)] = i
            i += 1
    return ret


COLORS = _get_color_pair_mapping()
del _get_color_pair_mapping


def _has_colors() -> bool:
    return curses.has_colors and curses.COLORS >= 16


def _color(fg: int, bg: int) -> int:
    if _has_colors():
        if bg > fg:
            return curses.A_REVERSE | curses.color_pair(COLORS[(bg, fg)])
        else:
            return curses.color_pair(COLORS[(fg, bg)])
    else:
        if bg > fg:
            return curses.A_REVERSE | curses.color_pair(0)
        else:
            return curses.color_pair(0)


def _init_colors(stdscr: 'curses._CursesWindow') -> None:
    curses.use_default_colors()
    if not _has_colors():
        return
    for (fg, bg), pair in COLORS.items():
        if pair == 0:  # cannot reset pair 0
            continue
        curses.init_pair(pair, fg, bg)


class Status:
    def __init__(self) -> None:
        self._status = ''
        self._action_counter = -1

    def update(self, status: str) -> None:
        self._status = status
        self._action_counter = 25

    def clear(self) -> None:
        self._status = ''

    def draw(self, stdscr: 'curses._CursesWindow', margin: Margin) -> None:
        if margin.footer or self._status:
            stdscr.insstr(curses.LINES - 1, 0, ' ' * curses.COLS)
            if self._status:
                status = f' {self._status} '
                x = (curses.COLS - len(status)) // 2
                if x < 0:
                    x = 0
                    status = status.strip()
                stdscr.insstr(curses.LINES - 1, x, status, curses.A_REVERSE)

    def tick(self, margin: Margin) -> None:
        # when the window is only 1-tall, hide the status quicker
        if margin.footer:
            self._action_counter -= 1
        else:
            self._action_counter -= 24
        if self._action_counter < 0:
            self.clear()

    def cancelled(self) -> PromptResult:
        self.update('cancelled')
        return PromptResult.CANCELLED


class Prompt:
    def __init__(self, screen: 'Screen', prompt: str, lst: List[str]) -> None:
        self._screen = screen
        self._prompt = prompt
        self._lst = lst
        self._y = len(lst) - 1
        self._x = 0

    @property
    def _s(self) -> str:
        return self._lst[self._y]

    @_s.setter
    def _s(self, s: str) -> None:
        self._lst[self._y] = s

    def _render_prompt(self, *, base: Optional[str] = None) -> None:
        base = base or self._prompt
        if not base or curses.COLS < 7:
            prompt_s = ''
        elif len(base) > curses.COLS - 6:
            prompt_s = f'{base[:curses.COLS - 7]}…: '
        else:
            prompt_s = f'{base}: '
        width = curses.COLS - len(prompt_s)
        line = _scrolled_line(self._s, self._x, width)
        cmd = f'{prompt_s}{line}'
        self._screen.stdscr.insstr(curses.LINES - 1, 0, cmd, curses.A_REVERSE)
        x = len(prompt_s) + self._x - _line_x(self._x, width)
        self._screen.stdscr.move(curses.LINES - 1, x)

    def _up(self) -> None:
        self._y = max(0, self._y - 1)
        self._x = len(self._s)

    def _down(self) -> None:
        self._y = min(len(self._lst) - 1, self._y + 1)
        self._x = len(self._s)

    def _right(self) -> None:
        self._x = min(len(self._s), self._x + 1)

    def _left(self) -> None:
        self._x = max(0, self._x - 1)

    def _home(self) -> None:
        self._x = 0

    def _end(self) -> None:
        self._x = len(self._s)

    def _ctrl_left(self) -> None:
        if self._x <= 1:
            self._x = 0
        else:
            self._x -= 1
            tp = self._s[self._x - 1].isalnum()
            while self._x > 0 and tp == self._s[self._x - 1].isalnum():
                self._x -= 1

    def _ctrl_right(self) -> None:
        if self._x >= len(self._s) - 1:
            self._x = len(self._s)
        else:
            self._x += 1
            tp = self._s[self._x].isalnum()
            while self._x < len(self._s) and tp == self._s[self._x].isalnum():
                self._x += 1

    def _backspace(self) -> None:
        if self._x > 0:
            self._s = self._s[:self._x - 1] + self._s[self._x:]
            self._x -= 1

    def _delete(self) -> None:
        if self._x < len(self._s):
            self._s = self._s[:self._x] + self._s[self._x + 1:]

    def _cut_to_end(self) -> None:
        self._s = self._s[:self._x]

    def _resize(self) -> None:
        self._screen.resize()

    def _reverse_search(self) -> Union[None, str, PromptResult]:
        reverse_s = ''
        reverse_idx = self._y
        while True:
            reverse_failed = False
            for search_idx in range(reverse_idx, -1, -1):
                if reverse_s in self._lst[search_idx]:
                    reverse_idx = self._y = search_idx
                    self._x = len(self._s)
                    break
            else:
                reverse_failed = True

            if reverse_failed:
                base = f'{self._prompt}(failed reverse-search)`{reverse_s}`'
            else:
                base = f'{self._prompt}(reverse-search)`{reverse_s}`'

            self._render_prompt(base=base)
            key = _get_char(self._screen.stdscr)

            if key.keyname == b'KEY_RESIZE':
                self._screen.resize()
            elif key.keyname == b'KEY_BACKSPACE':
                reverse_s = reverse_s[:-1]
            elif isinstance(key.wch, str) and key.wch.isprintable():
                reverse_s += key.wch
            elif key.keyname == b'^R':
                reverse_idx = max(0, reverse_idx - 1)
            elif key.keyname == b'^C':
                return self._screen.status.cancelled()
            elif key.keyname == b'^M':
                return self._s
            else:
                return None

    def _cancel(self) -> PromptResult:
        return self._screen.status.cancelled()

    def _submit(self) -> str:
        return self._s

    DISPATCH = {
        # movement
        b'KEY_UP': _up,
        b'KEY_DOWN': _down,
        b'KEY_RIGHT': _right,
        b'KEY_LEFT': _left,
        b'KEY_HOME': _home,
        b'^A': _home,
        b'KEY_END': _end,
        b'^E': _end,
        b'kRIT5': _ctrl_right,
        b'kLFT5': _ctrl_left,
        # editing
        b'KEY_BACKSPACE': _backspace,
        b'KEY_DC': _delete,
        b'^K': _cut_to_end,
        # misc
        b'KEY_RESIZE': _resize,
        b'^R': _reverse_search,
        b'^M': _submit,
        b'^C': _cancel,
    }

    def _c(self, c: str) -> None:
        self._s = self._s[:self._x] + c + self._s[self._x:]
        self._x += 1

    def run(self) -> Union[PromptResult, str]:
        while True:
            self._render_prompt()
            key = _get_char(self._screen.stdscr)

            if key.keyname in Prompt.DISPATCH:
                ret = Prompt.DISPATCH[key.keyname](self)
                if ret is not None:
                    return ret
            elif isinstance(key.wch, str) and key.wch.isprintable():
                self._c(key.wch)


class History:
    def __init__(self) -> None:
        self._orig_len: Dict[str, int] = collections.defaultdict(int)
        self.data: Dict[str, List[str]] = collections.defaultdict(list)
        self.prev: Dict[str, str] = {}

    @contextlib.contextmanager
    def save(self) -> Generator[None, None, None]:
        history_dir = os.path.join(
            os.environ.get('XDG_DATA_HOME') or
            os.path.expanduser('~/.local/share'),
            'babi/history',
        )
        os.makedirs(history_dir, exist_ok=True)
        for filename in os.listdir(history_dir):
            with open(os.path.join(history_dir, filename)) as f:
                self.data[filename] = f.read().splitlines()
                self._orig_len[filename] = len(self.data[filename])
        try:
            yield
        finally:
            for k, v in self.data.items():
                new_history = v[self._orig_len[k]:]
                if new_history:
                    with open(os.path.join(history_dir, k), 'a+') as f:
                        f.write('\n'.join(new_history) + '\n')


def _restore_lines_eof_invariant(lines: MutableSequenceNoSlice) -> None:
    """The file lines will always contain a blank empty string at the end to
    simplify rendering.  This should be called whenever the end of the file
    might change.
    """
    if not lines or lines[-1] != '':
        lines.append('')


def _get_lines(sio: IO[str]) -> Tuple[List[str], str, bool, str]:
    sha256 = hashlib.sha256()
    lines = []
    newlines = collections.Counter({'\n': 0})  # default to `\n`
    for line in sio:
        sha256.update(line.encode())
        for ending in ('\r\n', '\n'):
            if line.endswith(ending):
                lines.append(line[:-1 * len(ending)])
                newlines[ending] += 1
                break
        else:
            lines.append(line)
    _restore_lines_eof_invariant(lines)
    (nl, _), = newlines.most_common(1)
    mixed = len({k for k, v in newlines.items() if v}) > 1
    return lines, nl, mixed, sha256.hexdigest()


class Action:
    def __init__(
            self, *, name: str, spy: ListSpy,
            start_x: int, start_y: int, start_modified: bool,
            end_x: int, end_y: int, end_modified: bool,
            final: bool,
    ):
        self.name = name
        self.spy = spy
        self.start_x = start_x
        self.start_y = start_y
        self.start_modified = start_modified
        self.end_x = end_x
        self.end_y = end_y
        self.end_modified = end_modified
        self.final = final

    def apply(self, file: 'File') -> 'Action':
        spy = ListSpy(file.lines)
        action = Action(
            name=self.name, spy=spy,
            start_x=self.end_x, start_y=self.end_y,
            start_modified=self.end_modified,
            end_x=self.start_x, end_y=self.start_y,
            end_modified=self.start_modified,
            final=True,
        )

        self.spy.undo(spy)
        file.x = self.start_x
        file.y = self.start_y
        file.modified = self.start_modified

        return action


def action(func: TCallable) -> TCallable:
    @functools.wraps(func)
    def action_inner(self: 'File', *args: Any, **kwargs: Any) -> Any:
        assert not isinstance(self.lines, ListSpy), 'nested edit/movement'
        self.mark_previous_action_as_final()
        return func(self, *args, **kwargs)
    return cast(TCallable, action_inner)


def edit_action(
        name: str,
        *,
        final: bool,
) -> Callable[[TCallable], TCallable]:
    def edit_action_decorator(func: TCallable) -> TCallable:
        @functools.wraps(func)
        def edit_action_inner(self: 'File', *args: Any, **kwargs: Any) -> Any:
            with self.edit_action_context(name, final=final):
                return func(self, *args, **kwargs)
        return cast(TCallable, edit_action_inner)
    return edit_action_decorator


class Found(NamedTuple):
    y: int
    match: Match[str]


class _SearchIter:
    def __init__(
            self,
            file: 'File',
            reg: Pattern[str],
            *,
            offset: int,
    ) -> None:
        self.file = file
        self.reg = reg
        self.offset = offset
        self.wrapped = False
        self._start_x = file.x + offset
        self._start_y = file.y

    def __iter__(self) -> '_SearchIter':
        return self

    def _stop_if_past_original(self, y: int, match: Match[str]) -> Found:
        if (
                self.wrapped and (
                    y > self._start_y or
                    y == self._start_y and match.start() >= self._start_x
                )
        ):
            raise StopIteration()
        return Found(y, match)

    def __next__(self) -> Tuple[int, Match[str]]:
        x = self.file.x + self.offset
        y = self.file.y

        match = self.reg.search(self.file.lines[y], x)
        if match:
            return self._stop_if_past_original(y, match)

        if self.wrapped:
            for line_y in range(y + 1, self._start_y + 1):
                match = self.reg.search(self.file.lines[line_y])
                if match:
                    return self._stop_if_past_original(line_y, match)
        else:
            for line_y in range(y + 1, len(self.file.lines)):
                match = self.reg.search(self.file.lines[line_y])
                if match:
                    return self._stop_if_past_original(line_y, match)

            self.wrapped = True

            for line_y in range(0, self._start_y + 1):
                match = self.reg.search(self.file.lines[line_y])
                if match:
                    return self._stop_if_past_original(line_y, match)

        raise StopIteration()


class File:
    def __init__(self, filename: Optional[str]) -> None:
        self.filename = filename
        self.modified = False
        self.lines: MutableSequenceNoSlice = []
        self.nl = '\n'
        self.file_y = self.y = self.x = self.x_hint = 0
        self.sha256: Optional[str] = None
        self.undo_stack: List[Action] = []
        self.redo_stack: List[Action] = []

    def ensure_loaded(self, status: Status) -> None:
        if self.lines:
            return

        if self.filename is not None and os.path.isfile(self.filename):
            with open(self.filename, newline='') as f:
                self.lines, self.nl, mixed, self.sha256 = _get_lines(f)
        else:
            if self.filename is not None:
                if os.path.lexists(self.filename):
                    status.update(f'{self.filename!r} is not a file')
                    self.filename = None
                else:
                    status.update('(new file)')
            sio = io.StringIO('')
            self.lines, self.nl, mixed, self.sha256 = _get_lines(sio)

        if mixed:
            status.update(f'mixed newlines will be converted to {self.nl!r}')
            self.modified = True

    def __repr__(self) -> str:
        attrs = ',\n    '.join(f'{k}={v!r}' for k, v in self.__dict__.items())
        return f'{type(self).__name__}(\n    {attrs},\n)'

    # movement

    def scroll_screen_if_needed(self, margin: Margin) -> None:
        # if the `y` is not on screen, make it so
        if self.file_y <= self.y < self.file_y + margin.body_lines:
            return

        self.file_y = max(self.y - margin.body_lines // 2, 0)

    def _scroll_amount(self) -> int:
        return int(curses.LINES / 2 + .5)

    def _set_x_after_vertical_movement(self) -> None:
        self.x = min(len(self.lines[self.y]), self.x_hint)

    def _increment_y(self, margin: Margin) -> None:
        self.y += 1
        if self.y >= self.file_y + margin.body_lines:
            self.file_y += self._scroll_amount()

    def _decrement_y(self, margin: Margin) -> None:
        self.y -= 1
        if self.y < self.file_y:
            self.file_y -= self._scroll_amount()
            self.file_y = max(self.file_y, 0)

    @action
    def up(self, margin: Margin) -> None:
        if self.y > 0:
            self._decrement_y(margin)
            self._set_x_after_vertical_movement()

    @action
    def down(self, margin: Margin) -> None:
        if self.y < len(self.lines) - 1:
            self._increment_y(margin)
            self._set_x_after_vertical_movement()

    @action
    def right(self, margin: Margin) -> None:
        if self.x >= len(self.lines[self.y]):
            if self.y < len(self.lines) - 1:
                self.x = 0
                self._increment_y(margin)
        else:
            self.x += 1
        self.x_hint = self.x

    @action
    def left(self, margin: Margin) -> None:
        if self.x == 0:
            if self.y > 0:
                self._decrement_y(margin)
                self.x = len(self.lines[self.y])
        else:
            self.x -= 1
        self.x_hint = self.x

    @action
    def home(self, margin: Margin) -> None:
        self.x = self.x_hint = 0

    @action
    def end(self, margin: Margin) -> None:
        self.x = self.x_hint = len(self.lines[self.y])

    @action
    def ctrl_up(self, margin: Margin) -> None:
        self.file_y = max(0, self.file_y - 1)
        self.y = min(self.y, self.file_y + margin.body_lines - 1)
        self._set_x_after_vertical_movement()

    @action
    def ctrl_down(self, margin: Margin) -> None:
        self.file_y = min(len(self.lines) - 1, self.file_y + 1)
        self.y = max(self.y, self.file_y)
        self._set_x_after_vertical_movement()

    @action
    def ctrl_right(self, margin: Margin) -> None:
        line = self.lines[self.y]
        # if we're at the second to last character, jump to end of line
        if self.x == len(line) - 1:
            self.x = self.x_hint = self.x + 1
        # if we're at the end of the line, jump forward to the next non-ws
        elif self.x == len(line):
            while (
                    self.y < len(self.lines) - 1 and (
                        self.x == len(self.lines[self.y]) or
                        self.lines[self.y][self.x].isspace()
                    )
            ):
                if self.x == len(self.lines[self.y]):
                    self._increment_y(margin)
                    self.x = self.x_hint = 0
                else:
                    self.x = self.x_hint = self.x + 1
        # if we're inside the line, jump to next position that's not our type
        else:
            self.x = self.x_hint = self.x + 1
            tp = line[self.x].isalnum()
            while self.x < len(line) and tp == line[self.x].isalnum():
                self.x = self.x_hint = self.x + 1

    @action
    def ctrl_left(self, margin: Margin) -> None:
        line = self.lines[self.y]
        # if we're at position 1 and it's not a space, go to the beginning
        if self.x == 1 and not line[:self.x].isspace():
            self.x = self.x_hint = 0
        # if we're at the beginning or it's all space up to here jump to the
        # end of the previous non-space line
        elif self.x == 0 or line[:self.x].isspace():
            self.x = self.x_hint = 0
            while self.y > 0 and (self.x == 0 or not self.lines[self.y]):
                self._decrement_y(margin)
                self.x = self.x_hint = len(self.lines[self.y])
        else:
            self.x = self.x_hint = self.x - 1
            tp = line[self.x - 1].isalnum()
            while self.x > 0 and tp == line[self.x - 1].isalnum():
                self.x = self.x_hint = self.x - 1

    @action
    def ctrl_home(self, margin: Margin) -> None:
        self.x = self.x_hint = 0
        self.y = self.file_y = 0

    @action
    def ctrl_end(self, margin: Margin) -> None:
        self.x = self.x_hint = 0
        self.y = len(self.lines) - 1
        self.scroll_screen_if_needed(margin)

    @action
    def go_to_line(self, lineno: int, margin: Margin) -> None:
        self.x = self.x_hint = 0
        if lineno == 0:
            self.y = 0
        elif lineno > len(self.lines):
            self.y = len(self.lines) - 1
        elif lineno < 0:
            self.y = max(0, lineno + len(self.lines))
        else:
            self.y = lineno - 1
        self.scroll_screen_if_needed(margin)

    @action
    def search(
            self,
            reg: Pattern[str],
            status: Status,
            margin: Margin,
    ) -> None:
        search = _SearchIter(self, reg, offset=1)
        try:
            line_y, match = next(iter(search))
        except StopIteration:
            status.update('no matches')
        else:
            if line_y == self.y and match.start() == self.x:
                status.update('this is the only occurrence')
            else:
                if search.wrapped:
                    status.update('search wrapped')
                self.y = line_y
                self.x = self.x_hint = match.start()
                self.scroll_screen_if_needed(margin)

    def replace(
            self,
            screen: 'Screen',
            reg: Pattern[str],
            replace: str,
    ) -> None:
        self.mark_previous_action_as_final()

        def highlight() -> None:
            y = screen.file.rendered_y(screen.margin)
            x = screen.file.rendered_x()
            maxlen = curses.COLS - x
            s = match[0]
            if len(s) >= maxlen:
                s = _scrolled_line(match[0], 0, maxlen)
            # TODO: probably should do this without `len(s)`
            screen.stdscr.chgat(y, x, len(s), curses.A_REVERSE)

        count = 0
        res: Union[str, PromptResult] = ''
        search = _SearchIter(self, reg, offset=0)
        for line_y, match in search:
            self.y = line_y
            self.x = self.x_hint = match.start()
            self.scroll_screen_if_needed(screen.margin)
            if res != 'a':  # make `a` replace the rest of them
                screen.draw()
                highlight()
                with screen.resize_cb(highlight):
                    res = screen.quick_prompt(
                        'replace [y(es), n(o), a(ll)]?', frozenset('yna'),
                    )
            if res in {'y', 'a'}:
                count += 1
                with self.edit_action_context('replace', final=True):
                    replaced = match.expand(replace)
                    line = screen.file.lines[line_y]
                    line = line[:match.start()] + replaced + line[match.end():]
                    screen.file.lines[line_y] = line
                    screen.file.modified = True
                search.offset = len(replaced)
            elif res == 'n':
                search.offset = 1
            else:
                assert res is PromptResult.CANCELLED
                return

        if res == '':  # we never went through the loop
            screen.status.update('no matches')
        else:
            occurrences = 'occurrence' if count == 1 else 'occurrences'
            screen.status.update(f'replaced {count} {occurrences}')

    @action
    def page_up(self, margin: Margin) -> None:
        if self.y < margin.body_lines:
            self.y = self.file_y = 0
        else:
            pos = max(self.file_y - margin.page_size, 0)
            self.y = self.file_y = pos
        self._set_x_after_vertical_movement()

    @action
    def page_down(self, margin: Margin) -> None:
        if self.file_y + margin.body_lines >= len(self.lines):
            self.y = len(self.lines) - 1
        else:
            pos = self.file_y + margin.page_size
            self.y = self.file_y = pos
        self._set_x_after_vertical_movement()

    # editing

    @edit_action('backspace text', final=False)
    def backspace(self, margin: Margin) -> None:
        # backspace at the beginning of the file does nothing
        if self.y == 0 and self.x == 0:
            pass
        # at the beginning of the line, we join the current line and
        # the previous line
        elif self.x == 0:
            victim = self.lines.pop(self.y)
            new_x = len(self.lines[self.y - 1])
            self.lines[self.y - 1] += victim
            self._decrement_y(margin)
            self.x = self.x_hint = new_x
            # deleting the fake end-of-file doesn't cause modification
            self.modified |= self.y < len(self.lines) - 1
            _restore_lines_eof_invariant(self.lines)
        else:
            s = self.lines[self.y]
            self.lines[self.y] = s[:self.x - 1] + s[self.x:]
            self.x = self.x_hint = self.x - 1
            self.modified = True

    @edit_action('delete text', final=False)
    def delete(self, margin: Margin) -> None:
        # noop at end of the file
        if self.y == len(self.lines) - 1:
            pass
        # if we're at the end of the line, collapse the line afterwards
        elif self.x == len(self.lines[self.y]):
            victim = self.lines.pop(self.y + 1)
            self.lines[self.y] += victim
            self.modified = True
        else:
            s = self.lines[self.y]
            self.lines[self.y] = s[:self.x] + s[self.x + 1:]
            self.modified = True

    @edit_action('line break', final=False)
    def enter(self, margin: Margin) -> None:
        s = self.lines[self.y]
        self.lines[self.y] = s[:self.x]
        self.lines.insert(self.y + 1, s[self.x:])
        self._increment_y(margin)
        self.x = self.x_hint = 0
        self.modified = True

    @edit_action('cut', final=False)
    def cut(self, cut_buffer: Tuple[str, ...]) -> Tuple[str, ...]:
        if self.y == len(self.lines) - 1:
            return ()
        else:
            victim = self.lines.pop(self.y)
            self.x = self.x_hint = 0
            self.modified = True
            return cut_buffer + (victim,)

    @edit_action('uncut', final=True)
    def uncut(self, cut_buffer: Tuple[str, ...], margin: Margin) -> None:
        for cut_line in cut_buffer:
            line = self.lines[self.y]
            before, after = line[:self.x], line[self.x:]
            self.lines[self.y] = before + cut_line
            self.lines.insert(self.y + 1, after)
            self._increment_y(margin)
            self.x = self.x_hint = 0

    DISPATCH = {
        # movement
        b'KEY_UP': up,
        b'KEY_DOWN': down,
        b'KEY_RIGHT': right,
        b'KEY_LEFT': left,
        b'KEY_HOME': home,
        b'^A': home,
        b'KEY_END': end,
        b'^E': end,
        b'KEY_PPAGE': page_up,
        b'^Y': page_up,
        b'KEY_NPAGE': page_down,
        b'^V': page_down,
        b'kUP5': ctrl_up,
        b'kDN5': ctrl_down,
        b'kRIT5': ctrl_right,
        b'kLFT5': ctrl_left,
        b'kHOM5': ctrl_home,
        b'kEND5': ctrl_end,
        # editing
        b'KEY_BACKSPACE': backspace,
        b'KEY_DC': delete,
        b'^M': enter,
    }

    @edit_action('text', final=False)
    def c(self, wch: str, margin: Margin) -> None:
        s = self.lines[self.y]
        self.lines[self.y] = s[:self.x] + wch + s[self.x:]
        self.x = self.x_hint = self.x + 1
        self.modified = True
        _restore_lines_eof_invariant(self.lines)

    def mark_previous_action_as_final(self) -> None:
        if self.undo_stack:
            self.undo_stack[-1].final = True

    @contextlib.contextmanager
    def edit_action_context(
            self, name: str,
            *,
            final: bool,
    ) -> Generator[None, None, None]:
        continue_last = (
            self.undo_stack and
            self.undo_stack[-1].name == name and
            not self.undo_stack[-1].final
        )
        if continue_last:
            spy = self.undo_stack[-1].spy
        else:
            if self.undo_stack:
                self.undo_stack[-1].final = True
            spy = ListSpy(self.lines)

        before_x, before_line = self.x, self.y
        before_modified = self.modified
        assert not isinstance(self.lines, ListSpy), 'recursive action?'
        orig, self.lines = self.lines, spy
        try:
            yield
        finally:
            self.lines = orig
            self.redo_stack.clear()
            if continue_last:
                self.undo_stack[-1].end_x = self.x
                self.undo_stack[-1].end_y = self.y
                self.undo_stack[-1].end_modified = self.modified
            elif spy.has_modifications:
                action = Action(
                    name=name, spy=spy,
                    start_x=before_x, start_y=before_line,
                    start_modified=before_modified,
                    end_x=self.x, end_y=self.y,
                    end_modified=self.modified,
                    final=final,
                )
                self.undo_stack.append(action)

    # positioning

    def rendered_y(self, margin: Margin) -> int:
        return self.y - self.file_y + margin.header

    def rendered_x(self) -> int:
        return self.x - _line_x(self.x, curses.COLS)

    def move_cursor(
            self,
            stdscr: 'curses._CursesWindow',
            margin: Margin,
    ) -> None:
        stdscr.move(self.rendered_y(margin), self.rendered_x())

    def draw(self, stdscr: 'curses._CursesWindow', margin: Margin) -> None:
        to_display = min(len(self.lines) - self.file_y, margin.body_lines)
        for i in range(to_display):
            line_idx = self.file_y + i
            line = self.lines[line_idx]
            x = self.x if line_idx == self.y else 0
            line = _scrolled_line(line, x, curses.COLS)
            stdscr.insstr(i + margin.header, 0, line)
        blankline = ' ' * curses.COLS
        for i in range(to_display, margin.body_lines):
            stdscr.insstr(i + margin.header, 0, blankline)


class Screen:
    def __init__(
            self,
            stdscr: 'curses._CursesWindow',
            files: List[File],
    ) -> None:
        self.stdscr = stdscr
        self.files = files
        self.i = 0
        self.history = History()
        self.status = Status()
        self.margin = Margin.from_current_screen()
        self.prevkey = Key('', b'')
        self.cut_buffer: Tuple[str, ...] = ()
        self._resize_cb: Optional[Callable[[], None]] = None

    @property
    def file(self) -> File:
        return self.files[self.i]

    def _draw_header(self) -> None:
        filename = self.file.filename or '<<new file>>'
        if self.file.modified:
            filename += ' *'
        if len(self.files) > 1:
            files = f'[{self.i + 1}/{len(self.files)}] '
            version_width = len(VERSION_STR) + 2 + len(files)
        else:
            files = ''
            version_width = len(VERSION_STR) + 2
        centered = filename.center(curses.COLS)[version_width:]
        s = f' {VERSION_STR} {files}{centered}{files}'
        self.stdscr.insstr(0, 0, s, curses.A_REVERSE)

    def draw(self) -> None:
        if self.margin.header:
            self._draw_header()
        self.file.draw(self.stdscr, self.margin)
        self.status.draw(self.stdscr, self.margin)

    @contextlib.contextmanager
    def resize_cb(self, f: Callable[[], None]) -> Generator[None, None, None]:
        assert self._resize_cb is None, self._resize_cb
        self._resize_cb = f
        try:
            yield
        finally:
            self._resize_cb = None

    def resize(self) -> None:
        curses.update_lines_cols()
        self.margin = Margin.from_current_screen()
        self.file.scroll_screen_if_needed(self.margin)
        self.draw()
        if self._resize_cb is not None:
            self._resize_cb()

    def quick_prompt(
            self,
            prompt: str,
            options: FrozenSet[str],
    ) -> Union[str, PromptResult]:
        while True:
            s = prompt.ljust(curses.COLS)
            if len(s) > curses.COLS:
                s = f'{s[:curses.COLS - 1]}…'
            self.stdscr.insstr(curses.LINES - 1, 0, s, curses.A_REVERSE)
            x = min(curses.COLS - 1, len(prompt) + 1)
            self.stdscr.move(curses.LINES - 1, x)

            key = _get_char(self.stdscr)
            if key.keyname == b'KEY_RESIZE':
                self.resize()
            elif key.keyname == b'^C':
                return self.status.cancelled()
            elif key.wch in options:
                assert isinstance(key.wch, str)  # mypy doesn't know
                return key.wch

    def prompt(
            self,
            prompt: str,
            *,
            allow_empty: bool = False,
            history: Optional[str] = None,
            default_prev: bool = False,
            default: Optional[str] = None,
    ) -> Union[str, PromptResult]:
        default = default or ''
        self.status.clear()
        if history is not None:
            history_data = [*self.history.data[history], default]
            if default_prev and history in self.history.prev:
                prompt = f'{prompt} [{self.history.prev[history]}]'
        else:
            history_data = [default]

        ret = Prompt(self, prompt, history_data).run()

        if ret is not PromptResult.CANCELLED and history is not None:
            if ret:  # only put non-empty things in history
                history_lst = self.history.data[history]
                if not history_lst or history_lst[-1] != ret:
                    history_lst.append(ret)
                    self.history.prev[history] = ret
            elif default_prev and history in self.history.prev:
                return self.history.prev[history]

        if not allow_empty and not ret:
            return self.status.cancelled()
        else:
            return ret

    def go_to_line(self) -> None:
        response = self.prompt('enter line number')
        if response is not PromptResult.CANCELLED:
            try:
                lineno = int(response)
            except ValueError:
                self.status.update(f'not an integer: {response!r}')
            else:
                self.file.go_to_line(lineno, self.margin)

    def current_position(self) -> None:
        line = f'line {self.file.y + 1}'
        col = f'col {self.file.x + 1}'
        line_count = max(len(self.file.lines) - 1, 1)
        lines_word = 'line' if line_count == 1 else 'lines'
        self.status.update(f'{line}, {col} (of {line_count} {lines_word})')

    def cut(self) -> None:
        if self.prevkey.keyname == b'^K':
            cut_buffer = self.cut_buffer
        else:
            cut_buffer = ()
        self.cut_buffer = self.file.cut(cut_buffer)

    def uncut(self) -> None:
        self.file.uncut(self.cut_buffer, self.margin)

    def _get_search_re(self, prompt: str) -> Union[Pattern[str], PromptResult]:
        response = self.prompt(prompt, history='search', default_prev=True)
        if response is PromptResult.CANCELLED:
            return response
        try:
            return re.compile(response)
        except re.error:
            self.status.update(f'invalid regex: {response!r}')
            return PromptResult.CANCELLED

    def _undo_redo(
            self,
            op: str,
            from_stack: List[Action],
            to_stack: List[Action],
    ) -> None:
        if not from_stack:
            self.status.update(f'nothing to {op}!')
        else:
            action = from_stack.pop()
            to_stack.append(action.apply(self.file))
            self.file.scroll_screen_if_needed(self.margin)
            self.status.update(f'{op}: {action.name}')

    def undo(self) -> None:
        self._undo_redo('undo', self.file.undo_stack, self.file.redo_stack)

    def redo(self) -> None:
        self._undo_redo('redo', self.file.redo_stack, self.file.undo_stack)

    def search(self) -> None:
        response = self._get_search_re('search')
        if response is not PromptResult.CANCELLED:
            self.file.search(response, self.status, self.margin)

    def replace(self) -> None:
        search_response = self._get_search_re('search (to replace)')
        if search_response is not PromptResult.CANCELLED:
            response = self.prompt(
                'replace with', history='replace', allow_empty=True,
            )
            if response is not PromptResult.CANCELLED:
                self.file.replace(self, search_response, response)

    def command(self) -> Optional[EditResult]:
        response = self.prompt('', history='command')
        if response == ':q':
            return EditResult.EXIT
        elif response == ':w':
            self.save()
        elif response == ':wq':
            self.save()
            return EditResult.EXIT
        elif response is not PromptResult.CANCELLED:
            self.status.update(f'invalid command: {response}')
        return None

    def save(self) -> Optional[PromptResult]:
        self.file.mark_previous_action_as_final()

        # TODO: make directories if they don't exist
        # TODO: maybe use mtime / stat as a shortcut for hashing below
        # TODO: strip trailing whitespace?
        # TODO: save atomically?
        if self.file.filename is None:
            filename = self.prompt('enter filename')
            if filename is PromptResult.CANCELLED:
                return PromptResult.CANCELLED
            else:
                self.file.filename = filename

        if os.path.isfile(self.file.filename):
            with open(self.file.filename) as f:
                *_, sha256 = _get_lines(f)
        else:
            sha256 = hashlib.sha256(b'').hexdigest()

        contents = self.file.nl.join(self.file.lines)
        sha256_to_save = hashlib.sha256(contents.encode()).hexdigest()

        # the file on disk is the same as when we opened it
        if sha256 not in (self.file.sha256, sha256_to_save):
            self.status.update('(file changed on disk, not implemented)')
            return PromptResult.CANCELLED

        with open(self.file.filename, 'w') as f:
            f.write(contents)

        self.file.modified = False
        self.file.sha256 = sha256_to_save
        num_lines = len(self.file.lines) - 1
        lines = 'lines' if num_lines != 1 else 'line'
        self.status.update(f'saved! ({num_lines} {lines} written)')

        # fix up modified state in undo / redo stacks
        for stack in (self.file.undo_stack, self.file.redo_stack):
            first = True
            for action in reversed(stack):
                action.end_modified = not first
                action.start_modified = True
                first = False
        return None

    def save_filename(self) -> Optional[PromptResult]:
        response = self.prompt('enter filename', default=self.file.filename)
        if response is PromptResult.CANCELLED:
            return PromptResult.CANCELLED
        else:
            self.file.filename = response
            return self.save()

    def quit_save_modified(self) -> Optional[EditResult]:
        if self.file.modified:
            response = self.quick_prompt(
                'file is modified - save [y(es), n(o)]?', frozenset('yn'),
            )
            if response == 'y':
                if self.save_filename() is not PromptResult.CANCELLED:
                    return EditResult.EXIT
                else:
                    return None
            elif response == 'n':
                return EditResult.EXIT
            else:
                assert response is PromptResult.CANCELLED
                return None
        return EditResult.EXIT

    def background(self) -> None:
        curses.endwin()
        os.kill(os.getpid(), signal.SIGSTOP)
        self.stdscr = _init_screen()
        self.resize()

    DISPATCH = {
        b'KEY_RESIZE': resize,
        b'^_': go_to_line,
        b'^C': current_position,
        b'^K': cut,
        b'^U': uncut,
        b'M-u': undo,
        b'M-U': redo,
        b'^W': search,
        b'^\\': replace,
        b'^[': command,
        b'^S': save,
        b'^O': save_filename,
        b'^X': quit_save_modified,
        b'kLFT3': lambda screen: EditResult.PREV,
        b'kRIT3': lambda screen: EditResult.NEXT,
        b'^Z': background,
    }


def _color_test(stdscr: 'curses._CursesWindow') -> None:
    header = f' {VERSION_STR}'
    header += '<< color test >>'.center(curses.COLS)[len(header):]
    stdscr.insstr(0, 0, header, curses.A_REVERSE)

    maxy, maxx = stdscr.getmaxyx()
    if maxy < 19 or maxx < 68:  # pragma: no cover (will be deleted)
        raise SystemExit('--color-test needs a window of at least 68 x 19')

    y = 1
    for fg in range(-1, 16):
        x = 0
        for bg in range(-1, 16):
            if bg > fg:
                s = f'*{COLORS[bg, fg]:3}'
            else:
                s = f' {COLORS[fg, bg]:3}'
            stdscr.addstr(y, x, s, _color(fg, bg))
            x += 4
        y += 1
    stdscr.get_wch()


# TODO: find a place to populate these, surely there's a database somewhere
SEQUENCE_KEYNAME = {
    '\x1bOH': b'KEY_HOME',
    '\x1bOF': b'KEY_END',
    '\x1b[1;3A': b'kUP3',  # M-Up
    '\x1b[1;3B': b'kDN3',  # M-Down
    '\x1b[1;3C': b'kRIT3',  # M-Right
    '\x1b[1;3D': b'kLFT3',  # M-Left
    '\x1b[1;5A': b'kUP5',  # ^Up
    '\x1b[1;5B': b'kDN5',  # ^Down
    '\x1b[1;5C': b'kRIT5',  # ^Right
    '\x1b[1;5D': b'kLFT5',  # ^Left
    '\x1b[1;5H': b'kHOM5',  # ^Home
    '\x1b[1;5F': b'kEND5',  # ^End
}


def _get_char(stdscr: 'curses._CursesWindow') -> Key:
    wch = stdscr.get_wch()
    if isinstance(wch, str) and wch == '\x1b':
        stdscr.nodelay(True)
        try:
            while True:
                try:
                    new_wch = stdscr.get_wch()
                    if isinstance(new_wch, str):
                        wch += new_wch
                    else:  # pragma: no cover (impossible?)
                        curses.unget_wch(new_wch)
                        break
                except curses.error:
                    break
        finally:
            stdscr.nodelay(False)

        if len(wch) == 2:
            return Key(wch, f'M-{wch[1]}'.encode())
        elif len(wch) > 1:
            keyname = SEQUENCE_KEYNAME.get(wch, b'unknown')
            return Key(wch, keyname)
    elif wch == '\x7f':  # pragma: no cover (macos)
        keyname = curses.keyname(curses.KEY_BACKSPACE)
        return Key(wch, keyname)

    key = wch if isinstance(wch, int) else ord(wch)
    keyname = curses.keyname(key)
    return Key(wch, keyname)


def _edit(screen: Screen) -> EditResult:
    screen.prevkey = Key('', b'')
    screen.file.ensure_loaded(screen.status)

    while True:
        screen.status.tick(screen.margin)
        screen.draw()
        screen.file.move_cursor(screen.stdscr, screen.margin)

        key = _get_char(screen.stdscr)

        if key.keyname in File.DISPATCH:
            File.DISPATCH[key.keyname](screen.file, screen.margin)
        elif key.keyname in Screen.DISPATCH:
            ret = Screen.DISPATCH[key.keyname](screen)
            if isinstance(ret, EditResult):
                return ret
        elif isinstance(key.wch, str) and key.wch.isprintable():
            screen.file.c(key.wch, screen.margin)
        else:
            screen.status.update(f'unknown key: {key}')

        screen.prevkey = key


def c_main(stdscr: 'curses._CursesWindow', args: argparse.Namespace) -> None:
    if args.color_test:
        return _color_test(stdscr)
    screen = Screen(stdscr, [File(f) for f in args.filenames or [None]])
    with screen.history.save():
        while screen.files:
            screen.i = screen.i % len(screen.files)
            res = _edit(screen)
            if res == EditResult.EXIT:
                del screen.files[screen.i]
                screen.status.clear()
            elif res == EditResult.NEXT:
                screen.i += 1
                screen.status.clear()
            elif res == EditResult.PREV:
                screen.i -= 1
                screen.status.clear()
            else:
                raise AssertionError(f'unreachable {res}')


def _init_screen() -> 'curses._CursesWindow':
    # set the escape delay so curses does not pause waiting for sequences
    if sys.version_info >= (3, 9):  # pragma: no cover
        curses.set_escdelay(25)
    else:  # pragma: no cover
        os.environ.setdefault('ESCDELAY', '25')

    stdscr = curses.initscr()
    curses.noecho()
    curses.cbreak()
    # <enter> is not transformed into '\n' so it can be differentiated from ^J
    curses.nonl()
    # ^S / ^Q / ^Z / ^\ are passed through
    curses.raw()
    stdscr.keypad(True)
    with contextlib.suppress(curses.error):
        curses.start_color()
    _init_colors(stdscr)
    return stdscr


@contextlib.contextmanager
def make_stdscr() -> Generator['curses._CursesWindow', None, None]:
    """essentially `curses.wrapper` but split out to implement ^Z"""
    stdscr = _init_screen()
    try:
        yield stdscr
    finally:
        curses.endwin()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument('--color-test', action='store_true')
    parser.add_argument('filenames', metavar='filename', nargs='*')
    args = parser.parse_args()
    with make_stdscr() as stdscr:
        c_main(stdscr, args)
    return 0


if __name__ == '__main__':
    exit(main())