memory_ledger.py

import copy
import hashlib
import json
from datetime import datetime, timezone
from functools import wraps
from typing import Any, Callable, Dict, List, Optional, TypeVar, cast

from google.cloud.firestore_v1 import transactional  # type: ignore[reportUnknownMemberType]  # firestore SDK stub gap

from database import _client as _client_mod
from database import projection_repair
from database.account_deletion_policy import account_deletion_blocks_access, normalize_account_deletion_status
from database.firestore_transaction_retry import run_with_transaction_contention_retry
from database.legal_holds import assert_no_destructive_operation_transaction
from database.memory_collections import MemoryCollections
from models.memories import confidence_fields_for_evidence
from models.memory_state_head import (
    trusted_memory_state_head_fields_from_control,
    trusted_memory_state_head_fields_from_state,
)

db = getattr(_client_mod, 'data_plane_db', _client_mod.db)

T = TypeVar("T")


def _typed_transactional(func: Callable[..., T]) -> Callable[..., T]:
    """Create an isolated SDK transaction wrapper for every invocation.

    Firestore's ``_Transactional`` wrapper stores mutable retry IDs. Reusing one
    module-level instance across request threads can cross-contaminate retries,
    so defer constructing it until each call while preserving the typed surface.
    """

    @wraps(func)
    def invoke(transaction: Any, *args: Any, **kwargs: Any) -> T:
        wrapped = cast(Callable[..., T], transactional(func))
        return wrapped(transaction, *args, **kwargs)

    return invoke


def _typed_doc(doc: Any) -> Dict[str, Any]:
    """Typed adapter for Firestore DocumentSnapshot.to_dict() (SDK stub gap)."""
    raw: object = doc.to_dict()
    return cast(Dict[str, Any], raw) if isinstance(raw, dict) else {}


users_collection = 'users'
memory_state_collection = 'memory_state'
memory_state_document = 'head'
memory_apply_control_document = 'apply_control'
memory_commits_collection = 'memory_commits'
MEMORY_COMMIT_PRIVACY_PURGE_PAGE_SIZE = 100


class HeadConflict(Exception):
    def __init__(self, expected_parent: Optional[str], current_head: Optional[str]):
        super().__init__(f"Memory ledger head moved from {expected_parent} to {current_head}")
        self.expected_parent = expected_parent
        self.current_head = current_head


class LegacyCommitPrivacyFence(RuntimeError):
    """A legacy ledger mutation cannot prove it is safe after deletion."""


def _document(database: Any, path: str) -> Any:
    document = getattr(database, 'document', None)
    if callable(document):
        return document(path)
    collection_name, document_id = path.split('/', 1)
    return database.collection(collection_name).document(document_id)


def _mutation_memory_ids(mutations: List[Dict[str, Any]]) -> set[str]:
    ids: set[str] = set()
    for entry in mutations:
        fact = entry.get('fact')
        if isinstance(fact, dict):
            fact_id = fact.get('id')
            if isinstance(fact_id, str) and fact_id:
                ids.add(fact_id)
        for key in ('fact_id', 'target_fact_id', 'by'):
            value = entry.get(key)
            if isinstance(value, str) and value:
                ids.add(value)
    return ids


def _assert_legacy_commit_privacy_fences(
    *,
    transaction: Any,
    database: Any,
    uid: str,
    mutations: List[Dict[str, Any]],
) -> None:
    """Serialize legacy commits with deletion and reject retired identities."""

    assert_no_destructive_operation_transaction(transaction, database, uid=uid)
    deletion_ref = _document(database, f'account_deletions/{uid}')
    deletion_snapshot = deletion_ref.get(transaction=transaction)
    deletion_payload = deletion_snapshot.to_dict() if getattr(deletion_snapshot, 'exists', False) else {}
    deletion_status = normalize_account_deletion_status(
        marker_exists=bool(getattr(deletion_snapshot, 'exists', False)),
        raw_status=deletion_payload.get('wipe_status') if isinstance(deletion_payload, dict) else None,
    )
    if account_deletion_blocks_access(deletion_status):
        raise LegacyCommitPrivacyFence('legacy memory commit blocked by account deletion')

    memory_ids = sorted(_mutation_memory_ids(mutations))
    if len(memory_ids) > 200:
        raise LegacyCommitPrivacyFence('legacy memory commit identity inventory is too large')
    collections = MemoryCollections(uid=uid)
    for memory_id in memory_ids:
        item_snapshot = _document(database, f'{collections.memory_items}/{memory_id}').get(transaction=transaction)
        item_exists = getattr(item_snapshot, 'exists', False)
        item_payload = item_snapshot.to_dict() if item_exists else None
        if item_exists:
            if not isinstance(item_payload, dict):
                raise LegacyCommitPrivacyFence('legacy memory commit canonical identity is unreadable')
            if item_payload.get('status') == 'tombstoned':
                raise LegacyCommitPrivacyFence('legacy memory commit references a privacy-deleted memory')
            # Still read the override: delete_batch writes suppression before the
            # canonical tombstone, so a live item can already be privacy-deleted.
        override_snapshot = _document(database, f'{collections.memory_historical_overrides}/{memory_id}').get(
            transaction=transaction
        )
        override_exists = getattr(override_snapshot, 'exists', False)
        override_payload = override_snapshot.to_dict() if override_exists else None
        if override_exists and not isinstance(override_payload, dict):
            raise LegacyCommitPrivacyFence('legacy memory commit canonical identity is unreadable')
        if isinstance(override_payload, dict) and override_payload.get('status') == 'tombstoned':
            raise LegacyCommitPrivacyFence('legacy memory commit references a privacy-deleted memory')


def _state_head_write_payload(
    *,
    transaction: Any,
    user_ref: Any,
    uid: str,
    state: Dict[str, Any],
    commit: Dict[str, Any],
) -> Dict[str, Any]:
    """Keep canonical trusted fields when the legacy ledger advances its own head.

    The legacy ledger's ``current_head_commit_id`` is independent from the
    canonical apply stream's ``head_commit_id``.  Older ledger writes replaced
    the shared document and erased the canonical fields.  If an old write has
    already clobbered them, repair only from the transactional canonical control
    record; do not fabricate a trusted generation.
    """
    trusted_fields = trusted_memory_state_head_fields_from_state(state, uid=uid)
    if trusted_fields is None:
        control_ref = user_ref.collection(memory_state_collection).document(memory_apply_control_document)
        control_snapshot = control_ref.get(transaction=transaction)
        control = _typed_doc(control_snapshot) if control_snapshot.exists else {}
        trusted_fields = trusted_memory_state_head_fields_from_control(control, uid=uid)

    payload = {
        'current_head_commit_id': commit['commit_id'],
        'projection_version': 1,
        'updated_at': commit['commit_time'],
    }
    if trusted_fields is not None:
        payload.update(trusted_fields)
    return payload


def mutation(mutation_type: str, **payload: Any) -> Dict[str, Any]:
    return {'type': mutation_type, **payload}


def add_fact(fact: Dict[str, Any]) -> Dict[str, Any]:
    return mutation('add_fact', fact=normalize_fact_for_ledger(fact))


def normalize_fact_for_ledger(fact: Dict[str, Any]) -> Dict[str, Any]:
    normalized = copy.deepcopy(fact)
    qualifiers = normalized.setdefault('qualifiers', {})
    if normalized.get('valid_at') is not None and qualifiers.get('valid_from') is None:
        qualifiers['valid_from'] = normalized.get('valid_at')
    if normalized.get('invalid_at') is not None and qualifiers.get('valid_to') is None:
        qualifiers['valid_to'] = normalized.get('invalid_at')
    return normalized


def supersede_fact(
    fact_id: str,
    by: Optional[str],
    kind: str = 'contradict',
    valid_interval: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    return mutation('supersede_fact', fact_id=fact_id, by=by, kind=kind, valid_interval=valid_interval or {})


def refine_fact(fact_id: str, arg_changes: Dict[str, Any]) -> Dict[str, Any]:
    return mutation('refine_fact', fact_id=fact_id, arg_changes=copy.deepcopy(arg_changes))


def retract_fact(fact_id: str, reason: str = '') -> Dict[str, Any]:
    return mutation('retract_fact', fact_id=fact_id, reason=reason)


def _value_references_memory_id(value: Any, target_ids: set[str]) -> bool:
    if isinstance(value, str):
        return value in target_ids
    if isinstance(value, dict):
        return any(_value_references_memory_id(item, target_ids) for item in value.values())
    if isinstance(value, (list, tuple)):
        return any(_value_references_memory_id(item, target_ids) for item in value)
    return False


def purge_legacy_memory_commits_for_memories(
    uid: str,
    memory_ids: List[str],
    *,
    firestore_client: Any = None,
) -> List[str]:
    """Delete legacy content-bearing commits that reference privacy-deleted IDs.

    Canonical apply commits in the shared collection are content-free and have
    no ``mutations`` field. Legacy commits embed entire facts, corrections, and
    evidence; deleting the document also removes its content-derived ID path.
    The legacy head is retired compatibility state and is never repaired into a
    second authority after this privacy purge.
    """

    target_ids = {memory_id for memory_id in memory_ids if memory_id}
    if not target_ids:
        return []
    database: Any = firestore_client or db
    collection = database.collection(users_collection).document(uid).collection(memory_commits_collection)
    query = collection.order_by('__name__')
    cursor: Any = None
    purged: List[str] = []
    while True:
        page_query = query.start_after(cursor) if cursor is not None else query
        page = list(page_query.limit(MEMORY_COMMIT_PRIVACY_PURGE_PAGE_SIZE).stream())
        if not page:
            break
        for document in page:
            payload = document.to_dict()
            mutations = payload.get('mutations') if isinstance(payload, dict) else None
            if isinstance(mutations, list) and _value_references_memory_id(mutations, target_ids):
                document.reference.delete()
                purged.append(str(document.id))
        if len(page) < MEMORY_COMMIT_PRIVACY_PURGE_PAGE_SIZE:
            break
        cursor = page[-1]
    return purged


def _iter_collection_documents(collection: Any) -> Any:
    query = collection.order_by('__name__')
    cursor: Any = None
    while True:
        page_query = query.start_after(cursor) if cursor is not None else query
        page = list(page_query.limit(MEMORY_COMMIT_PRIVACY_PURGE_PAGE_SIZE).stream())
        if not page:
            return
        yield from page
        if len(page) < MEMORY_COMMIT_PRIVACY_PURGE_PAGE_SIZE:
            return
        cursor = page[-1]


def purge_canonical_privacy_history_for_memories(
    uid: str,
    memory_ids: List[str],
    *,
    firestore_client: Any = None,
    preserve_source_replacement_receipts: bool = False,
) -> Dict[str, List[str]]:
    """Remove pre-delete canonical history whose IDs or payload encode content.

    The new deletion operation/commit/outbox events are content-free privacy
    history and remain. Older operations and their commits/outbox rows are
    removed whole because their document IDs digest the plaintext logical
    payload; in-place redaction would leave a dictionary oracle in the path.
    A mixed legacy commit is removed whole when any mutation references a
    deleted identity. This never removes unrelated canonical memory rows.
    """

    target_ids = {memory_id for memory_id in memory_ids if memory_id}
    if not target_ids:
        return {}
    database: Any = firestore_client or db
    user_ref = database.collection(users_collection).document(uid)
    removed: Dict[str, List[str]] = {}
    deleted_operation_ids: set[str] = set()

    for document in _iter_collection_documents(user_ref.collection('memory_operations')):
        payload = document.to_dict()
        if not isinstance(payload, dict):
            continue
        if _value_references_memory_id(payload, target_ids):
            document.reference.delete()
            deleted_operation_ids.add(str(document.id))
    if deleted_operation_ids:
        removed['memory_operations'] = sorted(deleted_operation_ids)

    removed_commits: List[str] = []
    for document in _iter_collection_documents(user_ref.collection(memory_commits_collection)):
        payload = document.to_dict()
        if not isinstance(payload, dict):
            continue
        if payload.get('operation_id') in deleted_operation_ids or _value_references_memory_id(payload, target_ids):
            document.reference.delete()
            removed_commits.append(str(document.id))
    if removed_commits:
        removed['memory_commits'] = sorted(removed_commits)

    removed_outbox: List[str] = []
    for document in _iter_collection_documents(user_ref.collection('memory_outbox')):
        payload = document.to_dict()
        if isinstance(payload, dict) and (
            payload.get('operation_id') in deleted_operation_ids or _value_references_memory_id(payload, target_ids)
        ):
            document.reference.delete()
            removed_outbox.append(str(document.id))
    if removed_outbox:
        removed['memory_outbox'] = sorted(removed_outbox)

    for collection_name in ('memory_ledger_reopens', 'jit_trigger_feedback', 'jit_proactivity_events'):
        removed_ids: List[str] = []
        for document in _iter_collection_documents(user_ref.collection(collection_name)):
            payload = document.to_dict()
            if isinstance(payload, dict) and _value_references_memory_id(payload, target_ids):
                document.reference.delete()
                removed_ids.append(str(document.id))
        if removed_ids:
            removed[collection_name] = sorted(removed_ids)

    for document in _iter_collection_documents(user_ref.collection('jit_proactivity_daily_budgets')):
        payload = document.to_dict()
        if not isinstance(payload, dict):
            continue
        planned = payload.get('planned_by_trigger')
        if not isinstance(planned, dict) or not target_ids.intersection(planned):
            continue
        document.reference.update(
            {'planned_by_trigger': {key: value for key, value in planned.items() if key not in target_ids}}
        )

    if not preserve_source_replacement_receipts:
        removed_replacements = purge_source_replacement_receipts_for_memories(
            uid,
            list(target_ids),
            firestore_client=database,
        )
        if removed_replacements:
            removed['memory_source_replacements'] = removed_replacements
    return removed


def purge_source_replacement_receipts_for_memories(
    uid: str,
    memory_ids: List[str],
    *,
    firestore_client: Any = None,
) -> List[str]:
    target_ids = {memory_id for memory_id in memory_ids if memory_id}
    if not target_ids:
        return []
    database: Any = firestore_client or db
    collection = database.collection(users_collection).document(uid).collection('memory_source_replacements')
    removed: List[str] = []
    for document in _iter_collection_documents(collection):
        payload = document.to_dict()
        if isinstance(payload, dict) and _value_references_memory_id(payload, target_ids):
            document.reference.delete()
            removed.append(str(document.id))
    return sorted(removed)


@_typed_transactional
def _finalize_canonical_privacy_tombstones_transaction(
    transaction: Any,
    database: Any,
    collections: MemoryCollections,
    target_ids: List[str],
) -> None:
    """Atomically remove one already-bounded tombstoned lineage.

    Tombstone creation is transaction-bounded below Firestore's mutation cap,
    so the corresponding item/evidence/assertion removal also fits.  Keeping
    this final physical step atomic is the retry inventory: a failed commit
    leaves every content-free tombstone in place so any lineage member can
    reconstruct the complete cleanup set on the next request.
    """

    item_refs: List[Any] = []
    evidence_refs: Dict[str, Any] = {}
    assertion_refs: List[Any] = []
    for memory_id in target_ids:
        item_ref = database.document(f"{collections.memory_items}/{memory_id}")
        snapshot = item_ref.get(transaction=transaction)
        payload = snapshot.to_dict() if getattr(snapshot, "exists", False) else None
        if not isinstance(payload, dict):
            continue
        if payload.get("status") != "tombstoned":
            raise RuntimeError("canonical privacy finalization requires a tombstoned item")
        item_refs.append(item_ref)
        assertion_refs.append(database.document(f"{collections.memory_graph_assertions}/{memory_id}"))
        evidence_rows = payload.get("evidence")
        if not isinstance(evidence_rows, list):
            continue
        for evidence in evidence_rows:
            if not isinstance(evidence, dict):
                continue
            evidence_id = evidence.get("evidence_id")
            if not isinstance(evidence_id, str) or not evidence_id or evidence_id in evidence_refs:
                continue
            evidence_ref = database.document(f"{collections.memory_evidence}/{evidence_id}")
            evidence_snapshot = evidence_ref.get(transaction=transaction)
            evidence_payload = evidence_snapshot.to_dict() if getattr(evidence_snapshot, "exists", False) else None
            if isinstance(evidence_payload, dict) and evidence_payload.get("source_state") == "tombstoned":
                evidence_refs[evidence_id] = evidence_ref

    # Firestore forbids reads after writes, so stage deletes only after the
    # complete lineage and evidence validation pass above.
    for evidence_ref in evidence_refs.values():
        transaction.delete(evidence_ref)
    for assertion_ref in assertion_refs:
        transaction.delete(assertion_ref)
    for item_ref in item_refs:
        transaction.delete(item_ref)


def finalize_canonical_privacy_tombstones(
    uid: str,
    memory_ids: List[str],
    *,
    firestore_client: Any = None,
    preserve_source_replacement_receipts: bool = False,
) -> None:
    """Remove deterministic IDs after every derived provider is scrubbed.

    The caller must first prove vector/search/graph absence. This final step
    removes content-derived item/evidence paths and all old history while the
    server-keyed V2 anti-resurrection receipt remains for 30 days.
    """

    target_ids = list(dict.fromkeys(memory_id for memory_id in memory_ids if memory_id))
    if not target_ids:
        return
    database: Any = firestore_client or db
    collections = MemoryCollections(uid=uid)
    purge_canonical_privacy_history_for_memories(
        uid,
        target_ids,
        firestore_client=database,
        preserve_source_replacement_receipts=preserve_source_replacement_receipts,
    )
    transaction = database.transaction()
    _finalize_canonical_privacy_tombstones_transaction(
        transaction,
        database,
        collections,
        target_ids,
    )


def add_evidence(fact_id: str, evidence: Dict[str, Any]) -> Dict[str, Any]:
    return mutation('add_evidence', fact_id=fact_id, evidence=copy.deepcopy(evidence))


def remove_evidence(fact_id: str, evidence_id: str) -> Dict[str, Any]:
    return mutation('remove_evidence', fact_id=fact_id, evidence_id=evidence_id)


def tombstone_evidence(
    fact_id: str,
    evidence_id: str,
    tombstoned_at: Optional[datetime] = None,
    reason: str = 'source_tombstoned',
) -> Dict[str, Any]:
    return mutation(
        'tombstone_evidence',
        fact_id=fact_id,
        evidence_id=evidence_id,
        tombstoned_at=tombstoned_at,
        reason=reason,
    )


def merge_entities(
    entity_a: str,
    entity_b: str,
    evidence: Optional[Dict[str, Any]] = None,
    confidence: float = 0.5,
) -> Dict[str, Any]:
    return mutation(
        'merge_entities',
        entity_a=entity_a,
        entity_b=entity_b,
        evidence=copy.deepcopy(evidence or {}),
        confidence=confidence,
    )


def split_entity(entity_id: str, into: List[Dict[str, Any]], reason: str = '') -> Dict[str, Any]:
    return mutation('split_entity', entity_id=entity_id, into=copy.deepcopy(into), reason=reason)


def reassign_fact_subject(fact_id: str, old: Optional[str], new: Optional[str]) -> Dict[str, Any]:
    return mutation('reassign_fact_subject', fact_id=fact_id, old=old, new=new)


def _json_default(value: Any) -> Any:
    if isinstance(value, datetime):
        return value.astimezone(timezone.utc).isoformat()
    return str(value)


def _canonical_json(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(',', ':'), default=_json_default)


def commit_id_for(parent_commit_id: Optional[str], mutations: List[Dict[str, Any]]) -> str:
    payload = {'parent_commit_id': parent_commit_id, 'mutations': mutations}
    return hashlib.sha256(_canonical_json(payload).encode('utf-8')).hexdigest()


def build_commit(
    parent_commit_id: Optional[str],
    mutations: List[Dict[str, Any]],
    *,
    run_id: Optional[str] = None,
    commit_time: Optional[datetime] = None,
) -> Dict[str, Any]:
    commit_time = commit_time or datetime.now(timezone.utc)
    commit_id = commit_id_for(parent_commit_id, mutations)
    return {
        'commit_id': commit_id,
        'parent_commit_id': parent_commit_id,
        'commit_time': commit_time,
        'run_id': run_id,
        'mutations': copy.deepcopy(mutations),
    }


def append_commit_to_history(
    state: Dict[str, Any],
    commits: Dict[str, Dict[str, Any]],
    parent_commit_id: Optional[str],
    mutations: List[Dict[str, Any]],
    *,
    run_id: Optional[str] = None,
    commit_time: Optional[datetime] = None,
    use_current_head: bool = False,
) -> Dict[str, Any]:
    current_head = state.get('current_head_commit_id')
    expected_parent = current_head if use_current_head else parent_commit_id
    commit = build_commit(expected_parent, mutations, run_id=run_id, commit_time=commit_time)

    if commit['commit_id'] in commits:
        return {'commit': commits[commit['commit_id']], 'applied': False}

    if current_head != expected_parent:
        raise HeadConflict(expected_parent, current_head)

    commits[commit['commit_id']] = commit
    state['current_head_commit_id'] = commit['commit_id']
    state['projection_version'] = 1
    state['updated_at'] = commit['commit_time']
    return {'commit': commit, 'applied': True}


def append_commit(
    uid: str,
    parent_commit_id: Optional[str],
    mutations: List[Dict[str, Any]],
    *,
    run_id: Optional[str] = None,
    commit_time: Optional[datetime] = None,
    projection_writer: Optional[Callable[[Any], None]] = None,
    use_current_head: bool = False,
    firestore_client: Any = None,
) -> Dict[str, Any]:
    database: Any = firestore_client or db
    result = run_with_transaction_contention_retry(
        database.transaction,
        lambda transaction: _append_commit_transaction(
            transaction,
            database,
            uid,
            parent_commit_id,
            mutations,
            run_id,
            commit_time,
            projection_writer,
            use_current_head,
        ),
        operation_name="memory_ledger_append",
    )
    if result.get('applied'):
        projection_repair.enqueue_projection_repairs(uid, result.get('commit'), firestore_client=database)
    return result


def append_commit_with_builder(
    uid: str,
    parent_commit_id: Optional[str],
    mutation_builder: Callable[[Any], Dict[str, Any]],
    *,
    run_id: Optional[str] = None,
    commit_time: Optional[datetime] = None,
    use_current_head: bool = False,
    firestore_client: Any = None,
) -> Dict[str, Any]:
    database: Any = firestore_client or db
    result = run_with_transaction_contention_retry(
        database.transaction,
        lambda transaction: _append_commit_with_builder_transaction(
            transaction,
            database,
            uid,
            parent_commit_id,
            mutation_builder,
            run_id,
            commit_time,
            use_current_head,
        ),
        operation_name="memory_ledger_append_with_builder",
    )
    if result.get('applied'):
        projection_repair.enqueue_projection_repairs(uid, result.get('commit'), firestore_client=database)
    return result


@_typed_transactional
def _append_commit_transaction(
    transaction: Any,
    database: Any,
    uid: str,
    parent_commit_id: Optional[str],
    mutations: List[Dict[str, Any]],
    run_id: Optional[str],
    commit_time: Optional[datetime],
    projection_writer: Optional[Callable[[Any], None]],
    use_current_head: bool,
) -> Dict[str, Any]:
    user_ref = database.collection(users_collection).document(uid)
    state_ref = user_ref.collection(memory_state_collection).document(memory_state_document)
    state_snapshot = state_ref.get(transaction=transaction)
    state: Dict[str, Any] = _typed_doc(state_snapshot) if state_snapshot.exists else {}
    current_head = state.get('current_head_commit_id')
    expected_parent = current_head if use_current_head else parent_commit_id

    commit = build_commit(expected_parent, mutations, run_id=run_id, commit_time=commit_time)
    commit_ref = user_ref.collection(memory_commits_collection).document(commit['commit_id'])
    commit_snapshot = commit_ref.get(transaction=transaction)

    if commit_snapshot.exists:
        return {'commit': _typed_doc(commit_snapshot), 'applied': False}

    if current_head != expected_parent:
        raise HeadConflict(expected_parent, current_head)

    _assert_legacy_commit_privacy_fences(
        transaction=transaction,
        database=database,
        uid=uid,
        mutations=mutations,
    )

    # Build the state-head payload (including any apply_control fallback read)
    # before staging any writes: Firestore forbids a transactional read after a
    # write and raises ReadAfterWriteError otherwise.
    state_head_payload = _state_head_write_payload(
        transaction=transaction,
        user_ref=user_ref,
        uid=uid,
        state=state,
        commit=commit,
    )

    if projection_writer:
        projection_writer(transaction)

    transaction.set(commit_ref, commit)
    transaction.set(state_ref, state_head_payload)
    return {'commit': commit, 'applied': True}


@_typed_transactional
def _append_commit_with_builder_transaction(
    transaction: Any,
    database: Any,
    uid: str,
    parent_commit_id: Optional[str],
    mutation_builder: Callable[[Any], Dict[str, Any]],
    run_id: Optional[str],
    commit_time: Optional[datetime],
    use_current_head: bool,
) -> Dict[str, Any]:
    user_ref = database.collection(users_collection).document(uid)
    state_ref = user_ref.collection(memory_state_collection).document(memory_state_document)
    state_snapshot = state_ref.get(transaction=transaction)
    state: Dict[str, Any] = _typed_doc(state_snapshot) if state_snapshot.exists else {}
    current_head = state.get('current_head_commit_id')
    expected_parent = current_head if use_current_head else parent_commit_id

    if current_head != expected_parent:
        raise HeadConflict(expected_parent, current_head)

    built = mutation_builder(transaction)
    mutations: List[Dict[str, Any]] = cast(List[Dict[str, Any]], built.get('mutations') or [])
    projection_writer = built.get('projection_writer')
    _assert_legacy_commit_privacy_fences(
        transaction=transaction,
        database=database,
        uid=uid,
        mutations=mutations,
    )
    commit = build_commit(expected_parent, mutations, run_id=run_id, commit_time=commit_time)
    commit_ref = user_ref.collection(memory_commits_collection).document(commit['commit_id'])
    commit_snapshot = commit_ref.get(transaction=transaction)

    if commit_snapshot.exists:
        return {'commit': _typed_doc(commit_snapshot), 'applied': False}

    # Build the state-head payload (including any apply_control fallback read)
    # before staging any writes: Firestore forbids a transactional read after a
    # write and raises ReadAfterWriteError otherwise.
    state_head_payload = _state_head_write_payload(
        transaction=transaction,
        user_ref=user_ref,
        uid=uid,
        state=state,
        commit=commit,
    )

    if projection_writer:
        projection_writer(transaction)

    transaction.set(commit_ref, commit)
    transaction.set(state_ref, state_head_payload)
    return {'commit': commit, 'applied': True}


def read_head(uid: str) -> Optional[str]:
    state_ref = (
        db.collection(users_collection)
        .document(uid)
        .collection(memory_state_collection)
        .document(memory_state_document)
    )
    state_snapshot = state_ref.get()
    if not state_snapshot.exists:
        return None
    state: Dict[str, Any] = _typed_doc(state_snapshot)
    return state.get('current_head_commit_id')


def fold_commits(commits: List[Dict[str, Any]], valid_time: Optional[datetime] = None) -> Dict[str, Dict[str, Any]]:
    facts: Dict[str, Dict[str, Any]] = {}
    for commit in sorted(
        commits, key=lambda item: item.get('commit_time') or datetime.min.replace(tzinfo=timezone.utc)
    ):
        mutations: List[Dict[str, Any]] = cast(List[Dict[str, Any]], commit.get('mutations') or [])
        for item in mutations:
            _apply_mutation(facts, item, commit.get('commit_time'))
    if valid_time is None:
        return {fact_id: fact for fact_id, fact in facts.items() if fact.get('invalid_at') is None}
    return {fact_id: fact for fact_id, fact in facts.items() if _fact_valid_at(fact, valid_time)}


def _resolve_arg_change(value: Any) -> Any:
    if isinstance(value, dict) and "to" in value:
        return cast(Dict[str, Any], value).get("to")
    return cast(Any, value)


def _apply_mutation(facts: Dict[str, Dict[str, Any]], item: Dict[str, Any], commit_time: Optional[datetime]) -> None:
    mutation_type = item.get('type')
    if mutation_type == 'add_fact':
        fact: Dict[str, Any] = cast(Dict[str, Any], copy.deepcopy(item.get('fact') or {}))
        fact_id = fact.get('id')
        if fact_id:
            facts[fact_id] = fact
        return

    fact_id = item.get('fact_id')
    if not fact_id or fact_id not in facts:
        return

    if mutation_type == 'supersede_fact':
        facts[fact_id]['superseded_by'] = item.get('by')
        valid_interval: Dict[str, Any] = cast(Dict[str, Any], item.get('valid_interval') or {})
        invalid_at: Any = valid_interval.get('valid_to') or commit_time
        facts[fact_id]['invalid_at'] = invalid_at
        facts[fact_id].setdefault('qualifiers', {})['valid_to'] = invalid_at
        return

    if mutation_type == 'retract_fact':
        facts[fact_id]['invalid_at'] = commit_time
        facts[fact_id]['retraction_reason'] = item.get('reason')
        facts[fact_id]['redaction_status'] = 'payload_tombstoned'
        facts[fact_id]['content'] = None
        facts[fact_id]['arguments'] = {}
        return

    if mutation_type == 'refine_fact':
        _apply_arg_changes(facts[fact_id], cast(Dict[str, Any], item.get('arg_changes') or {}))
        return

    if mutation_type == 'add_evidence':
        evidence: Any = facts[fact_id].setdefault('evidence', [])
        new_evidence = item.get('evidence')
        if new_evidence and new_evidence not in evidence:
            evidence.append(copy.deepcopy(new_evidence))
        return

    if mutation_type == 'remove_evidence':
        evidence_id = item.get('evidence_id')
        evidence_items: Any = facts[fact_id].get('evidence') or []
        facts[fact_id]['evidence'] = [
            evidence
            for evidence in evidence_items
            if not isinstance(evidence, dict) or cast(Dict[str, Any], evidence).get('evidence_id') != evidence_id
        ]
        return

    if mutation_type == 'tombstone_evidence':
        evidence_id = item.get('evidence_id')
        current_evidence: Any = facts[fact_id].get('evidence') or []
        for evidence in current_evidence:
            if isinstance(evidence, dict) and cast(Dict[str, Any], evidence).get('evidence_id') == evidence_id:
                evidence_dict: Dict[str, Any] = cast(Dict[str, Any], evidence)
                evidence_dict['redaction_status'] = 'tombstoned'
                evidence_dict['tombstoned_at'] = item.get('tombstoned_at') or commit_time
                evidence_dict['tombstone_reason'] = item.get('reason')
        active_evidence: List[Dict[str, Any]] = [
            cast(Dict[str, Any], evidence)
            for evidence in current_evidence
            if isinstance(evidence, dict)
            and cast(Dict[str, Any], evidence).get('redaction_status', 'active') != 'tombstoned'
        ]
        facts[fact_id].update(
            confidence_fields_for_evidence(
                active_evidence,
                facts[fact_id].get('subject_attribution', 'unknown'),
                existing_capture_confidence=facts[fact_id].get('capture_confidence'),
            )
        )
        return

    if mutation_type == 'reassign_fact_subject':
        facts[fact_id]['subject_entity_id'] = item.get('new')


def _apply_arg_changes(fact: Dict[str, Any], arg_changes: Dict[str, Any]) -> None:
    arguments: Any = fact.setdefault('arguments', {})
    for key, value in arg_changes.items():
        resolved: Any = _resolve_arg_change(value)
        if key == 'content':
            fact['content'] = resolved
            continue
        arguments[key] = resolved


def _fact_valid_at(fact: Dict[str, Any], valid_time: datetime) -> bool:
    qualifiers: Dict[str, Any] = cast(Dict[str, Any], fact.get('qualifiers') or {})
    valid_from: Any = qualifiers.get('valid_from') or fact.get('valid_at')
    valid_to: Any = qualifiers.get('valid_to') or fact.get('invalid_at')
    if isinstance(valid_from, datetime) and valid_time < valid_from:
        return False
    if isinstance(valid_to, datetime) and valid_time > valid_to:
        return False
    return True


def replay_to(
    uid: str,
    commit_time: Optional[datetime] = None,
    valid_time: Optional[datetime] = None,
) -> Dict[str, Dict[str, Any]]:
    commits_ref = db.collection(users_collection).document(uid).collection(memory_commits_collection)
    commits: List[Dict[str, Any]] = []
    for doc in commits_ref.order_by('commit_time').stream():
        commit: Dict[str, Any] = _typed_doc(doc)
        commit_time_value: Any = commit.get('commit_time')
        if commit_time is None or commit_time_value <= commit_time:
            commits.append(commit)
    return fold_commits(commits, valid_time=valid_time)


def diff(commit_a: Dict[str, Any], commit_b: Dict[str, Any]) -> List[Dict[str, Any]]:
    if commit_b.get('parent_commit_id') == commit_a.get('commit_id'):
        return copy.deepcopy(commit_b.get('mutations') or [])
    return []


def diff_commits(
    commits: List[Dict[str, Any]], from_commit_id: Optional[str], to_commit_id: str
) -> List[Dict[str, Any]]:
    by_parent: Dict[Optional[str], Dict[str, Any]] = {commit.get('parent_commit_id'): commit for commit in commits}
    out: List[Dict[str, Any]] = []
    cursor = from_commit_id
    while cursor != to_commit_id and cursor in by_parent:
        next_commit = by_parent[cursor]
        out.extend(copy.deepcopy(next_commit.get('mutations') or []))
        cursor = next_commit.get('commit_id')
    return out if cursor == to_commit_id else []