homelab-brain/chatlogger/store.py

"""SQLite store for the Arakawa Concierge chat logger.

One DB file holds all watched groups. Schema is created on first open.
Full-text search via FTS5 across message text, captions, sender names.
"""
from __future__ import annotations

import json
import os
import sqlite3
import time
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Iterable

SCHEMA = """
CREATE TABLE IF NOT EXISTS messages (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    chat_id INTEGER NOT NULL,
    chat_title TEXT,
    chat_type TEXT,
    message_id INTEGER NOT NULL,
    user_id INTEGER,
    user_first_name TEXT,
    user_last_name TEXT,
    user_username TEXT,
    text TEXT,
    media_type TEXT,
    media_caption TEXT,
    media_file_id TEXT,
    reply_to_message_id INTEGER,
    forward_from TEXT,
    timestamp_utc INTEGER NOT NULL,
    edit_timestamp_utc INTEGER,
    raw_json TEXT,
    UNIQUE(chat_id, message_id)
);

CREATE INDEX IF NOT EXISTS idx_chat_ts ON messages(chat_id, timestamp_utc);
CREATE INDEX IF NOT EXISTS idx_user_ts ON messages(user_id, timestamp_utc);

CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
    text, media_caption, user_first_name, user_username,
    content='messages', content_rowid='id',
    tokenize='unicode61 remove_diacritics 1'
);

CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
    INSERT INTO messages_fts(rowid, text, media_caption, user_first_name, user_username)
    VALUES (new.id, new.text, new.media_caption, new.user_first_name, new.user_username);
END;

CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
    INSERT INTO messages_fts(messages_fts, rowid, text, media_caption, user_first_name, user_username)
    VALUES ('delete', old.id, old.text, old.media_caption, old.user_first_name, old.user_username);
END;

CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
    INSERT INTO messages_fts(messages_fts, rowid, text, media_caption, user_first_name, user_username)
    VALUES ('delete', old.id, old.text, old.media_caption, old.user_first_name, old.user_username);
    INSERT INTO messages_fts(rowid, text, media_caption, user_first_name, user_username)
    VALUES (new.id, new.text, new.media_caption, new.user_first_name, new.user_username);
END;

CREATE TABLE IF NOT EXISTS message_edits (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    chat_id INTEGER NOT NULL,
    message_id INTEGER NOT NULL,
    text_before TEXT,
    text_after TEXT,
    edit_timestamp_utc INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_edits_msg ON message_edits(chat_id, message_id);
"""


@dataclass
class StoredMessage:
    id: int
    chat_id: int
    chat_title: str | None
    message_id: int
    user_id: int | None
    user_display: str
    text: str | None
    media_type: str | None
    media_caption: str | None
    reply_to_message_id: int | None
    timestamp_utc: int
    edit_timestamp_utc: int | None


class Store:
    def __init__(self, db_path: str):
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
        with self._conn() as c:
            c.executescript(SCHEMA)
            c.execute("PRAGMA journal_mode=WAL")
            c.execute("PRAGMA synchronous=NORMAL")

    @contextmanager
    def _conn(self):
        conn = sqlite3.connect(self.db_path, timeout=10, isolation_level=None)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()

    @staticmethod
    def _user_display(row: sqlite3.Row | dict[str, Any]) -> str:
        first = (row["user_first_name"] or "").strip()
        last = (row["user_last_name"] or "").strip()
        uname = (row["user_username"] or "").strip()
        full = " ".join(p for p in (first, last) if p)
        if full and uname:
            return f"{full} (@{uname})"
        return full or (f"@{uname}" if uname else f"user:{row['user_id']}")

    def add_message(
        self,
        *,
        chat_id: int,
        chat_title: str | None,
        chat_type: str | None,
        message_id: int,
        user_id: int | None,
        user_first_name: str | None,
        user_last_name: str | None,
        user_username: str | None,
        text: str | None,
        media_type: str | None,
        media_caption: str | None,
        media_file_id: str | None,
        reply_to_message_id: int | None,
        forward_from: str | None,
        timestamp_utc: int,
        raw_json: str,
    ) -> bool:
        with self._conn() as c:
            try:
                c.execute(
                    """INSERT INTO messages (
                        chat_id, chat_title, chat_type, message_id,
                        user_id, user_first_name, user_last_name, user_username,
                        text, media_type, media_caption, media_file_id,
                        reply_to_message_id, forward_from, timestamp_utc, raw_json
                    ) VALUES (?,?,?,?, ?,?,?,?, ?,?,?,?, ?,?,?,?)""",
                    (
                        chat_id, chat_title, chat_type, message_id,
                        user_id, user_first_name, user_last_name, user_username,
                        text, media_type, media_caption, media_file_id,
                        reply_to_message_id, forward_from, timestamp_utc, raw_json,
                    ),
                )
                return True
            except sqlite3.IntegrityError:
                return False

    def update_message(
        self,
        *,
        chat_id: int,
        message_id: int,
        new_text: str | None,
        new_caption: str | None,
        edit_timestamp_utc: int,
    ) -> bool:
        with self._conn() as c:
            row = c.execute(
                "SELECT id, text, media_caption FROM messages WHERE chat_id=? AND message_id=?",
                (chat_id, message_id),
            ).fetchone()
            if not row:
                return False
            text_before = row["text"] or row["media_caption"]
            text_after = new_text if new_text is not None else new_caption
            c.execute(
                """UPDATE messages SET text=?, media_caption=?, edit_timestamp_utc=?
                   WHERE chat_id=? AND message_id=?""",
                (new_text, new_caption, edit_timestamp_utc, chat_id, message_id),
            )
            c.execute(
                """INSERT INTO message_edits (chat_id, message_id, text_before, text_after, edit_timestamp_utc)
                   VALUES (?,?,?,?,?)""",
                (chat_id, message_id, text_before, text_after, edit_timestamp_utc),
            )
            return True

    def recent(self, chat_id: int | None = None, limit: int = 50) -> list[StoredMessage]:
        with self._conn() as c:
            if chat_id is None:
                rows = c.execute(
                    "SELECT * FROM messages ORDER BY timestamp_utc DESC LIMIT ?",
                    (limit,),
                ).fetchall()
            else:
                rows = c.execute(
                    "SELECT * FROM messages WHERE chat_id=? ORDER BY timestamp_utc DESC LIMIT ?",
                    (chat_id, limit),
                ).fetchall()
        return [self._row_to_msg(r) for r in reversed(rows)]

    def since(self, chat_id: int | None, since_ts: int, limit: int = 1000) -> list[StoredMessage]:
        with self._conn() as c:
            if chat_id is None:
                rows = c.execute(
                    "SELECT * FROM messages WHERE timestamp_utc >= ? ORDER BY timestamp_utc ASC LIMIT ?",
                    (since_ts, limit),
                ).fetchall()
            else:
                rows = c.execute(
                    "SELECT * FROM messages WHERE chat_id=? AND timestamp_utc >= ? ORDER BY timestamp_utc ASC LIMIT ?",
                    (chat_id, since_ts, limit),
                ).fetchall()
        return [self._row_to_msg(r) for r in rows]

    def search(self, query: str, chat_id: int | None = None, limit: int = 30) -> list[StoredMessage]:
        with self._conn() as c:
            sql = (
                "SELECT m.* FROM messages m "
                "JOIN messages_fts f ON f.rowid = m.id "
                "WHERE messages_fts MATCH ?"
            )
            params: list[Any] = [query]
            if chat_id is not None:
                sql += " AND m.chat_id = ?"
                params.append(chat_id)
            sql += " ORDER BY m.timestamp_utc DESC LIMIT ?"
            params.append(limit)
            rows = c.execute(sql, params).fetchall()
        return [self._row_to_msg(r) for r in rows]

    def stats(self) -> dict[str, Any]:
        with self._conn() as c:
            total = c.execute("SELECT COUNT(*) AS n FROM messages").fetchone()["n"]
            per_chat = c.execute(
                """SELECT chat_id, MAX(chat_title) AS title, COUNT(*) AS n,
                          MIN(timestamp_utc) AS first_ts, MAX(timestamp_utc) AS last_ts
                   FROM messages GROUP BY chat_id ORDER BY n DESC"""
            ).fetchall()
            edits = c.execute("SELECT COUNT(*) AS n FROM message_edits").fetchone()["n"]
        return {
            "total_messages": total,
            "total_edits": edits,
            "chats": [dict(r) for r in per_chat],
            "db_size_bytes": os.path.getsize(self.db_path) if os.path.exists(self.db_path) else 0,
        }

    def chats(self) -> list[dict[str, Any]]:
        with self._conn() as c:
            rows = c.execute(
                """SELECT chat_id, MAX(chat_title) AS title, MAX(chat_type) AS type,
                          COUNT(*) AS n, MAX(timestamp_utc) AS last_ts
                   FROM messages GROUP BY chat_id ORDER BY last_ts DESC"""
            ).fetchall()
        return [dict(r) for r in rows]

    def _row_to_msg(self, row: sqlite3.Row) -> StoredMessage:
        return StoredMessage(
            id=row["id"],
            chat_id=row["chat_id"],
            chat_title=row["chat_title"],
            message_id=row["message_id"],
            user_id=row["user_id"],
            user_display=self._user_display(row),
            text=row["text"],
            media_type=row["media_type"],
            media_caption=row["media_caption"],
            reply_to_message_id=row["reply_to_message_id"],
            timestamp_utc=row["timestamp_utc"],
            edit_timestamp_utc=row["edit_timestamp_utc"],
        )


def format_messages(msgs: Iterable[StoredMessage], tz_offset_hours: int = 0) -> str:
    """Render messages as a chat-style transcript (chronological)."""
    out: list[str] = []
    for m in msgs:
        ts = time.gmtime(m.timestamp_utc + tz_offset_hours * 3600)
        stamp = time.strftime("%Y-%m-%d %H:%M", ts)
        body = m.text or m.media_caption or (f"[{m.media_type}]" if m.media_type else "")
        if m.media_type and (m.text or m.media_caption):
            body = f"[{m.media_type}] {body}"
        reply = f" ↩#{m.reply_to_message_id}" if m.reply_to_message_id else ""
        edit = " (bearbeitet)" if m.edit_timestamp_utc else ""
        out.append(f"[{stamp}] {m.user_display}{reply}{edit}: {body}")
    return "\n".join(out)