messageLogger/src/store.py

"""What the plugin knows per message: server state (edits, deletion, copy
pointers, hidden marks) and deletions this client saw itself.

Both are cached in SQLite and loaded per chat on first use: a chat opened
again shows its marks at once, before (or without) the server, and a
message the plugin kept on screen still looks deleted after a restart.

Per chat the range of message ids already checked with the server (and
when) is kept too (`coverage`): see hooks/chat.py for when it is extended,
renewed or expired. An expired chat keeps its cache (shown meanwhile) until
the next successful check replaces it. Independently, cached answers older
than the configured number of days are removed, checked once a day
(`purge_old`).
Messages removed from view (hidden) are marked on the server, so every
client skips them; a change waits in `hidden_pending` until the server took
it (and overrides the server's mark meanwhile).
"""

import json
import os
import sqlite3
import threading
import time

import diag


class State:
    __slots__ = (
        "edit_count",
        "deleted",
        "logged",
        "history",
        "source_date",
        "sender_id",
        "hidden",
        "raw",
    )

    def __init__(self):
        self.edit_count = 0
        self.deleted = False
        self.logged = None  # {"chat_id", "message_id", "topic_id"}
        self.history = None
        self.source_date = None  # ISO string
        self.sender_id = None
        self.hidden = False  # removed from view (kept on the server)
        self.raw = None  # the server's answer for it

    @classmethod
    def from_server(cls, data):
        state = cls()
        state.edit_count = int(data.get("edit_count") or 0)
        state.deleted = data.get("deleted_at") is not None
        state.logged = data.get("logged")
        state.history = data.get("history")
        state.source_date = data.get("source_date")
        state.sender_id = data.get("sender_id")
        state.hidden = data.get("hidden_at") is not None
        state.raw = data
        return state


_lock = threading.RLock()
# (account, dialog_id) -> {message_id: State}
_states = {}
# (account, dialog_id) -> set of message ids deleted while we watched
_local = {}
_loaded = set()
# (account, message_id) -> monotonic time the user deleted it here
_own = {}
# (account, dialog_id) -> (low, high, checked_at)
_coverage = {}
_purged_on = None  # local date of the last age purge
_db = None
_OWN_TTL = 120


def _conn():
    global _db
    if _db is None:
        from file_utils import get_files_dir

        directory = os.path.join(get_files_dir(), "message_logger")
        os.makedirs(directory, exist_ok=True)
        _db = sqlite3.connect(
            os.path.join(directory, "state.db"), check_same_thread=False
        )
        _db.execute(
            "CREATE TABLE IF NOT EXISTS deleted ("
            "account INTEGER, dialog_id INTEGER, message_id INTEGER, at REAL, "
            "PRIMARY KEY (account, dialog_id, message_id))"
        )
        _db.execute(
            "CREATE TABLE IF NOT EXISTS states ("
            "account INTEGER, dialog_id INTEGER, message_id INTEGER, data TEXT, "
            "at REAL, PRIMARY KEY (account, dialog_id, message_id))"
        )
        _db.execute("CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT)")
        _db.execute(
            "CREATE TABLE IF NOT EXISTS hidden_pending ("
            "account INTEGER, dialog_id INTEGER, message_id INTEGER, hidden INTEGER, "
            "at REAL, PRIMARY KEY (account, dialog_id, message_id))"
        )
        # Before hidden marks were kept on the server: sent there once.
        if _db.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name='hidden'"
        ).fetchone():
            _db.execute(
                "INSERT OR IGNORE INTO hidden_pending "
                "SELECT account, dialog_id, message_id, 1, at FROM hidden"
            )
            _db.execute("DROP TABLE hidden")
        _db.execute(
            "CREATE TABLE IF NOT EXISTS coverage ("
            "account INTEGER, dialog_id INTEGER, low INTEGER, high INTEGER, at REAL, "
            "PRIMARY KEY (account, dialog_id))"
        )
        _db.commit()
    return _db


def _local_ids(account, dialog_id):
    """Local deletions of a chat; its cached states are loaded alongside."""
    key = (account, dialog_id)
    if key not in _loaded:
        _loaded.add(key)
        try:
            db = _conn()
            rows = db.execute(
                "SELECT message_id FROM deleted WHERE account=? AND dialog_id=?",
                key,
            ).fetchall()
            _local.setdefault(key, set()).update(r[0] for r in rows)
            states = _states.setdefault(key, {})
            for message_id, data in db.execute(
                "SELECT message_id, data FROM states WHERE account=? AND dialog_id=?",
                key,
            ).fetchall():
                states.setdefault(message_id, State.from_server(json.loads(data)))
            row = db.execute(
                "SELECT low, high, at FROM coverage WHERE account=? AND dialog_id=?",
                key,
            ).fetchone()
            if row is not None:
                _coverage.setdefault(key, tuple(row))
        except Exception as exc:
            diag.error("store.load", exc)
    return _local.setdefault(key, set())


def get(account, dialog_id, message_id):
    with _lock:
        local = _local_ids(account, dialog_id)  # loads the chat's cache
        state = _states.get((account, dialog_id), {}).get(message_id)
        if message_id in local:
            if state is None:
                state = State()
                _states.setdefault((account, dialog_id), {})[message_id] = state
            state.deleted = True
        return state


def is_deleted(account, dialog_id, message_id):
    state = get(account, dialog_id, message_id)
    return state is not None and state.deleted


def is_edited(account, dialog_id, message_id):
    state = get(account, dialog_id, message_id)
    return state is not None and state.edit_count > 0


def put_server_states(account, dialog_id, items):
    """Server state of messages (sync, lookups), kept and cached."""
    if not items:
        return
    with _lock:
        _local_ids(account, dialog_id)  # load the chat's cache first
        states = _states.setdefault((account, dialog_id), {})
        for data in items:
            states[int(data["message_id"])] = State.from_server(data)
        now = time.time()
        try:
            db = _conn()
            db.executemany(
                "INSERT OR REPLACE INTO states VALUES (?, ?, ?, ?, ?)",
                [
                    (account, dialog_id, int(d["message_id"]), json.dumps(d), now)
                    for d in items
                ],
            )
            db.commit()
        except Exception as exc:
            diag.error("store.put_server_states", exc)


def cached_deleted(account, dialog_id, low, high, hidden=False):
    """Cached server answers of deleted messages with a logged copy in
    [low, high] (what a restore needs, without asking the server again);
    hidden ones only when `hidden`."""
    with _lock:
        _local_ids(account, dialog_id)
        return [
            state.raw
            for message_id, state in _states.get((account, dialog_id), {}).items()
            if low <= message_id <= high
            and (hidden or not is_hidden(account, dialog_id, message_id))
            and state.raw is not None
            and state.raw.get("deleted_at")
            and state.raw.get("logged")
        ]


def coverage(account, dialog_id):
    """(low, high, checked_at) of the ids already checked, else None."""
    with _lock:
        _local_ids(account, dialog_id)
        return _coverage.get((account, dialog_id))


def set_coverage(account, dialog_id, low, high, at):
    with _lock:
        _coverage[(account, dialog_id)] = (int(low), int(high), float(at))
        try:
            db = _conn()
            db.execute(
                "INSERT OR REPLACE INTO coverage VALUES (?, ?, ?, ?, ?)",
                (account, dialog_id, int(low), int(high), float(at)),
            )
            db.commit()
        except Exception as exc:
            diag.error("store.set_coverage", exc)


def expire_coverage(account, dialog_id):
    """An edit/delete happened in the chat: check it again on next look."""
    current = coverage(account, dialog_id)
    if current is not None:
        set_coverage(account, dialog_id, current[0], current[1], 0.0)


def replace_window(account, dialog_id, low, high, items, capped, full=False):
    """A sync's answer for messages [low, high]: cached marks it no longer
    reports there are dropped (and returned) — all of the chat's when `full`
    (a fresh check replaces the old cache). A capped answer (the newest marks
    only) is complete from its oldest mark up."""
    if capped and items:
        low = max(low, min(int(d["message_id"]) for d in items))
    keep = {int(d["message_id"]) for d in items}
    with _lock:
        _local_ids(account, dialog_id)
        states = _states.setdefault((account, dialog_id), {})
        stale = [
            i for i in states if (full or low <= i <= high) and i not in keep
        ]
        for i in stale:
            del states[i]
        try:
            db = _conn()
            db.executemany(
                "DELETE FROM states WHERE account=? AND dialog_id=? AND message_id=?",
                [(account, dialog_id, i) for i in stale],
            )
            db.commit()
        except Exception as exc:
            diag.error("store.replace_window", exc)
    put_server_states(account, dialog_id, items)
    return stale


def mark_deleted(account, dialog_id, message_ids):
    """Deletions this client saw (kept on screen or confirmed missing)."""
    ids = [int(i) for i in message_ids]
    with _lock:
        _local_ids(account, dialog_id).update(ids)
        now = time.time()
        try:
            db = _conn()
            db.executemany(
                "INSERT OR IGNORE INTO deleted VALUES (?, ?, ?, ?)",
                [(account, dialog_id, i, now) for i in ids],
            )
            db.commit()
        except Exception as exc:
            diag.error("store.mark_deleted", exc)


def forget_deleted(account, dialog_id, message_ids):
    """The user deleted kept messages themselves: stop tracking them."""
    ids = [int(i) for i in message_ids]
    with _lock:
        _local_ids(account, dialog_id).difference_update(ids)
        states = _states.get((account, dialog_id), {})
        for i in ids:
            states.pop(i, None)
        try:
            db = _conn()
            for table in ("deleted", "states"):
                db.executemany(
                    f"DELETE FROM {table} WHERE account=? AND dialog_id=? AND message_id=?",
                    [(account, dialog_id, i) for i in ids],
                )
            db.commit()
        except Exception as exc:
            diag.error("store.forget_deleted", exc)


def reset(account=None, dialog_id=None):
    """Forgets cached marks, local deletion records and checked ranges of a
    chat (or of every chat when `dialog_id` is None); the next look asks the
    server again. Messages the user removed from view stay hidden.
    Returns the number of marks forgotten."""
    where, args = "", ()
    if dialog_id is not None:
        where, args = " WHERE account=? AND dialog_id=?", (int(account), int(dialog_id))
    with _lock:
        removed = 0
        try:
            db = _conn()
            for table in ("states", "deleted", "coverage"):
                count = db.execute(f"DELETE FROM {table}{where}", args).rowcount
                if table != "coverage":
                    removed += count
            db.commit()
        except Exception as exc:
            diag.error("store.reset", exc)
        if dialog_id is None:
            _states.clear()
            _local.clear()
            _coverage.clear()
            _loaded.clear()
        else:
            key = (int(account), int(dialog_id))
            for cache in (_states, _local, _coverage):
                cache.pop(key, None)
            _loaded.discard(key)
        return removed


_pending = None  # (account, dialog_id, message_id) -> hidden, not yet on the server
_revealed = set()  # (account, dialog_id) showing their hidden messages


def _pending_marks():
    global _pending
    if _pending is None:
        found = {}
        try:
            for account, dialog_id, message_id, hidden in _conn().execute(
                "SELECT account, dialog_id, message_id, hidden FROM hidden_pending"
            ).fetchall():
                found[(account, dialog_id, message_id)] = bool(hidden)
        except Exception as exc:
            diag.error("store.hidden_pending", exc)
        _pending = found
    return _pending


def set_hidden(account, dialog_id, message_ids, hidden=True):
    """Removes messages from view (or shows them again) here at once; the
    change waits for the server in `hidden_pending` (see `pending_hidden`)."""
    ids = [int(i) for i in message_ids]
    with _lock:
        pending = _pending_marks()
        states = _states.get((account, dialog_id), {})
        for i in ids:
            pending[(account, dialog_id, i)] = bool(hidden)
            state = states.get(i)
            if state is not None:
                state.hidden = bool(hidden)
        try:
            db = _conn()
            now = time.time()
            db.executemany(
                "INSERT OR REPLACE INTO hidden_pending VALUES (?, ?, ?, ?, ?)",
                [(account, dialog_id, i, int(bool(hidden)), now) for i in ids],
            )
            db.commit()
        except Exception as exc:
            diag.error("store.set_hidden", exc)


def pending_hidden():
    """Hidden marks not on the server yet: {(account, dialog_id, hidden): [ids]}."""
    with _lock:
        out = {}
        for (account, dialog_id, message_id), hidden in _pending_marks().items():
            out.setdefault((account, dialog_id, hidden), []).append(message_id)
        return out


def hidden_sent(account, dialog_id, message_ids, hidden):
    """The server took these marks: they are its state now."""
    ids = [int(i) for i in message_ids]
    with _lock:
        pending = _pending_marks()
        sent = [i for i in ids if pending.get((account, dialog_id, i)) == bool(hidden)]
        for i in sent:
            del pending[(account, dialog_id, i)]
        states = _states.get((account, dialog_id), {})
        changed = []
        for i in sent:
            state = states.get(i)
            if state is not None and state.raw is not None:
                state.raw["hidden_at"] = (
                    state.raw.get("hidden_at") or _now_iso() if hidden else None
                )
                state.hidden = bool(hidden)
                changed.append((account, dialog_id, i, json.dumps(state.raw)))
        try:
            db = _conn()
            db.executemany(
                "DELETE FROM hidden_pending WHERE account=? AND dialog_id=? AND message_id=?",
                [(account, dialog_id, i) for i in sent],
            )
            db.executemany(
                "UPDATE states SET data=? WHERE account=? AND dialog_id=? AND message_id=?",
                [(data, a, d, i) for a, d, i, data in changed],
            )
            db.commit()
        except Exception as exc:
            diag.error("store.hidden_sent", exc)


def _now_iso():
    import datetime

    return datetime.datetime.now(datetime.timezone.utc).isoformat()


def is_hidden(account, dialog_id, message_id):
    with _lock:
        pending = _pending_marks().get((account, dialog_id, int(message_id)))
        if pending is not None:
            return pending
        state = _states.get((account, dialog_id), {}).get(int(message_id))
        return state is not None and state.hidden


def hidden_deleted(account, dialog_id):
    """Ids of the chat's deleted messages that are hidden (cached)."""
    with _lock:
        _local_ids(account, dialog_id)
        return [
            i
            for i, state in _states.get((account, dialog_id), {}).items()
            if state.deleted and is_hidden(account, dialog_id, i)
        ]


def reveal(account, dialog_id, on):
    with _lock:
        if on:
            _revealed.add((account, dialog_id))
        else:
            _revealed.discard((account, dialog_id))


def is_revealed(account, dialog_id):
    return (account, dialog_id) in _revealed


_removed_dialogs = set()  # (account, dialog_id) of removed chats listed


def set_removed_dialog(account, dialog_id, listed):
    """A removed chat listed among the chats (hooks/removed_chats.py): its
    state comes with its history, the chat is not synced."""
    with _lock:
        if listed:
            _removed_dialogs.add((int(account), int(dialog_id)))
        else:
            _removed_dialogs.discard((int(account), int(dialog_id)))


def is_removed_dialog(account, dialog_id):
    return (int(account), int(dialog_id)) in _removed_dialogs


def get_meta(key):
    try:
        with _lock:
            row = _conn().execute("SELECT value FROM meta WHERE key=?", (key,)).fetchone()
        return row[0] if row else None
    except Exception as exc:
        diag.error("store.get_meta", exc)
        return None


def put_meta(key, value):
    try:
        with _lock:
            db = _conn()
            db.execute("INSERT OR REPLACE INTO meta VALUES (?, ?)", (key, value))
            db.commit()
    except Exception as exc:
        diag.error("store.put_meta", exc)


def note_own_deletion(account, message_ids):
    now = time.monotonic()
    with _lock:
        for key in [k for k, at in _own.items() if now - at > _OWN_TTL]:
            del _own[key]
        for i in message_ids:
            _own[(account, int(i))] = now


def is_own_deletion(account, message_id):
    with _lock:
        return (account, int(message_id)) in _own


def purge_old(days, force=False):
    """Removes cached server answers (marks, checked ranges) older than
    `days` (None: never), at most once per local calendar day unless
    `force`. Local deletion records are kept: they are what this client saw,
    not a copy of the server."""
    import datetime

    global _purged_on
    today = datetime.date.today().isoformat()
    if not force and _purged_on == today:
        return 0
    with _lock:
        try:
            db = _conn()
            row = db.execute("SELECT value FROM meta WHERE key='purged_on'").fetchone()
            if not force and row is not None and row[0] == today:
                _purged_on = today
                return 0
            removed = 0
            if days is not None:
                cutoff = time.time() - days * 86400
                removed = db.execute("DELETE FROM states WHERE at < ?", (cutoff,)).rowcount
                db.execute("DELETE FROM coverage WHERE at < ?", (cutoff,))
            db.execute(
                "INSERT OR REPLACE INTO meta VALUES ('purged_on', ?)", (today,)
            )
            db.commit()
            _purged_on = today
        except Exception as exc:
            diag.error("store.purge_old", exc)
            return 0
        if removed:
            # Reload chats from the database on next use.
            _states.clear()
            _coverage.clear()
            _loaded.clear()
            _local.clear()
        diag.log(f"cache purge (older than {days} days): {removed} marks removed")
        return removed