diora-web/tui/diora_tui/marks.py

"""Highlights in the terminal text: from the web's anchors to spans of the displayed text, and
back from a selection to an anchor.

An anchor addresses characters of a block's *raw* text (every text node, whitespace as in the
file); the TUI shows the block with whitespace collapsed (blocks.collapse_map relates the two).
The raw text is read from the file for the blocks involved only, a chapter at a time.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path

from rich.style import Style

from . import annotations, blocks

SELECTION_STYLE = Style(reverse=True)
NOTE_STYLE = Style(underline=True, dim=False)
_WORD_RE = re.compile(r"\S+")


def highlight_style(color: str | None) -> Style:
    return Style(color="black", bgcolor=annotations.COLORS.get(color or "yellow", annotations.COLORS["yellow"]))


# ---- words (the cursor moves by word)

def word_starts(text: str) -> list[int]:
    return [m.start() for m in _WORD_RE.finditer(text)]


def word_at(text: str, pos: int) -> tuple[int, int] | None:
    """(start, end) of the word at or after `pos`, or None if there is none."""
    for m in _WORD_RE.finditer(text):
        if m.end() > pos:
            return m.start(), m.end()
    return None


def next_word(text: str, pos: int) -> int | None:
    return next((s for s in word_starts(text) if s > pos), None)


def prev_word(text: str, pos: int) -> int | None:
    return next((s for s in reversed(word_starts(text)) if s < pos), None)


@dataclass
class Maps:
    raw: str
    norm_of_raw: list[int]
    raw_of_norm: list[int]
    length: int


class MarkIndex:
    """Relates a book's displayed blocks to the raw text anchors count in."""

    def __init__(self, book: blocks.FlatBook) -> None:
        self.book = book
        self._raws: dict[int, list[str]] = {}  # chapter -> raw text of its blocks
        self._maps: dict[int, Maps | None] = {}

    def maps(self, block_index: int) -> Maps | None:
        """None when the displayed text isn't the plain collapse of the raw text (a footnote
        marker with spaces inside): such a block can't be addressed exactly."""
        if block_index in self._maps:
            return self._maps[block_index]
        b = self.book.blocks[block_index]
        chapter = b.chapter_index
        if chapter not in self._raws:
            self._raws[chapter] = blocks.raw_block_texts(self.book.path, chapter)
        first = self.book.chapter_start_block[chapter]
        raws = self._raws[chapter]
        result = None
        if 0 <= block_index - first < len(raws):
            raw = raws[block_index - first]
            norm_of_raw, raw_of_norm, length = blocks.collapse_map(raw)
            if length == len(b.text) and "".join(raw[i] if not blocks._JS_SPACE_RE.fullmatch(raw[i]) else " " for i in raw_of_norm) == b.text:
                result = Maps(raw, norm_of_raw, raw_of_norm, length)
        self._maps[block_index] = result
        return result

    # ---- annotation -> spans

    def spans(self, items: list[dict]) -> dict[int, list[tuple[int, int, Style]]]:
        """Per block the (start, end, style) spans of the highlights and notes."""
        out: dict[int, list[tuple[int, int, Style]]] = {}
        n = len(self.book.blocks)
        for h in items:
            a = h.get("anchor") or {}
            if "startBlockIndex" not in a:  # an old anchor (XPath): the web falls back to the quote, so do we
                found = self._find_quote(a)
                if found is not None:
                    bi, s0, e0 = found
                    out.setdefault(bi, []).append((s0, e0, NOTE_STYLE if h.get("type") == "note" else highlight_style(h.get("color"))))
                continue
            try:
                sb = int(a["startBlockIndex"])
                eb = int(a.get("endBlockIndex", sb))
                sc, ec = int(a.get("startChar", 0)), int(a.get("endChar", 0))
            except (KeyError, TypeError, ValueError):
                continue
            if not (0 <= sb < n and sb <= eb < n):
                continue
            style = NOTE_STYLE if h.get("type") == "note" else highlight_style(h.get("color"))
            for bi, s, e in self._ranges(sb, sc, eb, ec, a.get("quote") or ""):
                if e > s:
                    out.setdefault(bi, []).append((s, e, style))
        return out

    def _find_quote(self, anchor: dict) -> tuple[int, int, int] | None:
        """(block, start, end) of the first place in the anchor's chapter that holds its quote."""
        needle = " ".join((anchor.get("quote") or "").split())
        src = anchor.get("chapterSrc")
        if not needle or src not in self.book.chapter_srcs:
            return None
        chapter = self.book.chapter_srcs.index(src)
        first = self.book.chapter_start_block[chapter]
        last = self.book.chapter_start_block[chapter + 1] if chapter + 1 < len(self.book.chapter_start_block) else len(self.book.blocks)
        for bi in range(first, last):
            i = self.book.blocks[bi].text.find(needle)
            if i >= 0:
                return bi, i, i + len(needle)
        return None

    def _ranges(self, sb: int, sc: int, eb: int, ec: int, quote: str):
        text = self.book.blocks[sb].text
        m = self.maps(sb)
        if m is None or (eb != sb and self.maps(eb) is None):
            # not addressable by offsets: find the quote
            i = text.find(" ".join(quote.split())) if quote else -1
            if i >= 0:
                yield sb, i, i + len(" ".join(quote.split()))
            return
        s = m.norm_of_raw[min(sc, len(m.raw))]
        if sb == eb:
            yield sb, s, m.norm_of_raw[min(ec, len(m.raw))]
            return
        yield sb, s, m.length
        for bi in range(sb + 1, eb):
            yield bi, 0, len(self.book.blocks[bi].text)
        me = self.maps(eb)
        yield eb, 0, me.norm_of_raw[min(ec, len(me.raw))]

    # ---- selection -> anchor

    def anchor(self, start: tuple[int, int], end: tuple[int, int]) -> dict | None:
        """The web anchor for the displayed range start..end ((block, char) each, end exclusive),
        or None when a block can't be addressed exactly."""
        (sb, sc), (eb, ec) = start, end
        ms, me = self.maps(sb), self.maps(eb)
        if ms is None or me is None:
            return None
        s_raw = ms.raw_of_norm[sc] if sc < ms.length else len(ms.raw)
        e_raw = me.raw_of_norm[ec - 1] + 1 if ec > 0 else 0
        if sb == eb:
            quote = ms.raw[s_raw:e_raw]
        else:
            quote = ms.raw[s_raw:] + me.raw[:e_raw]
        chapter_src = self.book.chapter_srcs[self.book.blocks[sb].chapter_index] if self.book.chapter_srcs else ""
        return annotations.make_anchor(chapter_src, sb, s_raw, eb, e_raw, quote)