chat.py

import copy
import hashlib
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Iterator, List, Optional, cast

from google.api_core.exceptions import AlreadyExists, Conflict, FailedPrecondition, NotFound
from google.cloud import firestore, firestore_v1
from google.cloud.firestore_v1 import FieldFilter

from database.firestore_index_registry import (
    CURRENT_CHAT_SESSION_ORDERED_QUERY,
    CURRENT_CHAT_SESSION_QUERY,
)

# Sessions are per-user and per-app, so this is a ceiling on a small collection
# rather than a page size; it exists so a pathological account cannot turn one
# lookup into an unbounded read.
CURRENT_CHAT_SESSION_SCAN_LIMIT = 200

from models.chat import Message
from utils import encryption
from ._client import db
from .helpers import prepare_for_read, prepare_for_write, set_data_protection_level
from database.read_boundary import parse_snapshot_or_none

logger = logging.getLogger(__name__)

BATCH_LIMIT = 500  # Firestore hard limit
DELETE_MESSAGES_BATCH_LIMIT = 200  # Leaves room for one session-counter write per deleted message.
DELETE_MESSAGES_CONFLICT_RETRIES = 3
CHAT_HISTORY_BASE_VISIBLE_MESSAGES = 10
CHAT_HISTORY_APPEND_EPOCH_MESSAGES = 8
# Maximum number of reported (hidden) rows to over-fetch per raw Firestore
# query when reading cache-aligned history. Keeps the raw read bounded even
# when a user has thousands of lifetime reported messages; the newest page
# rarely contains more reported rows than this cap.
CHAT_HISTORY_REPORTED_RAW_SCAN_CAP = 50
# Extra documents a visible page may stream *beyond* the rows it would need if none
# were reported. The floor is the page itself, never this: the previous raw
# ``.offset(n).limit(m)`` query already streamed n + m documents, so budgeting
# ``needed + slack`` can only read more than before by the slack, and can never fail
# to service an offset the old query serviced. Capping the total instead made a deep
# offset return an empty page, which the router reads as end-of-results -- the same
# defect this scan exists to fix.
CHAT_MESSAGES_VISIBLE_PAGE_SCAN_SLACK = 1000


class ClientMessageIdPayloadConflict(ValueError):
    """The same client id was reused for a different immutable message."""


class MessageReconcileCursorError(ValueError):
    """A desktop journal cursor is absent or outside the authenticated scope."""


def _typed_doc(doc: Any) -> Dict[str, Any]:
    """Typed adapter for a Firestore DocumentSnapshot.to_dict() result.

    Returns an empty dict when the document has no fields (None payload),
    so callers can safely mutate and read keys without Optional checks.
    """
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else {}


# *********************************
# ******* ENCRYPTION HELPERS ******
# *********************************


def _encrypt_chat_data(chat_data: Dict[str, Any], uid: str) -> Dict[str, Any]:
    data = copy.deepcopy(chat_data)

    if 'text' in data and isinstance(data['text'], str):
        data['text'] = encryption.encrypt(data['text'], uid)
    return data


def _decrypt_chat_data(chat_data: Dict[str, Any], uid: str) -> Dict[str, Any]:
    data = copy.deepcopy(chat_data)

    if 'text' in data and isinstance(data['text'], str):
        try:
            data['text'] = encryption.decrypt(data['text'], uid)
        except Exception:
            pass

    return data


def _prepare_data_for_write(data: Dict[str, Any], uid: str, level: str) -> Dict[str, Any]:
    if level == 'enhanced':
        return _encrypt_chat_data(data, uid)
    return data


def _prepare_message_for_read(message_data: Dict[str, Any], uid: str) -> Dict[str, Any]:
    level = message_data.get('data_protection_level')
    if level == 'enhanced':
        return _decrypt_chat_data(message_data, uid)

    return message_data


def decrypt_message_payload(message_data: Dict[str, Any], uid: str) -> Dict[str, Any]:
    """Public read-path decryption for one raw message document.

    Callers outside this module (the feedback-report context hydrator) need the
    same enhanced-protection handling `get_message` applies, but starting from a
    raw snapshot dict they already hold. Exposed rather than reaching into the
    private helper so the encryption contract has one owner.
    """
    return _prepare_message_for_read(message_data, uid)


# *****************************
# ********** CRUD *************
# *****************************


@set_data_protection_level(data_arg_name='message_data')
@prepare_for_write(data_arg_name='message_data', prepare_func=_prepare_data_for_write)
def add_message(uid: str, message_data: Dict[str, Any]) -> Dict[str, Any]:
    del message_data['memories']
    user_ref = db.collection('users').document(uid)
    user_ref.collection('messages').add(message_data)
    return message_data


def add_app_message(text: str, app_id: str, uid: str, conversation_id: Optional[str] = None) -> Message:
    """Add a chat message an app posted for the user, linking it to that app's chat session so it
    appears in the chat feed. get_messages filters by chat_session_id whenever a session exists, so
    a message stored without one is never returned on that path."""
    chat_session = get_chat_session(uid, app_id=app_id)
    chat_session_id = chat_session['id'] if chat_session else None

    ai_message = Message(
        id=str(uuid.uuid4()),
        text=text,
        created_at=datetime.now(timezone.utc),
        sender='ai',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageSender enum
        app_id=app_id,
        from_external_integration=False,
        type='text',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageType enum
        memories_id=[conversation_id] if conversation_id else [],
        chat_session_id=chat_session_id,
    )
    add_message(uid, ai_message.model_dump())
    if chat_session_id:
        add_message_to_chat_session(uid, chat_session_id, ai_message.id)
    return ai_message


def add_integration_chat_message(text: str, app_id: Optional[str], uid: str) -> Message:
    """Add a chat message from an external integration (e.g. notification API),
    linking it to the user's existing chat session so it appears in the chat feed."""
    chat_session = get_chat_session(uid, app_id=app_id)
    chat_session_id = chat_session['id'] if chat_session else None

    ai_message = Message(
        id=str(uuid.uuid4()),
        text=text,
        created_at=datetime.now(timezone.utc),
        sender='ai',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageSender enum
        app_id=app_id,
        from_external_integration=True,
        type='text',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageType enum
        chat_session_id=chat_session_id,
    )
    add_message(uid, ai_message.model_dump())
    if chat_session_id:
        add_message_to_chat_session(uid, chat_session_id, ai_message.id)
    return ai_message


def add_summary_message(text: str, uid: str) -> Message:
    ai_message = Message(
        id=str(uuid.uuid4()),
        text=text,
        created_at=datetime.now(timezone.utc),
        sender='ai',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageSender enum
        app_id=None,
        from_external_integration=False,
        type='day_summary',  # type: ignore[reportArgumentType]  # pydantic accepts str for MessageType enum
        memories_id=[],
    )
    add_message(uid, ai_message.model_dump())
    return ai_message


@prepare_for_read(decrypt_func=_prepare_message_for_read)
def get_app_messages(
    uid: str, app_id: str, limit: int = 20, include_conversations: bool = False
) -> List[Dict[str, Any]]:
    """Return an app's newest visible messages, up to ``limit``.

    ``reported`` is intentionally filtered in Python because the legacy data
    model permits the field to be absent.  The Firestore limit must therefore
    not run before that visibility rule: a reported row inside the raw page
    would otherwise consume a caller-visible slot and leave an older visible
    message unfetched.  This follows the same bounded visible-row scan as
    ``get_messages`` below.
    """
    visible_limit = max(0, int(limit))
    if visible_limit == 0:
        return []

    user_ref = db.collection('users').document(uid)
    query: Any = (
        user_ref.collection('messages')
        .where(filter=FieldFilter('plugin_id', '==', app_id))
        .order_by('created_at', direction=firestore.Query.DESCENDING)
    )
    # A clean page needs exactly ``visible_limit`` raw rows.  Bound only the
    # extra rows needed to cross reported records, so this cannot become an
    # unbounded history read while a deep run of reported rows still has a
    # flat allowance to cross.
    scan_budget = visible_limit + CHAT_MESSAGES_VISIBLE_PAGE_SCAN_SLACK
    scanned = 0
    reported_row_seen = False
    cursor_snapshot: Any = None
    messages: List[Dict[str, Any]] = []
    conversations_id: set[str] = set()

    while scanned < scan_budget and len(messages) < visible_limit:
        # A clean page reads exactly its requested visible rows, even when it
        # needs more than one capped 100-row batch. Once a reported row has
        # appeared, use capped batches to cross a dense hidden run without
        # turning a missing visible row into one read per document.
        batch_limit = min(100, scan_budget - scanned)
        if not reported_row_seen:
            batch_limit = min(batch_limit, max(1, visible_limit - len(messages)))
        page_query = query.start_after(cursor_snapshot) if cursor_snapshot is not None else query
        documents = list(page_query.limit(batch_limit).stream())
        if not documents:
            break

        for document in documents:
            scanned += 1
            cursor_snapshot = document
            message: Dict[str, Any] = _typed_doc(document)
            if message.get('reported') is True:
                reported_row_seen = True
                continue
            messages.append(message)
            conversations_id.update(message.get('memories_id', []))
            if len(messages) == visible_limit:
                break

        if len(documents) < batch_limit:
            break

    if not include_conversations:
        return messages

    # Fetch all conversations at once
    conversations: Dict[str, Any] = {}
    conversations_ref = user_ref.collection('conversations')
    doc_refs = [conversations_ref.document(str(conversation_id)) for conversation_id in conversations_id]
    docs = db.get_all(doc_refs)
    for doc in docs:
        if doc.exists:
            conversation: Dict[str, Any] = _typed_doc(doc)
            conversations[conversation['id']] = conversation

    # Attach conversations to messages
    for message in messages:
        message['memories'] = [
            conversations[conversation_id]
            for conversation_id in message.get('memories_id', [])
            if conversation_id in conversations
        ]

    return messages


@prepare_for_read(decrypt_func=_prepare_message_for_read)
def get_messages(
    uid: str,
    limit: int = 20,
    offset: int = 0,
    include_conversations: bool = False,
    app_id: Optional[str] = None,
    chat_session_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """Return a visible chat page with offset counted in non-reported rows.

    Firestore cannot apply ``reported != True`` cheaply for legacy docs that omit
    the field, so reported rows are filtered in Python. Applying ``limit`` /
    ``offset`` on the raw query first makes pages short and advances past
    visible messages the client never saw. Scan with a bounded budget (same
    spirit as ``get_messages_reconcile_page``) until ``offset`` visible rows are
    skipped and ``limit`` visible rows are collected.
    """
    logger.info(f'get_messages {uid} {limit} {offset} {app_id} {include_conversations}')
    user_ref = db.collection('users').document(uid)
    messages_ref = user_ref.collection('messages')
    if chat_session_id:
        # Session-scoped query: filter by session only, skip plugin_id filter
        # because the session already determines which app the messages belong to.
        messages_ref = messages_ref.where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
    else:
        # App-scoped query: filter by plugin_id (None = main chat)
        messages_ref = messages_ref.where(filter=FieldFilter('plugin_id', '==', app_id))

    query: Any = messages_ref.order_by('created_at', direction=firestore.Query.DESCENDING)

    # A page with no reported rows needs exactly this many documents, which is what the
    # old raw query streamed. Bound the *slack* on top of it, not the page itself.
    needed = max(offset, 0) + max(limit, 0)
    # Flat slack, not proportional. Scaling it with the page size gave a small page a
    # tiny allowance (limit=2 -> 6 documents), so a dense run of reported rows still
    # returned an empty page the router reads as end-of-results. The read cost is set
    # by the batch sizing below, not by this ceiling, so a flat allowance costs a clean
    # page nothing and only bounds how far a page that meets reported rows may scan.
    scan_budget = needed + CHAT_MESSAGES_VISIBLE_PAGE_SCAN_SLACK
    scanned = 0
    visible_skipped = 0
    messages: List[Dict[str, Any]] = []
    conversations_id: set[str] = set()
    files_id: set[str] = set()
    cursor_snapshot: Any = None

    while scanned < scan_budget and len(messages) < limit:
        # Read exactly what the page needs before reading any slack. Without this the
        # first batch was a flat 100 documents, so the chat-send path's limit=5 and
        # limit=15 reads streamed ~20x the documents they used to. Slack is only paid
        # for by a page that actually met a reported row.
        batch_limit = min(100, scan_budget - scanned)
        if scanned == 0:
            batch_limit = min(batch_limit, max(1, needed))
        page_query = query.start_after(cursor_snapshot) if cursor_snapshot is not None else query
        documents = list(page_query.limit(batch_limit).stream())
        if not documents:
            break

        for document in documents:
            scanned += 1
            cursor_snapshot = document
            message: Dict[str, Any] = _typed_doc(document)
            if message.get('reported') is True:
                continue
            if visible_skipped < offset:
                visible_skipped += 1
                continue
            messages.append(message)
            conversations_id.update(message.get('memories_id', []))
            files_id.update(message.get('files_id', []))
            if len(messages) == limit:
                break

        if len(documents) < batch_limit:
            break

    if not include_conversations:
        return messages

    # Fetch all conversations at once
    conversations: Dict[str, Any] = {}
    conversations_ref = user_ref.collection('conversations')
    doc_refs = [conversations_ref.document(str(conversation_id)) for conversation_id in conversations_id]
    docs = db.get_all(doc_refs)
    for doc in docs:
        if doc.exists:
            conversation: Dict[str, Any] = _typed_doc(doc)
            conversations[conversation['id']] = conversation

    # Attach conversations to messages
    for message in messages:
        message['memories'] = [
            conversations[conversation_id]
            for conversation_id in message.get('memories_id', [])
            if conversation_id in conversations
        ]

    # Fetch file chat
    files: Dict[str, Any] = {}
    files_ref = user_ref.collection('files')
    files_ref = [files_ref.document(str(file_id)) for file_id in files_id]
    doc_files = db.get_all(files_ref)
    for doc in doc_files:
        if doc.exists:
            file: Dict[str, Any] = _typed_doc(doc)
            files[file['id']] = file

    # Attach files to messages
    for message in messages:
        message['files'] = [files[file_id] for file_id in message.get('files_id', []) if file_id in files]

    return messages


def cache_aligned_history_limit(total_visible_messages: int) -> int:
    """Return a bounded history size whose start moves only at epoch boundaries.

    A fixed newest-N window changes at the front on every chat turn, invalidating
    Anthropic's cumulative message-prefix cache. This policy keeps at least the
    existing ten-message continuity window and lets it grow append-only for eight
    messages before resetting to ten. The request therefore carries 10..17
    messages, never less history than before and never an unbounded transcript.
    """
    if total_visible_messages < 0:
        raise ValueError('total_visible_messages must be non-negative')
    if total_visible_messages <= CHAT_HISTORY_BASE_VISIBLE_MESSAGES:
        return total_visible_messages
    return CHAT_HISTORY_BASE_VISIBLE_MESSAGES + (
        (total_visible_messages - CHAT_HISTORY_BASE_VISIBLE_MESSAGES) % CHAT_HISTORY_APPEND_EPOCH_MESSAGES
    )


def get_cache_aligned_messages(
    uid: str,
    *,
    app_id: Optional[str] = None,
    chat_session_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """Read a cache-aligned, scope-safe chat history in newest-first order.

    Reported messages are excluded from the visible count and read. Over-fetching
    by the scoped reported count guarantees the target number of visible messages
    even when hidden records fall inside the selected raw Firestore page.
    """
    user_ref = db.collection('users').document(uid)
    scoped_ref = user_ref.collection('messages')
    if chat_session_id:
        scoped_ref = scoped_ref.where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
    else:
        scoped_ref = scoped_ref.where(filter=FieldFilter('plugin_id', '==', app_id))

    total_result = scoped_ref.count().get()
    total = int(total_result[0][0].value) if total_result and total_result[0] else 0
    reported_result = scoped_ref.where(filter=FieldFilter('reported', '==', True)).count().get()
    reported = int(reported_result[0][0].value) if reported_result and reported_result[0] else 0
    visible_total = max(0, total - reported)
    visible_limit = cache_aligned_history_limit(visible_total)
    if visible_limit == 0:
        return []

    # Cap the raw Firestore read so a large lifetime reported count cannot
    # cause unbounded document reads on every chat send. The over-fetch only
    # needs to cover reported rows that fall inside the newest raw page, not
    # the lifetime total.
    reported_overfetch = min(reported, CHAT_HISTORY_REPORTED_RAW_SCAN_CAP)
    raw_limit = min(total, visible_limit + reported_overfetch)
    return get_messages(
        uid,
        limit=raw_limit,
        app_id=app_id,
        chat_session_id=chat_session_id,
    )[:visible_limit]


@prepare_for_read(decrypt_func=_prepare_message_for_read)
def get_messages_reconcile_page(
    uid: str,
    *,
    limit: int,
    cursor_message_id: Optional[str] = None,
    app_id: Optional[str] = None,
    chat_session_id: Optional[str] = None,
) -> tuple[List[Dict[str, Any]], Optional[str], bool]:
    """Return a stable, owner-scoped keyset page for the desktop journal.

    `get_messages` remains offset-compatible for existing clients. This path is
    deliberately separate: a Firestore document snapshot cursor cannot skip or
    duplicate rows when a newer message is inserted between page requests.
    Reported rows are scanned but not returned, with a bounded scan budget and a
    cursor over the last inspected document so filtering cannot stall progress.
    """
    if limit < 1 or limit > 100:
        raise ValueError('reconcile page limit must be between 1 and 100')

    user_ref = db.collection('users').document(uid)
    messages_collection = user_ref.collection('messages')
    query: Any = messages_collection
    if chat_session_id:
        query = query.where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
    else:
        query = query.where(filter=FieldFilter('plugin_id', '==', app_id))
    query = query.order_by('created_at', direction=firestore.Query.DESCENDING)

    cursor_snapshot: Any = None
    if cursor_message_id:
        cursor_snapshot = messages_collection.document(cursor_message_id).get()
        if not cursor_snapshot.exists:
            raise MessageReconcileCursorError('message cursor does not exist')
        cursor_payload = _typed_doc(cursor_snapshot)
        cursor_in_scope = (
            cursor_payload.get('chat_session_id') == chat_session_id
            if chat_session_id
            else cursor_payload.get('plugin_id') == app_id
        )
        if not cursor_in_scope or cursor_payload.get('created_at') is None:
            raise MessageReconcileCursorError('message cursor is outside the requested scope')

    # Four pages of reported rows may be traversed per request. The returned
    # cursor lets a caller resume if the bounded budget is exhausted.
    scan_budget = min(1000, max(100, limit * 4))
    scanned = 0
    messages: List[Dict[str, Any]] = []
    next_cursor = cursor_message_id
    has_more = False

    while scanned < scan_budget and len(messages) < limit:
        batch_limit = min(100, scan_budget - scanned)
        page_query = query.start_after(cursor_snapshot) if cursor_snapshot is not None else query
        documents = list(page_query.limit(batch_limit).stream())
        if not documents:
            has_more = False
            break

        reached_return_limit = False
        for document in documents:
            scanned += 1
            cursor_snapshot = document
            next_cursor = str(document.id)
            message = _typed_doc(document)
            if message.get('reported') is not True:
                messages.append(message)
                if len(messages) == limit:
                    reached_return_limit = True
                    break

        if reached_return_limit:
            # An extra empty page is harmless when the last returned row was the
            # collection tail; claiming continuation avoids a racey count query.
            has_more = True
            break
        if len(documents) < batch_limit:
            has_more = False
            break
        has_more = True

    if scanned >= scan_budget and len(messages) < limit:
        has_more = True
    if next_cursor == cursor_message_id and not messages:
        next_cursor = None
    return messages, next_cursor, has_more


def get_message_count(uid: str) -> int:
    """Return the number of chat messages visible to the user.

    Reported messages are hidden from every chat view (``get_messages`` and ``get_app_messages``
    skip ``reported == True``), so this stat excludes them too; otherwise it would exceed the number
    of messages the user can actually see anywhere. Uses count() aggregation (total minus the
    reported subset) rather than streaming every message. A ``reported == False`` count would be
    wrong because legacy messages may omit the field.
    """
    messages_ref = db.collection('users').document(uid).collection('messages')
    total_res = messages_ref.count().get()
    total = int(total_res[0][0].value) if total_res and total_res[0] else 0
    reported_res = messages_ref.where(filter=FieldFilter('reported', '==', True)).count().get()
    reported = int(reported_res[0][0].value) if reported_res and reported_res[0] else 0
    return max(0, total - reported)


def iter_all_messages(uid: str, batch_size: int = 1000) -> Iterator[Dict[str, Any]]:
    """Yield all chat messages for a user, decrypted, in batches. Used for streaming data export."""
    user_ref = db.collection('users').document(uid)
    msgs_ref = user_ref.collection('messages').order_by('created_at', direction=firestore.Query.DESCENDING)
    cursor = None
    while True:
        batch_ref = msgs_ref.limit(batch_size)
        if cursor is not None:
            batch_ref = batch_ref.start_after(cursor)
        batch: List[Dict[str, Any]] = []
        snapshots = list(batch_ref.stream())
        for doc in snapshots:
            msg: Dict[str, Any] = _typed_doc(doc)
            msg['id'] = doc.id
            msg = _prepare_message_for_read(msg, uid) or msg
            batch.append(msg)
        yield from batch
        if len(snapshots) < batch_size:
            break
        cursor = snapshots[-1]


def get_message(uid: str, message_id: str) -> tuple[Message, str] | None:
    user_ref = db.collection('users').document(uid)
    message_ref = user_ref.collection('messages').where('id', '==', message_id).limit(1).stream()
    message_doc = next(message_ref, None)
    if not message_doc:
        return None

    message = parse_snapshot_or_none(
        Message,
        message_doc,
        payload_from_snapshot=lambda snapshot: _prepare_message_for_read(_typed_doc(snapshot), uid),
    )
    if message is None:
        return None

    return message, message_doc.id


def report_message(uid: str, msg_doc_id: str) -> Dict[str, str]:
    user_ref = db.collection('users').document(uid)
    message_ref = user_ref.collection('messages').document(msg_doc_id)
    try:
        message_ref.update({'reported': True})
        return {"message": "Message reported"}
    except Exception as e:
        logger.error(f"Update failed: {e}")
        return {"message": f"Update failed: {e}"}


def update_message_rating(uid: str, message_id: str, rating: Optional[int]) -> Optional[Dict[str, Any]]:
    """
    Update the rating on a message document.

    Returns the already-streamed message snapshot on success so analytics can
    copy identifiers (notification kind, app_id) without a second Firestore read.
    Returns None if the message does not exist.
    """
    user_ref = db.collection('users').document(uid)
    message_ref = user_ref.collection('messages').where('id', '==', message_id).limit(1).stream()
    message_doc = next(message_ref, None)
    if not message_doc:
        logger.warning(f"⚠️ Message {message_id} not found for user {uid}")
        return None

    snapshot = _typed_doc(message_doc)
    try:
        user_ref.collection('messages').document(message_doc.id).update({'rating': rating})
        logger.info(f"✅ Updated message {message_id} rating to {rating}")
        snapshot['rating'] = rating
        return snapshot
    except Exception as e:
        logger.error(f"❌ Failed to update message rating: {e}")
        return None


def batch_delete_messages(
    parent_doc_ref: Any, batch_size: int = 450, app_id: Optional[str] = None, chat_session_id: Optional[str] = None
) -> None:
    messages_ref = parent_doc_ref.collection('messages')
    messages_ref = messages_ref.where(filter=FieldFilter('plugin_id', '==', app_id))
    if chat_session_id:
        messages_ref = messages_ref.where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
    logger.info(f'batch_delete_messages {app_id}')

    while True:
        docs_stream = messages_ref.limit(batch_size).stream()
        docs_list: List[Any] = list(docs_stream)

        if not docs_list:
            logger.info("No more messages to delete")
            break

        batch = db.batch()
        for doc in docs_list:
            batch.delete(doc.reference)
        batch.commit()

        logger.info(f'Deleted {len(docs_list)} messages')

        if len(docs_list) < batch_size:
            logger.info("Processed all messages")
            break


def clear_chat(
    uid: str, app_id: Optional[str] = None, chat_session_id: Optional[str] = None
) -> Optional[Dict[str, str]]:
    try:
        user_ref = db.collection('users').document(uid)
        logger.info(f"Deleting messages for user: {uid}")
        if not user_ref.get().exists:
            return {"message": "User not found"}
        batch_delete_messages(user_ref, app_id=app_id, chat_session_id=chat_session_id)
        return None
    except Exception as e:
        return {"message": str(e)}


def add_multi_files(uid: str, files_data: List[Dict[str, Any]]) -> None:
    batch = db.batch()
    user_ref = db.collection('users').document(uid)

    for file_data in files_data:
        file_ref = user_ref.collection('files').document(file_data['id'])
        batch.set(file_ref, file_data)

    batch.commit()


def get_chat_files(uid: str, files_id: Optional[List[str]] = None) -> List[Dict[str, Any]]:
    files_ref = db.collection('users').document(uid).collection('files')

    if files_id is None:
        files_id = []

    # If no specific files requested, return all
    if len(files_id) == 0:
        return [_typed_doc(doc) for doc in files_ref.stream()]

    # Firestore IN operator supports max 30 values, so chunk the queries
    if len(files_id) <= 30:
        files_ref = files_ref.where(filter=FieldFilter('id', 'in', files_id))
        return [_typed_doc(doc) for doc in files_ref.stream()]

    # Chunk into batches of 30
    results: List[Dict[str, Any]] = []
    for i in range(0, len(files_id), 30):
        chunk = files_id[i : i + 30]
        chunk_ref = db.collection('users').document(uid).collection('files')
        chunk_ref = chunk_ref.where(filter=FieldFilter('id', 'in', chunk))
        results.extend([_typed_doc(doc) for doc in chunk_ref.stream()])

    return results


def get_chat_files_desc(uid: str, files_id: Optional[List[str]] = None, limit: int = 10) -> List[Dict[str, Any]]:
    """Get the most recent chat files ordered by created_at descending, optionally filtered by file IDs"""
    files_ref = db.collection('users').document(uid).collection('files')

    if files_id is None:
        files_id = []

    # If no specific files requested, return most recent files
    if len(files_id) == 0:
        files_ref = files_ref.order_by('created_at', direction=firestore.Query.DESCENDING).limit(limit)
        return [_typed_doc(doc) for doc in files_ref.stream()]

    # If specific files requested, filter by them first
    # Firestore IN operator supports max 30 values
    if len(files_id) <= 30:
        files_ref = files_ref.where(filter=FieldFilter('id', 'in', files_id))
        files_ref = files_ref.order_by('created_at', direction=firestore.Query.DESCENDING).limit(limit)
        return [_typed_doc(doc) for doc in files_ref.stream()]

    # Chunk into batches of 30 if more than 30 files
    results: List[Dict[str, Any]] = []
    for i in range(0, len(files_id), 30):
        chunk = files_id[i : i + 30]
        chunk_ref = db.collection('users').document(uid).collection('files')
        chunk_ref = chunk_ref.where(filter=FieldFilter('id', 'in', chunk))
        chunk_ref = chunk_ref.order_by('created_at', direction=firestore.Query.DESCENDING)
        results.extend([_typed_doc(doc) for doc in chunk_ref.stream()])

    # Sort all results by created_at and limit. Use a tz-aware sentinel for a missing created_at so it
    # never TypeError-compares against the tz-aware Firestore datetimes and sinks to the bottom of the
    # descending sort (same class as the review-queue tz sentinel in #9571).
    results.sort(key=lambda x: x.get('created_at', datetime.min.replace(tzinfo=timezone.utc)), reverse=True)
    return results[:limit]


def delete_multi_files(uid: str, files_data: List[Dict[str, Any]]) -> None:
    batch = db.batch()
    user_ref = db.collection('users').document(uid)

    for file_data in files_data:
        file_ref = user_ref.collection('files').document(file_data["id"])
        batch.delete(file_ref)

    batch.commit()


def add_chat_session(uid: str, chat_session_data: Dict[str, Any]) -> Dict[str, Any]:
    user_ref = db.collection('users').document(uid)
    user_ref.collection('chat_sessions').document(chat_session_data['id']).set(chat_session_data)
    return chat_session_data


def get_chat_session(uid: str, app_id: Optional[str] = None) -> Optional[Dict[str, Any]]:
    """The user's current chat session for an app.

    Newest-first: an unordered `.limit(1)` lets Firestore return any matching
    document, so once a user has more than one session for an app the "current"
    one is whichever the index happens to yield. Callers treat this as the
    session to read and append to, so an arbitrary pick silently splits a
    conversation across sessions.

    The ordering is applied after the read, not by `order_by`, because Firestore
    drops documents that lack the ordered field entirely. `add_chat_session`
    writes whatever dict it is handed, so a session with no `created_at` is
    representable — and ordering in the query would make those sessions
    invisible here, stranding a user's existing history behind a brand new
    session. A session with no timestamp sorts oldest, and its id breaks ties so
    the answer is stable across calls.
    """
    collection = db.collection('users').document(uid).collection('chat_sessions')
    ordered_sessions = (
        CURRENT_CHAT_SESSION_ORDERED_QUERY.build(
            collection,
            {'app_id': app_id},
            field_filter_factory=FieldFilter,
        )
        .order_by('created_at', direction=firestore.Query.DESCENDING)
        .order_by('__name__', direction=firestore.Query.DESCENDING)
        .limit(1)
        .stream()
    )
    ordered_docs = [_typed_doc(session) for session in ordered_sessions]
    if ordered_docs:
        return max(
            ordered_docs,
            key=lambda data: (
                data.get('created_at') is not None,
                data.get('created_at') or datetime.min.replace(tzinfo=timezone.utc),
                str(data.get('id') or ''),
            ),
        )

    legacy_session = (
        CURRENT_CHAT_SESSION_QUERY.build(
            collection,
            {'app_id': app_id},
            field_filter_factory=FieldFilter,
        )
        .order_by('__name__', direction=firestore.Query.ASCENDING)
        .limit(1)
        .stream()
    )

    legacy_docs = [_typed_doc(session) for session in legacy_session]
    if len(legacy_docs) > 1:
        legacy_docs = legacy_docs[:1]

    newest: Optional[Dict[str, Any]] = None
    newest_key: Optional[tuple] = None
    for data in legacy_docs:
        # `_typed_doc` returns {} for a document with no fields, which sorts as
        # untimestamped and loses to anything real rather than being skipped.
        created = data.get('created_at')
        key = (created is not None, created or datetime.min.replace(tzinfo=timezone.utc), str(data.get('id') or ''))
        if newest_key is None or key > newest_key:
            newest, newest_key = data, key

    return newest


def get_chat_session_by_id(uid: str, chat_session_id: str) -> Optional[Dict[str, Any]]:
    """Get a specific chat session by its ID"""
    user_ref = db.collection('users').document(uid)
    session_ref = user_ref.collection('chat_sessions').document(chat_session_id)
    session_doc = session_ref.get()

    if session_doc.exists:
        data = session_doc.to_dict()
        data['id'] = chat_session_id
        return _normalize_chat_session(data)

    return None


def delete_chat_session(uid: str, chat_session_id: str, cascade_messages: bool = False) -> Optional[bool]:
    user_ref = db.collection('users').document(uid)
    session_ref = user_ref.collection('chat_sessions').document(chat_session_id)

    if cascade_messages:
        if not session_ref.get().exists:
            return False
        msg_col = user_ref.collection('messages')
        query = msg_col.where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
        while True:
            docs: List[Any] = list(query.limit(BATCH_LIMIT).stream())
            if not docs:
                break
            batch = db.batch()
            for doc in docs:
                batch.delete(msg_col.document(doc.id))
            batch.commit()

    session_ref.delete()
    return None


def _update_chat_session_if_exists(uid: str, chat_session_id: str, values: Dict[str, Any], what: str) -> bool:
    """Apply a derived-state update to a chat session, tolerating a deleted session.

    The message/file id lists and the OpenAI ids are derived state the session
    document owns. Every writer below runs after a multi-second LLM call, and
    DELETE /v2/messages (clear chat) deletes the session it read at the start of
    that same window — so a concurrent clear, or a client retrying the slow
    request, leaves these writes pointing at a tombstone. Firestore's update()
    then raises NotFound and the user's chat call 500s even though the work it
    was reporting already succeeded. A session that no longer exists has nothing
    to record.

    Returns True when the update was applied.
    """
    session_ref = db.collection('users').document(uid).collection('chat_sessions').document(chat_session_id)
    try:
        session_ref.update(values)
        return True
    except NotFound:
        logger.warning(f"chat session {chat_session_id} no longer exists; skipping {what}")
        return False


def add_message_to_chat_session(uid: str, chat_session_id: str, message_id: str) -> None:
    _update_chat_session_if_exists(
        uid, chat_session_id, {"message_ids": firestore.ArrayUnion([message_id])}, "message link"
    )


def add_files_to_chat_session(uid: str, chat_session_id: str, file_ids: List[str]) -> None:
    if not file_ids:
        return

    _update_chat_session_if_exists(uid, chat_session_id, {"file_ids": firestore.ArrayUnion(file_ids)}, "file link")


# **************************************
# ********* MIGRATION HELPERS **********
# **************************************


def get_chats_to_migrate(uid: str, target_level: str) -> List[Dict[str, Any]]:
    """
    Finds all chat messages that are not at the target protection level by fetching all documents
    and filtering them in memory. This simplifies the code but may be less performant for
    users with a very large number of documents.
    """
    messages_ref = db.collection('users').document(uid).collection('messages')
    all_messages = messages_ref.select(['data_protection_level']).stream()

    to_migrate: List[Dict[str, Any]] = []
    for doc in all_messages:
        doc_data: Dict[str, Any] = _typed_doc(doc)
        current_level = doc_data.get('data_protection_level', 'standard')
        if target_level != current_level:
            to_migrate.append({'id': doc.id, 'type': 'chat'})

    return to_migrate


def migrate_chats_level_batch(uid: str, message_doc_ids: List[str], target_level: str) -> None:
    """
    Migrates a batch of chat messages to the target protection level.
    """
    batch = db.batch()
    messages_ref = db.collection('users').document(uid).collection('messages')
    doc_refs = [messages_ref.document(msg_id) for msg_id in message_doc_ids]
    doc_snapshots = db.get_all(doc_refs)

    for doc_snapshot in doc_snapshots:
        if not doc_snapshot.exists:
            logger.warning(f"Message {doc_snapshot.id} not found, skipping.")
            continue

        message_data: Dict[str, Any] = _typed_doc(doc_snapshot)
        current_level = message_data.get('data_protection_level', 'standard')

        if current_level == target_level:
            continue

        plain_data: Dict[str, Any] = _prepare_message_for_read(message_data, uid)
        plain_text = plain_data.get('text')
        migrated_text = plain_text
        if target_level == 'enhanced':
            if isinstance(plain_text, str):
                migrated_text = encryption.encrypt(plain_text, uid)

        update_data: Dict[str, Any] = {'data_protection_level': target_level, 'text': migrated_text}
        batch.update(doc_snapshot.reference, update_data)

    batch.commit()


# ============================================================================
# CHAT SESSIONS (v2)
#
# v2 sessions support: title, preview, message_count, starred, updated_at.
# v1 sessions store: message_ids, file_ids (legacy docs may still carry openai_thread_id).
# Both schemas coexist in the same Firestore collection.
# Both MUST write plugin_id alongside app_id for cross-platform query compat.
# ============================================================================


def _normalize_chat_session(data: Optional[dict]) -> Optional[dict]:
    """Guarantee a v2 chat-session dict satisfies ``ChatSessionResponse``.

    Firestore holds sessions written by several code paths (Python v2, the Rust
    desktop backend, legacy docs). Some rows are missing fields the response
    model requires (``title``, ``created_at``, ``message_count``, ``starred``),
    which makes FastAPI raise ``ResponseValidationError`` (HTTP 500). Fill safe
    defaults so listing/reading sessions never 500 on an incomplete doc.
    """
    if data is None:
        return None
    data.setdefault('title', 'New Chat')
    data.setdefault('preview', None)
    data.setdefault('message_count', 0)
    data.setdefault('starred', False)
    # created_at/updated_at are required datetimes; fall back to each other when
    # one is missing (the list query orders by updated_at, so it is present there).
    if data.get('created_at') is None:
        data['created_at'] = data.get('updated_at') or datetime.now(timezone.utc)
    if data.get('updated_at') is None:
        data['updated_at'] = data.get('created_at') or datetime.now(timezone.utc)
    return data


def create_chat_session(uid: str, title: Optional[str] = None, app_id: Optional[str] = None) -> Dict[str, Any]:
    session_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc)
    doc: Dict[str, Any] = {
        'id': session_id,
        'title': title or 'New Chat',
        'preview': None,
        'created_at': now,
        'updated_at': now,
        'app_id': app_id,
        'plugin_id': app_id,  # Python chat.py queries chat_sessions by plugin_id
        'message_count': 0,
        'starred': False,
    }
    db.collection('users').document(uid).collection('chat_sessions').document(session_id).set(doc)
    return doc


def acquire_chat_session(uid: str, app_id: Optional[str] = None) -> str:
    """Get or create a chat session for the given app_id (None = main chat).

    Queries by plugin_id to match both Python chat.py and Rust backend behavior.
    For main chat (app_id=None), matches sessions where plugin_id is None.
    """
    session = get_chat_session(uid, app_id=app_id)
    if session:
        return session['id']
    session = create_chat_session(uid, app_id=app_id)
    return session['id']


def get_chat_sessions(
    uid: str,
    app_id: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
    starred: Optional[bool] = None,
) -> List[Dict[str, Any]]:
    col = db.collection('users').document(uid).collection('chat_sessions')
    # Order by updated_at — v2 sessions always have this field.
    # Legacy v1 sessions (missing updated_at) are excluded by Firestore,
    # which is correct since this endpoint serves v2 clients only.
    query = col.order_by('updated_at', direction=firestore.Query.DESCENDING)

    # Always filter — when app_id is None this returns only default-chat sessions
    query = query.where(filter=FieldFilter('plugin_id', '==', app_id))
    if starred is not None:
        query = query.where(filter=FieldFilter('starred', '==', starred))

    query = query.offset(offset).limit(limit)
    items: List[Dict[str, Any]] = []
    for doc in query.stream():
        data: Dict[str, Any] = _typed_doc(doc)
        data['id'] = doc.id
        normalized = _normalize_chat_session(data)
        if normalized is not None:
            items.append(normalized)
    return items


def update_chat_session(
    uid: str,
    session_id: str,
    title: Optional[str] = None,
    starred: Optional[bool] = None,
) -> Optional[Dict[str, Any]]:
    ref = db.collection('users').document(uid).collection('chat_sessions').document(session_id)
    if not ref.get().exists:
        return None
    updates: Dict[str, Any] = {'updated_at': datetime.now(timezone.utc)}
    if title is not None:
        updates['title'] = title
    if starred is not None:
        updates['starred'] = starred
    ref.update(updates)
    result: Dict[str, Any] = _typed_doc(ref.get())
    result['id'] = session_id
    return _normalize_chat_session(result)


# ============================================================================
# MESSAGES (v2)
#
# Persistence-only message writes (no LLM streaming).  They write the same
# field set as the Message model for cross-platform compatibility:
#   plugin_id, app_id, type='text', chat_session_id, from_external_integration
#
# When session_id is not provided, acquire_chat_session() auto-creates one.
# ============================================================================


def save_message(
    uid: str,
    text: str,
    sender: str,
    app_id: Optional[str] = None,
    session_id: Optional[str] = None,
    metadata: Optional[str] = None,
    content_blocks: Optional[List[Dict[str, Any]]] = None,
    client_message_id: Optional[str] = None,
    message_source: str = 'desktop_chat',
    journal_revision: Optional[int] = None,
) -> Dict[str, Any]:
    """Save a chat message for the desktop app.

    Writes all fields expected by chat.py's Message model so messages are
    visible across platforms.  Auto-acquires a session if none provided.
    """
    msg_id = client_message_id or str(uuid.uuid4())
    now = datetime.now(timezone.utc)
    requested_session_id = session_id
    idempotency_payload_hash = _message_idempotency_payload_hash(
        text=text,
        sender=sender,
        app_id=app_id,
        session_id=requested_session_id,
        metadata=metadata,
        content_blocks=content_blocks,
        message_source=message_source,
    )

    message_ref = db.collection('users').document(uid).collection('messages').document(msg_id)
    if client_message_id:
        existing_message = message_ref.get()
        if existing_message.exists:
            existing_result = _apply_existing_message_revision(
                message_ref,
                text=text,
                sender=sender,
                app_id=app_id,
                session_id=requested_session_id,
                metadata=metadata,
                content_blocks=content_blocks,
                message_source=message_source,
                payload_hash=idempotency_payload_hash,
                journal_revision=journal_revision,
            )
            if existing_result is not None:
                return _message_revision_response(msg_id, existing_result, now)

    # Auto-acquire session (matches Rust backend behavior)
    if not session_id:
        session_id = acquire_chat_session(uid, app_id=app_id)

    doc: Dict[str, Any] = {
        'id': msg_id,
        'text': text,
        'created_at': now,
        'sender': sender,
        'type': 'text',  # Desktop messages are always type 'text'
        'app_id': app_id,
        'plugin_id': app_id,  # chat.py queries messages by plugin_id
        'session_id': session_id,
        'chat_session_id': session_id,  # chat.py uses this field name
        'from_external_integration': False,
        'rating': None,
        'reported': False,
        'memories_id': [],
        'metadata': metadata,
        'message_source': message_source,
    }
    if content_blocks is not None:
        doc['content_blocks'] = content_blocks
    if client_message_id:
        doc['client_message_id'] = client_message_id
        doc['client_message_payload_hash'] = idempotency_payload_hash
        if journal_revision is not None:
            doc['journal_revision'] = journal_revision
    created = True
    if client_message_id:
        try:
            message_ref.create(doc)
        except (AlreadyExists, Conflict):
            existing_result = _apply_existing_message_revision(
                message_ref,
                text=text,
                sender=sender,
                app_id=app_id,
                session_id=requested_session_id,
                metadata=metadata,
                content_blocks=content_blocks,
                message_source=message_source,
                payload_hash=idempotency_payload_hash,
                journal_revision=journal_revision,
            )
            if existing_result is None:
                raise ClientMessageIdPayloadConflict('client_message_id disappeared during revision arbitration')
            return _message_revision_response(msg_id, existing_result, now)
    else:
        message_ref.set(doc)

    # Update session message_count and preview (skip if session was deleted).
    # Retried client_message_id saves are idempotent and must not bump counters.
    if session_id and created:
        session_ref = db.collection('users').document(uid).collection('chat_sessions').document(session_id)
        if session_ref.get().exists:
            session_ref.update(
                {
                    'updated_at': now,
                    'message_count': firestore.Increment(1),
                    'preview': text[:100] if text else None,
                }
            )

    return {
        'id': msg_id,
        'created_at': now.isoformat(),
        'session_id': session_id,
        'created': created,
        'updated': False,
        'journal_revision': journal_revision,
    }


def _apply_existing_message_revision(
    message_ref: Any,
    *,
    text: str,
    sender: str,
    app_id: Optional[str],
    session_id: Optional[str],
    metadata: Optional[str],
    content_blocks: Optional[List[Dict[str, Any]]],
    message_source: str,
    payload_hash: str,
    journal_revision: Optional[int],
) -> Optional[Dict[str, Any]]:
    """Atomically arbitrate an idempotent retry or monotonic journal enrichment."""
    transaction = db.transaction()

    @firestore_v1.transactional
    def apply(write_transaction: Any) -> Optional[Dict[str, Any]]:
        snapshot = message_ref.get(transaction=write_transaction)
        if not snapshot.exists:
            return None
        existing = _typed_doc(snapshot)
        _assert_message_identity(
            existing,
            sender=sender,
            app_id=app_id,
            session_id=session_id,
            message_source=message_source,
        )
        stored_revision = int(existing.get('journal_revision') or 1)
        if journal_revision is None:
            _assert_idempotent_message_payload(
                existing,
                text=text,
                sender=sender,
                app_id=app_id,
                session_id=session_id,
                metadata=metadata,
                content_blocks=content_blocks,
                message_source=message_source,
                payload_hash=payload_hash,
            )
            existing['_revision_updated'] = False
            existing['journal_revision'] = existing.get('journal_revision')
            return existing
        if journal_revision < stored_revision:
            existing['_revision_updated'] = False
            return existing
        if journal_revision == stored_revision:
            _assert_idempotent_message_payload(
                existing,
                text=text,
                sender=sender,
                app_id=app_id,
                session_id=session_id,
                metadata=metadata,
                content_blocks=content_blocks,
                message_source=message_source,
                payload_hash=payload_hash,
            )
            existing['_revision_updated'] = False
            return existing
        patch: Dict[str, Any] = {
            'text': text,
            'metadata': metadata,
            'client_message_payload_hash': payload_hash,
            'journal_revision': journal_revision,
        }
        if content_blocks is not None:
            patch['content_blocks'] = content_blocks
        write_transaction.update(message_ref, patch)
        existing.update(patch)
        existing['_revision_updated'] = True
        return existing

    return apply(transaction)


def _assert_message_identity(
    existing: Dict[str, Any],
    *,
    sender: str,
    app_id: Optional[str],
    session_id: Optional[str],
    message_source: str,
) -> None:
    mismatched = existing.get('sender') != sender
    mismatched = mismatched or existing.get('message_source', 'desktop_chat') != message_source
    existing_app_ids = [existing[field] for field in ('app_id', 'plugin_id') if field in existing] or [None]
    mismatched = mismatched or any(existing_app_id != app_id for existing_app_id in existing_app_ids)
    if session_id is not None:
        existing_session_ids = [
            existing[field] for field in ('chat_session_id', 'session_id') if field in existing
        ] or [None]
        mismatched = mismatched or any(
            existing_session_id != session_id for existing_session_id in existing_session_ids
        )
    if mismatched:
        raise ClientMessageIdPayloadConflict('client_message_id already exists with a different canonical identity')


def _message_revision_response(msg_id: str, existing: Dict[str, Any], now: datetime) -> Dict[str, Any]:
    existing_created_at = existing.get('created_at')
    if existing_created_at is not None and hasattr(existing_created_at, 'isoformat'):
        existing_created_at = existing_created_at.isoformat()
    return {
        'id': msg_id,
        'created_at': existing_created_at or now.isoformat(),
        'session_id': existing.get('chat_session_id') or existing.get('session_id'),
        'created': False,
        'updated': bool(existing.get('_revision_updated')),
        'journal_revision': existing.get('journal_revision'),
    }


def _assert_idempotent_message_payload(
    existing: Dict[str, Any],
    *,
    text: str,
    sender: str,
    app_id: Optional[str],
    session_id: Optional[str],
    metadata: Optional[str],
    message_source: str,
    payload_hash: str,
    content_blocks: Optional[List[Dict[str, Any]]] = None,
) -> None:
    """Reject an idempotency-key collision without exposing message content."""
    existing_payload_hash = existing.get('client_message_payload_hash')
    if existing_payload_hash is not None:
        if existing_payload_hash != payload_hash:
            raise ClientMessageIdPayloadConflict('client_message_id already exists with a different payload')
        return

    # Legacy messages predate the request-payload fingerprint. Preserve retries
    # for them while validating every field whose original request value can be
    # reconstructed from the stored row. New writes always use the exact hash.
    mismatched = existing.get('text') != text or existing.get('sender') != sender
    mismatched = mismatched or existing.get('metadata') != metadata
    if content_blocks is not None:
        mismatched = mismatched or existing.get('content_blocks') != content_blocks
    mismatched = mismatched or existing.get('message_source', 'desktop_chat') != message_source
    existing_app_ids = [existing[field] for field in ('app_id', 'plugin_id') if field in existing] or [None]
    mismatched = mismatched or any(existing_app_id != app_id for existing_app_id in existing_app_ids)
    if session_id is not None:
        existing_session_ids = [
            existing[field] for field in ('chat_session_id', 'session_id') if field in existing
        ] or [None]
        mismatched = mismatched or any(
            existing_session_id != session_id for existing_session_id in existing_session_ids
        )
    if mismatched:
        raise ClientMessageIdPayloadConflict('client_message_id already exists with a different payload')


def _message_idempotency_payload_hash(
    *,
    text: str,
    sender: str,
    app_id: Optional[str],
    session_id: Optional[str],
    metadata: Optional[str],
    message_source: str,
    content_blocks: Optional[List[Dict[str, Any]]] = None,
) -> str:
    """Return a stable digest of the caller-controlled immutable payload."""
    payload: Dict[str, Any] = {
        'app_id': app_id,
        'message_source': message_source,
        'metadata': metadata,
        'sender': sender,
        'session_id': session_id,
        'text': text,
    }
    if content_blocks is not None:
        payload['content_blocks'] = content_blocks
    canonical = json.dumps(payload, ensure_ascii=False, separators=(',', ':'), sort_keys=True)
    return f'sha256:{hashlib.sha256(canonical.encode("utf-8")).hexdigest()}'


def delete_messages(uid: str, app_id: Optional[str] = None, session_id: Optional[str] = None) -> int:
    """Delete messages and apply inverse session metadata updates atomically."""
    user_ref = db.collection('users').document(uid)
    col = user_ref.collection('messages')
    if session_id:
        # Session-scoped delete: filter by session only (same logic as get_messages)
        query = col.where(filter=FieldFilter('chat_session_id', '==', session_id))
    else:
        # App-scoped delete: filter by plugin_id (None = main chat)
        query = col.where(filter=FieldFilter('plugin_id', '==', app_id))

    deleted = 0
    consecutive_conflicts = 0
    while True:
        docs: List[Any] = list(query.limit(DELETE_MESSAGES_BATCH_LIMIT).stream())
        if not docs:
            break

        deleted_by_session: Dict[str, int] = {}
        deleted_message_ids_by_session: Dict[str, List[str]] = {}
        deleted_previews_by_session: Dict[str, set[str]] = {}
        for doc in docs:
            data = _typed_doc(doc)
            message_session_id = data.get('chat_session_id') or data.get('session_id')
            if isinstance(message_session_id, str) and message_session_id:
                deleted_by_session[message_session_id] = deleted_by_session.get(message_session_id, 0) + 1
                stored_message_id = data.get('id')
                deleted_message_ids_by_session.setdefault(message_session_id, []).append(
                    stored_message_id if isinstance(stored_message_id, str) and stored_message_id else doc.id
                )
                text = data.get('text')
                if isinstance(text, str) and text:
                    deleted_previews_by_session.setdefault(message_session_id, set()).add(text[:100])

        session_snapshots: Dict[str, Any] = {}
        for message_session_id in deleted_by_session:
            session_snapshots[message_session_id] = (
                user_ref.collection('chat_sessions').document(message_session_id).get()
            )

        batch = db.batch()
        for doc in docs:
            delete_option = db.write_option(last_update_time=doc.update_time)
            batch.delete(col.document(doc.id), option=delete_option)

        for message_session_id, deleted_from_session in deleted_by_session.items():
            session_snapshot = session_snapshots[message_session_id]
            if not session_snapshot.exists:
                continue
            session_data = _typed_doc(session_snapshot)
            stored_count = session_data.get('message_count')
            updates: Dict[str, Any] = {}
            if isinstance(session_data.get('message_ids'), list):
                updates['message_ids'] = firestore.ArrayRemove(deleted_message_ids_by_session[message_session_id])
            if isinstance(stored_count, int) and stored_count > 0:
                decrement = min(stored_count, deleted_from_session)
                updates['message_count'] = firestore.Increment(-decrement)
            current_preview = session_data.get('preview')
            if isinstance(current_preview, str) and current_preview in deleted_previews_by_session.get(
                message_session_id, set()
            ):
                updates['preview'] = None
            if not updates:
                continue
            session_ref = user_ref.collection('chat_sessions').document(message_session_id)
            option = db.write_option(last_update_time=session_snapshot.update_time)
            batch.update(session_ref, updates, option=option)

        try:
            batch.commit()
        except FailedPrecondition:
            # Another clear or a concurrent session mutation won the race.
            # The batch is atomic, so re-query before applying any decrement.
            consecutive_conflicts += 1
            if consecutive_conflicts >= DELETE_MESSAGES_CONFLICT_RETRIES:
                raise
            continue
        deleted += len(docs)
        consecutive_conflicts = 0

    return deleted