diora-web/tui/diora_tui/annotations.py

"""Highlights, notes and bookmarks of a book, as the web reader keeps them.

Format (static/js/app/72-highlights.js, 68-bookmarks.js): per book and kind one JSON list,
encrypted with the book key (AES-GCM, {ct, iv}), replaced as a whole on the server:

  highlight  {id, type: 'highlight'|'note', anchor, color, note, createdAt}
  anchor     {type: 'epub', chapterSrc, startBlockIndex, startChar, endBlockIndex, endChar, quote}
  bookmark   {id, label, anchor (the chapter's file), scrollFraction, createdAt}  newest first

The server does no merging, so the apps write their whole list. The TUI keeps what it last
saw (`base`), and before it writes it merges its list with the server's current one by id
(three-way: what it added, deleted or changed since `base`, on top of what the server holds), so
a highlight another device made meanwhile is not wiped out. Its edits are queued (`dirty`) while
offline.
"""

from __future__ import annotations

import functools
import json
import os
import threading
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

from . import crypto, paths

COLORS = {"yellow": "#f1c40f", "green": "#2ecc71", "blue": "#3498db", "red": "#e63946"}
COLOR_NAMES = {"yellow": "Gelb", "green": "Grün", "blue": "Blau", "red": "Rot"}
KINDS = ("highlights", "bookmarks")
MAX_BYTES = {"highlights": 700 * 1024, "bookmarks": 100 * 1024}  # the server's size guards (raw bytes)
QUOTE_MAX = 200

_DIR = paths.data_dir() / "annotations"


def _now_iso() -> str:
    """Like the browser's toISOString(): 2026-10-26T12:00:00.123Z."""
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")


# ---- the items

def new_highlight(anchor: dict, color: str = "yellow", note: str = "") -> dict:
    return {"id": str(uuid.uuid4()), "type": "highlight", "anchor": anchor, "color": color, "note": note, "createdAt": _now_iso()}


def new_note(anchor: dict, text: str) -> dict:
    """A freeform note: no mark in the text, its anchor only says where it belongs."""
    return {"id": str(uuid.uuid4()), "type": "note", "anchor": anchor, "color": None, "note": text, "createdAt": _now_iso()}


def new_bookmark(label: str, anchor: str, scroll_fraction: float) -> dict:
    return {"id": str(uuid.uuid4()), "label": label, "anchor": anchor, "scrollFraction": scroll_fraction, "createdAt": _now_iso()}


def make_anchor(chapter_src: str, start_block: int, start_char: int, end_block: int, end_char: int, quote: str) -> dict:
    return {"type": "epub", "chapterSrc": chapter_src, "startBlockIndex": start_block, "startChar": start_char,
            "endBlockIndex": end_block, "endChar": end_char, "quote": quote[:QUOTE_MAX]}


def sort_key(h: dict) -> tuple:
    a = h.get("anchor") or {}
    return (a.get("startBlockIndex", 10**9), a.get("startChar", 0))


# ---- merge

def merge(base: list[dict] | None, ours: list[dict], theirs: list[dict], *, prepend_added: bool = False) -> list[dict]:
    """The server's list (`theirs`) with our changes since `base` laid on it, by id: what we
    deleted is gone, what we added is in, what we changed is ours (also against their change).
    `base` None: we never saw the server's list, so nothing counts as deleted."""
    base_by = {x["id"]: x for x in (base or []) if isinstance(x, dict) and "id" in x}
    ours_by = {x["id"]: x for x in ours if isinstance(x, dict) and "id" in x}
    deleted = set(base_by) - set(ours_by) if base is not None else set()
    added = [x for x in ours if x["id"] not in base_by]
    changed = {i for i, x in ours_by.items() if i in base_by and x != base_by[i]}
    out, seen = [], set()
    for x in theirs:
        i = x.get("id")
        if i in deleted:
            continue
        out.append(ours_by[i] if i in changed else x)
        seen.add(i)
    new = [x for x in added if x["id"] not in seen]
    # a change of ours to something they deleted stays: losing an edit is worse than a revived item
    revived = [ours_by[i] for i in changed if i not in seen and ours_by[i] not in new]
    return (new + revived + out) if prepend_added else (out + revived + new)


# ---- the blobs

def decode(key_b64: str, blob: dict) -> list[dict]:
    """The server's {ct, iv} (ct may be null: none yet) as a list."""
    if not blob or not blob.get("ct"):
        return []
    data = json.loads(crypto.decrypt(key_b64, blob["iv"], blob["ct"]).decode("utf-8"))
    return [x for x in data if isinstance(x, dict) and x.get("id")] if isinstance(data, list) else []


def encode(key_b64: str, items: list[dict]) -> dict:
    iv, ct = crypto.encrypt(key_b64, json.dumps(items, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
    return {"ct": ct, "iv": iv}


# ---- local state

@dataclass
class BookAnnotations:
    highlights: list[dict] = field(default_factory=list)
    bookmarks: list[dict] = field(default_factory=list)
    base: dict[str, list[dict] | None] = field(default_factory=lambda: {"highlights": None, "bookmarks": None})
    dirty: dict[str, bool] = field(default_factory=lambda: {"highlights": False, "bookmarks": False})

    def items(self, kind: str) -> list[dict]:
        return self.highlights if kind == "highlights" else self.bookmarks


def _locked(fn):
    @functools.wraps(fn)
    def wrapper(self, *a, **kw):
        with self._lock:
            return fn(self, *a, **kw)
    return wrapper


class AnnotationStore:
    """One JSON file per book key; shared by the reader (edits) and the sync workers."""

    def __init__(self, directory: Path | None = None) -> None:
        self.dir = directory or _DIR
        self._lock = threading.RLock()
        self._books: dict[str, BookAnnotations] = {}

    def _file(self, key: str) -> Path:
        return self.dir / f"{key}.json"

    def _load(self, key: str) -> BookAnnotations:
        if key in self._books:
            return self._books[key]
        ann = BookAnnotations()
        try:
            raw = json.loads(self._file(key).read_text(encoding="utf-8"))
            ann.highlights = list(raw.get("highlights") or [])
            ann.bookmarks = list(raw.get("bookmarks") or [])
            ann.base = {k: raw.get("base", {}).get(k) for k in KINDS}
            ann.dirty = {k: bool(raw.get("dirty", {}).get(k)) for k in KINDS}
        except (OSError, ValueError, AttributeError):
            pass
        self._books[key] = ann
        return ann

    def _save(self, key: str) -> None:
        ann = self._books[key]
        self.dir.mkdir(parents=True, exist_ok=True)
        tmp = self._file(key).with_suffix(".tmp")
        tmp.write_text(json.dumps({"highlights": ann.highlights, "bookmarks": ann.bookmarks, "base": ann.base,
                                   "dirty": ann.dirty}, ensure_ascii=False), encoding="utf-8")
        os.replace(tmp, self._file(key))

    @_locked
    def get(self, key: str) -> BookAnnotations:
        """A copy (the lists are copied, not the items: don't change those)."""
        a = self._load(key)
        return BookAnnotations(list(a.highlights), list(a.bookmarks),
                               {k: (list(v) if v is not None else None) for k, v in a.base.items()}, dict(a.dirty))

    @_locked
    def dirty_keys(self) -> list[str]:
        keys = set(self._books)
        if self.dir.exists():
            keys |= {p.stem for p in self.dir.glob("*.json")}
        return sorted(k for k in keys if any(self._load(k).dirty.values()))

    # ---- edits (they make the kind dirty: queued for the server)

    @_locked
    def add(self, key: str, kind: str, item: dict) -> None:
        ann = self._load(key)
        if kind == "bookmarks":
            ann.bookmarks.insert(0, item)  # newest first, like the web
        else:
            ann.highlights.append(item)
        ann.dirty[kind] = True
        self._save(key)

    @_locked
    def delete(self, key: str, kind: str, item_id: str) -> bool:
        ann = self._load(key)
        items = ann.items(kind)
        kept = [x for x in items if x.get("id") != item_id]
        if len(kept) == len(items):
            return False
        items[:] = kept
        ann.dirty[kind] = True
        self._save(key)
        return True

    @_locked
    def set_note(self, key: str, item_id: str, text: str) -> bool:
        ann = self._load(key)
        for i, h in enumerate(ann.highlights):
            if h.get("id") == item_id:
                if h.get("note") == text:
                    return False
                ann.highlights[i] = {**h, "note": text}
                ann.dirty["highlights"] = True
                self._save(key)
                return True
        return False

    # ---- the server

    @_locked
    def take_remote(self, key: str, kind: str, theirs: list[dict]) -> None:
        """The server's current list: ours becomes it, with our queued edits laid on top."""
        ann = self._load(key)
        if ann.dirty[kind]:
            merged = merge(ann.base[kind], ann.items(kind), theirs, prepend_added=kind == "bookmarks")
            ann.items(kind)[:] = merged
        else:
            ann.items(kind)[:] = theirs
        ann.base[kind] = list(theirs)
        self._save(key)

    @_locked
    def for_push(self, key: str, kind: str, theirs: list[dict]) -> tuple[list[dict], list[dict]]:
        """(what to write over the server's current list `theirs`, our list as it is now)."""
        ann = self._load(key)
        ours = list(ann.items(kind))
        return merge(ann.base[kind], ours, theirs, prepend_added=kind == "bookmarks"), ours

    @_locked
    def pushed(self, key: str, kind: str, written: list[dict], snapshot: list[dict]) -> None:
        """The server holds `written`, made from our list `snapshot`. Ours becomes `written`;
        what was edited while it was on its way stays queued on top."""
        ann = self._load(key)
        current = ann.items(kind)
        ann.base[kind] = list(written)
        if current == snapshot:
            current[:] = written
            ann.dirty[kind] = False
        else:
            current[:] = merge(snapshot, current, written, prepend_added=kind == "bookmarks")
            ann.dirty[kind] = True
        self._save(key)


_default: AnnotationStore | None = None


def default() -> AnnotationStore:
    global _default
    if _default is None:
        _default = AnnotationStore()
    return _default