ForcePilot/backend/package/yuxi/channels/adapters/imessage/echo_cache.py

from __future__ import annotations

import hashlib
import time
from dataclasses import dataclass

TEXT_TTL_S = 4.0
TEXT_BACKED_BY_ID_TTL_S = 4.0
MESSAGE_ID_TTL_S = 60.0
SELF_CHAT_TTL_S = 10.0
MAX_SELF_CHAT_ENTRIES = 512


@dataclass
class _SentEntry:
    message_id: str
    text: str
    chat_guid: str
    sent_at: float


@dataclass
class _SelfChatEntry:
    text_hash: str
    chat_guid: str
    recorded_at: float


class EchoCache:
    def __init__(
        self,
        text_ttl_s: float = TEXT_TTL_S,
        text_backed_by_id_ttl_s: float = TEXT_BACKED_BY_ID_TTL_S,
        message_id_ttl_s: float = MESSAGE_ID_TTL_S,
        self_chat_ttl_s: float = SELF_CHAT_TTL_S,
        max_self_chat_entries: int = MAX_SELF_CHAT_ENTRIES,
    ):
        self._text_ttl = text_ttl_s
        self._text_backed_by_id_ttl = text_backed_by_id_ttl_s
        self._msg_id_ttl = message_id_ttl_s
        self._self_chat_ttl = self_chat_ttl_s
        self._max_self_chat_entries = max_self_chat_entries

        self._by_message_id: dict[str, _SentEntry] = {}
        self._by_text: dict[str, list[_SentEntry]] = {}
        self._by_text_backed_by_id: dict[str, _SentEntry] = {}

        self._self_chat_cache: list[_SelfChatEntry] = []
        self._last_self_chat_cleanup = time.monotonic()

    def record_sent(self, message_id: str, text: str, chat_guid: str) -> None:
        now = time.monotonic()
        entry = _SentEntry(message_id=message_id, text=text, chat_guid=chat_guid, sent_at=now)
        self._by_message_id[message_id] = entry

        text_key = _text_key(text, chat_guid)
        if text_key not in self._by_text:
            self._by_text[text_key] = []
        self._by_text[text_key].append(entry)

        if message_id:
            self._by_text_backed_by_id[message_id] = entry

        self._cleanup()

    def is_echo(self, message_id: str | None = None, text: str | None = None, chat_guid: str = "") -> bool:
        self._cleanup()

        if message_id and message_id in self._by_message_id:
            return True

        if message_id and message_id in self._by_text_backed_by_id:
            return True

        if text:
            text_key = _text_key(text, chat_guid)
            entries = self._by_text.get(text_key, [])
            if entries:
                return True

        return False

    def is_self_chat_echo(
        self,
        sender_handle: str,
        self_handle: str,
        chat_guid: str,
    ) -> bool:
        if not self_handle:
            return False
        sender_clean = sender_handle.strip().replace(" ", "").lstrip("+")
        self_clean = self_handle.strip().replace(" ", "").lstrip("+")
        return sender_clean == self_clean

    def record_self_chat_text(self, text: str, chat_guid: str) -> None:
        now = time.monotonic()
        text_hash = hashlib.sha256(f"{chat_guid}:{text.strip()[:200]}".encode()).hexdigest()
        entry = _SelfChatEntry(text_hash=text_hash, chat_guid=chat_guid, recorded_at=now)
        self._self_chat_cache.append(entry)
        self._cleanup_self_chat()
        if len(self._self_chat_cache) > self._max_self_chat_entries:
            self._self_chat_cache = self._self_chat_cache[-self._max_self_chat_entries :]

    def is_in_self_chat_cache(self, text: str, chat_guid: str) -> bool:
        self._cleanup_self_chat()
        text_hash = hashlib.sha256(f"{chat_guid}:{text.strip()[:200]}".encode()).hexdigest()
        for entry in self._self_chat_cache:
            if entry.text_hash == text_hash and entry.chat_guid == chat_guid:
                return True
        return False

    def _cleanup_self_chat(self) -> None:
        now = time.monotonic()
        self._self_chat_cache = [
            entry for entry in self._self_chat_cache if now - entry.recorded_at <= self._self_chat_ttl
        ]

    def _cleanup(self) -> None:
        now = time.monotonic()

        expired_msg_ids = [mid for mid, entry in self._by_message_id.items() if now - entry.sent_at > self._msg_id_ttl]
        for mid in expired_msg_ids:
            del self._by_message_id[mid]

        expired_text_backed = [
            mid
            for mid, entry in self._by_text_backed_by_id.items()
            if now - entry.sent_at > self._text_backed_by_id_ttl
        ]
        for mid in expired_text_backed:
            del self._by_text_backed_by_id[mid]

        for text_key in list(self._by_text):
            self._by_text[text_key] = [
                entry for entry in self._by_text[text_key] if now - entry.sent_at <= self._text_ttl
            ]
            if not self._by_text[text_key]:
                del self._by_text[text_key]

        if now - self._last_self_chat_cleanup > self._self_chat_ttl:
            self._cleanup_self_chat()
            self._last_self_chat_cleanup = now


def _text_key(text: str, chat_guid: str) -> str:
    return f"{chat_guid}:{text.strip()[:200]}"