nanobot/nanobot/channels/matrix/runtime.py

"""Matrix (Element) channel — inbound sync + outbound message/media delivery."""

# pyright: reportMissingTypeStubs=false

import asyncio
import html
import json
import mimetypes
import re
import sys
import time
from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Literal, Protocol, TypeAlias, cast
from urllib.parse import quote, unquote, urlparse

from pydantic import Field

from nanobot.security.workspace_policy import is_path_within

try:
    import aiohttp
    import nh3
    from mistune import HTMLRenderer, create_markdown
    from nio import (
        Api,
        AsyncClient,
        AsyncClientConfig,
        InviteEvent,
        JoinError,
        JoinResponse,
        KeyVerificationCancel,
        KeyVerificationEvent,
        KeyVerificationKey,
        KeyVerificationMac,
        KeyVerificationStart,
        LoginResponse,
        MatrixRoom,
        RoomEncryptedMedia,
        RoomMessage,
        RoomMessageMedia,
        RoomMessageText,
        RoomSendError,
        RoomSendResponse,
        RoomTypingError,
        SyncError,
        SyncResponse,
        ToDeviceError,
        UploadError,
    )
    from nio.crypto.attachments import decrypt_attachment
    from nio.exceptions import EncryptionError
except ImportError as e:
    raise ImportError(
        "Matrix dependencies not installed. Run: nanobot plugins enable matrix"
    ) from e

from nanobot.bus.events import OutboundMessage
from nanobot.bus.outbound_events import ProgressEvent
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.paths import get_data_dir, get_media_dir
from nanobot.config.schema import Base
from nanobot.utils.helpers import safe_filename
from nanobot.utils.logging_bridge import redirect_lib_logging

TYPING_NOTICE_TIMEOUT_MS = 30_000
# Must stay below TYPING_NOTICE_TIMEOUT_MS so the indicator doesn't expire mid-processing.
TYPING_KEEPALIVE_INTERVAL_MS = 20_000
MATRIX_HTML_FORMAT = "org.matrix.custom.html"
_ATTACH_MARKER = "[attachment: {}]"
_ATTACH_TOO_LARGE = "[attachment: {} - too large]"
_ATTACH_FAILED = "[attachment: {} - download failed]"
_ATTACH_UPLOAD_FAILED = "[attachment: {} - upload failed]"
_DEFAULT_ATTACH_NAME = "attachment"
_MSGTYPE_MAP = {"m.image": "image", "m.audio": "audio", "m.video": "video", "m.file": "file"}

MATRIX_MEDIA_EVENT_FILTER = (RoomMessageMedia, RoomEncryptedMedia)
MatrixMediaEvent: TypeAlias = RoomMessageMedia | RoomEncryptedMedia


class _MatrixCallbackRegistrar(Protocol):
    """Runtime callback surface whose upstream stubs reject valid filtered handlers."""

    def add_event_callback(self, callback: Callable[..., Any], event_filter: Any) -> None: ...
    def add_to_device_callback(
        self,
        callback: Callable[..., Any],
        event_filter: Any,
    ) -> None: ...
    def add_response_callback(self, callback: Callable[..., Any], response_filter: Any) -> None: ...


class _MediaTooLargeError(Exception):
    """Raised when an inbound Matrix media download exceeds the configured cap."""

MATRIX_MARKDOWN = create_markdown(
    renderer=HTMLRenderer(escape=True, allow_harmful_protocols=("mxc://",)),
    plugins=["table", "strikethrough", "url", "superscript", "subscript"],
)

MATRIX_ALLOWED_HTML_TAGS = {
    "p", "a", "strong", "em", "del", "code", "pre", "blockquote",
    "ul", "ol", "li", "h1", "h2", "h3", "h4", "h5", "h6",
    "hr", "br", "table", "thead", "tbody", "tr", "th", "td",
    "caption", "sup", "sub", "img",
}
MATRIX_ALLOWED_HTML_ATTRIBUTES: dict[str, set[str]] = {
    "a": {"href"}, "code": {"class"}, "ol": {"start"},
    "img": {"src", "alt", "title", "width", "height"},
}
MATRIX_ALLOWED_URL_SCHEMES = {"https", "http", "matrix", "mailto", "mxc"}
_MXC_IMAGE_PLACEHOLDER_PREFIX = "https://nanobot.invalid/matrix-mxc/"
_MXC_MARKDOWN_IMAGE_RE = re.compile(
    r"(?P<prefix>!\[[^\]]*\]\()"
    r"(?P<value>mxc://[^\s)]+)"
    r"(?P<suffix>(?:\s+[^)]*)?\))"
)
_MXC_IMAGE_SRC_RE = re.compile(
    r"(?P<prefix>\bsrc=)(?P<quote>[\"'])(?P<value>mxc://[^\"']+)(?P=quote)",
    re.IGNORECASE,
)
_MXC_PLACEHOLDER_SRC_RE = re.compile(
    rf'src="{re.escape(_MXC_IMAGE_PLACEHOLDER_PREFIX)}([^"]+)"'
)


def _filter_matrix_html_attribute(tag: str, attr: str, value: str) -> str | None:
    """Filter attribute values to a safe Matrix-compatible subset."""
    if tag == "a" and attr == "href":
        return value if value.lower().startswith(("https://", "http://", "matrix:", "mailto:")) else None
    if tag == "img" and attr == "src":
        lowered = value.lower()
        if lowered.startswith("mxc://") or lowered.startswith(_MXC_IMAGE_PLACEHOLDER_PREFIX):
            return value
        return None
    if tag == "code" and attr == "class":
        classes = [c for c in value.split() if c.startswith("language-") and not c.startswith("language-_")]
        return " ".join(classes) if classes else None
    return value


MATRIX_HTML_CLEANER = nh3.Cleaner(
    tags=MATRIX_ALLOWED_HTML_TAGS,
    attributes=MATRIX_ALLOWED_HTML_ATTRIBUTES,
    attribute_filter=_filter_matrix_html_attribute,
    url_schemes=MATRIX_ALLOWED_URL_SCHEMES,
    strip_comments=True,
    link_rel="noopener noreferrer",
)


def _mask_mxc_markdown_image_sources(text: str) -> str:
    def repl(match: re.Match[str]) -> str:
        value = quote(match.group("value"), safe="")
        return (
            f"{match.group('prefix')}"
            f"{_MXC_IMAGE_PLACEHOLDER_PREFIX}{value}"
            f"{match.group('suffix')}"
        )

    return _MXC_MARKDOWN_IMAGE_RE.sub(repl, text)


def _mask_mxc_image_sources(rendered_html: str) -> str:
    def repl(match: re.Match[str]) -> str:
        value = quote(match.group("value"), safe="")
        return (
            f'{match.group("prefix")}{match.group("quote")}'
            f"{_MXC_IMAGE_PLACEHOLDER_PREFIX}{value}"
            f'{match.group("quote")}'
        )

    return _MXC_IMAGE_SRC_RE.sub(repl, rendered_html)


def _unmask_mxc_image_sources(cleaned_html: str) -> str:
    def repl(match: re.Match[str]) -> str:
        value = html.escape(unquote(match.group(1)), quote=True)
        return f'src="{value}"'

    return _MXC_PLACEHOLDER_SRC_RE.sub(repl, cleaned_html)


@dataclass
class _StreamBuf:
    """
    Represents a buffer for managing LLM response stream data.

    :ivar text: Stores the text content of the buffer.
    :type text: str
    :ivar event_id: Identifier for the associated event. None indicates no
        specific event association.
    :type event_id: str | None
    :ivar last_edit: Timestamp of the most recent edit to the buffer.
    :type last_edit: float
    """
    text: str = ""
    event_id: str | None = None
    last_edit: float = 0.0

def _render_markdown_html(text: str) -> str | None:
    """Render markdown to sanitized HTML; returns None for plain text."""
    try:
        masked_text = _mask_mxc_markdown_image_sources(text)
        rendered = _mask_mxc_image_sources(cast(str, MATRIX_MARKDOWN(masked_text)))
        formatted = _unmask_mxc_image_sources(MATRIX_HTML_CLEANER.clean(rendered).strip())
    except Exception:
        return None
    if not formatted:
        return None
    # Skip formatted_body for plain <p>text</p> to keep payload minimal.
    if formatted.startswith("<p>") and formatted.endswith("</p>"):
        inner = formatted[3:-4]
        if "<" not in inner and ">" not in inner:
            return None
    return formatted


def _build_matrix_text_content(
    text: str,
    event_id: str | None = None,
    thread_relates_to: dict[str, object] | None = None,
) -> dict[str, object]:
    """
    Constructs and returns a dictionary representing the matrix text content with optional
    HTML formatting and reference to an existing event for replacement. This function is
    primarily used to create content payloads compatible with the Matrix messaging protocol.

    :param text: The plain text content to include in the message.
    :type text: str
    :param event_id: Optional ID of the event to replace. If provided, the function will
        include information indicating that the message is a replacement of the specified
        event.
    :type event_id: str | None
    :param thread_relates_to: Optional Matrix thread relation metadata. For edits this is
        stored in ``m.new_content`` so the replacement remains in the same thread.
    :type thread_relates_to: dict[str, object] | None
    :return: A dictionary containing the matrix text content, potentially enriched with
        HTML formatting and replacement metadata if applicable.
    :rtype: dict[str, object]
    """
    content: dict[str, object] = {"msgtype": "m.text", "body": text, "m.mentions": {}}
    if html := _render_markdown_html(text):
        content["format"] = MATRIX_HTML_FORMAT
        content["formatted_body"] = html
    if event_id:
        new_content: dict[str, object] = {
            "body": text,
            "msgtype": "m.text",
        }
        content["m.new_content"] = new_content
        content["m.relates_to"] = {
            "rel_type": "m.replace",
            "event_id": event_id,
        }
        if thread_relates_to:
            new_content["m.relates_to"] = thread_relates_to
    elif thread_relates_to:
        content["m.relates_to"] = thread_relates_to

    return content


def _matrix_stream_key(chat_id: str, stream_id: str | None) -> str:
    return chat_id if stream_id is None else f"{chat_id}\0{stream_id}"


class MatrixConfig(Base):
    """Matrix (Element) channel configuration."""

    enabled: bool = False
    homeserver: str = "https://matrix.org"
    user_id: str = ""
    password: str = ""
    access_token: str = ""
    device_id: str = ""
    e2ee_enabled: bool = Field(default=sys.platform != "win32", alias="e2eeEnabled")
    sas_verification: bool = Field(default=False, alias="sasVerification")
    sync_stop_grace_seconds: int = 2
    max_media_bytes: int = 20 * 1024 * 1024
    max_concurrent_media_downloads: int = 2
    allow_from: list[str] = Field(default_factory=list)
    group_policy: Literal["open", "mention", "allowlist"] = "open"
    group_allow_from: list[str] = Field(default_factory=list)
    allow_room_mentions: bool = False
    streaming: bool = False


class MatrixChannel(BaseChannel):
    """Matrix (Element) channel using long-polling sync."""

    name = "matrix"
    display_name = "Matrix"
    _STREAM_EDIT_INTERVAL = 2 # min seconds between edit_message_text calls
    monotonic_time: Callable[[], float] = staticmethod(time.monotonic)

    @classmethod
    def default_config(cls) -> dict[str, Any]:
        return MatrixConfig().model_dump(by_alias=True)

    def __init__(
        self,
        config: Any,
        bus: MessageBus,
        *,
        restrict_to_workspace: bool = False,
        workspace: str | Path | None = None,
    ):
        if isinstance(config, dict):
            config = MatrixConfig.model_validate(config)
        super().__init__(config, bus)
        self.client: AsyncClient | None = None
        self._sync_task: asyncio.Task[None] | None = None
        self._typing_tasks: dict[str, asyncio.Task[None]] = {}
        self._restrict_to_workspace = bool(restrict_to_workspace)
        self._workspace = (
            Path(workspace).expanduser().resolve(strict=False) if workspace is not None else None
        )
        self._server_upload_limit_bytes: int | None = None
        self._server_upload_limit_checked = False
        self._stream_bufs: dict[str, _StreamBuf] = {}
        self._started_at_ms: int = 0
        self._media_download_semaphore = asyncio.Semaphore(
            max(1, int(self.config.max_concurrent_media_downloads))
        )


    async def start(self) -> None:
        """Start Matrix client and begin sync loop."""
        self._running = True
        self._started_at_ms = int(time.time() * 1000)
        redirect_lib_logging("nio", level="WARNING")

        self.store_path = get_data_dir() / "matrix-store"
        self.store_path.mkdir(parents=True, exist_ok=True)
        self.session_path = self.store_path / "session.json"

        # Replace ':' with '_' to produce a Windows-safe filename
        safe_store_name = self.config.user_id.replace(":", "_") + f"_{self.config.device_id}.db"

        self.client = AsyncClient(
            homeserver=self.config.homeserver,
            user=self.config.user_id,
            store_path=str(self.store_path),
            config=AsyncClientConfig(
                store_sync_tokens=True,
                encryption_enabled=self.config.e2ee_enabled,
                store_name=safe_store_name,
            ),
        )

        self._register_event_callbacks()
        self._register_to_device_callbacks()
        self._register_response_callbacks()

        if not self.config.e2ee_enabled:
            self.logger.warning("E2EE disabled; encrypted rooms may be undecryptable.")

        if self.config.password:
            if self.config.access_token or self.config.device_id:
                self.logger.warning("Password-based login active; access_token and device_id fields will be ignored.")

            create_new_session = True
            if self.session_path.exists():
                self.logger.info("Found session.json at {}; attempting to use existing session...", self.session_path)
                try:
                    with open(self.session_path, "r", encoding="utf-8") as f:
                        session = json.load(f)
                    self.client.user_id = self.config.user_id
                    self.client.access_token = session["access_token"]
                    self.client.device_id = session["device_id"]
                    self.client.load_store()
                    self.logger.info("Successfully loaded from existing session")
                    create_new_session = False
                except Exception as e:
                    self.logger.warning("Failed to load from existing session: {}", e)
                    self.logger.info("Falling back to password login...")

            if create_new_session:
                self.logger.info("Using password login...")
                resp = await self.client.login(self.config.password)
                if isinstance(resp, LoginResponse):
                    self.logger.info("Logged in using a password; saving details to disk")
                    self._write_session_to_disk(resp)
                else:
                    self.logger.error("Failed to log in: {}", resp)
                    return

        elif self.config.access_token and self.config.device_id:
            try:
                self.client.user_id = self.config.user_id
                self.client.access_token = self.config.access_token
                self.client.device_id = self.config.device_id
                self.client.load_store()
                self.logger.info("Successfully loaded from existing session")
            except Exception as e:
                self.logger.warning("Failed to load from existing session: {}", e)

        else:
            self.logger.warning("Unable to load a session due to missing password, access_token, or device_id; encryption may not work")
            return

        self._sync_task = asyncio.create_task(self._sync_loop())

    def _require_client(self) -> AsyncClient:
        if self.client is None:
            raise RuntimeError("Matrix client is not started")
        return self.client

    def _callback_registrar(self) -> _MatrixCallbackRegistrar:
        # matrix-nio's callback annotations do not model filtered subtype or
        # async handlers, although the runtime API supports both.
        return cast(_MatrixCallbackRegistrar, self._require_client())

    async def stop(self) -> None:
        """Stop the Matrix channel with graceful sync shutdown."""
        self._running = False
        for room_id in list(self._typing_tasks):
            await self._stop_typing_keepalive(room_id, clear_typing=False)
        if self.client:
            self.client.stop_sync_forever()
        if self._sync_task:
            try:
                await asyncio.wait_for(asyncio.shield(self._sync_task),
                                       timeout=self.config.sync_stop_grace_seconds)
            except (asyncio.TimeoutError, asyncio.CancelledError):
                self._sync_task.cancel()
                with suppress(asyncio.CancelledError):
                    await self._sync_task
        if self.client:
            await self.client.close()

    def _write_session_to_disk(self, resp: LoginResponse) -> None:
        """Save login session to disk for persistence across restarts."""
        session = {
            "access_token": resp.access_token,
            "device_id": resp.device_id,
        }
        try:
            with open(self.session_path, "w", encoding="utf-8") as f:
                json.dump(session, f, indent=2)
            self.logger.info("Session saved to {}", self.session_path)
        except Exception as e:
            self.logger.warning("Failed to save session: {}", e)

    def _is_workspace_path_allowed(self, path: Path) -> bool:
        """Check path is inside workspace (when restriction enabled)."""
        if not self._restrict_to_workspace or not self._workspace:
            return True
        return is_path_within(path, self._workspace)

    def _collect_outbound_media_candidates(self, media: list[str]) -> list[Path]:
        """Deduplicate and resolve outbound attachment paths."""
        seen: set[str] = set()
        candidates: list[Path] = []
        for raw in media:
            raw_value = cast(object, raw)
            if not isinstance(raw_value, str) or not raw_value.strip():
                continue
            path = Path(raw_value.strip()).expanduser()
            try:
                key = str(path.resolve(strict=False))
            except OSError:
                key = str(path)
            if key not in seen:
                seen.add(key)
                candidates.append(path)
        return candidates

    @staticmethod
    def _build_outbound_attachment_content(
        *, filename: str, mime: str, size_bytes: int,
        mxc_url: str, encryption_info: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Build Matrix content payload for an uploaded file/image/audio/video."""
        prefix = mime.split("/")[0]
        msgtype = {"image": "m.image", "audio": "m.audio", "video": "m.video"}.get(prefix, "m.file")
        content: dict[str, Any] = {
            "msgtype": msgtype, "body": filename, "filename": filename,
            "info": {"mimetype": mime, "size": size_bytes}, "m.mentions": {},
        }
        if encryption_info:
            content["file"] = {**encryption_info, "url": mxc_url}
        else:
            content["url"] = mxc_url
        return content

    def _is_encrypted_room(self, room_id: str) -> bool:
        if not self.client:
            return False
        room = getattr(self.client, "rooms", {}).get(room_id)
        return bool(getattr(room, "encrypted", False))

    async def _send_room_content(self, room_id: str,
                                 content: dict[str, Any]) -> None | RoomSendResponse | RoomSendError:
        """Send m.room.message with E2EE options."""
        if not self.client:
            return None
        kwargs: dict[str, Any] = {"room_id": room_id, "message_type": "m.room.message", "content": content}

        if self.config.e2ee_enabled:
            kwargs["ignore_unverified_devices"] = True
        response = await self.client.room_send(**kwargs)
        return response

    async def _resolve_server_upload_limit_bytes(self) -> int | None:
        """Query homeserver upload limit once per channel lifecycle."""
        if self._server_upload_limit_checked:
            return self._server_upload_limit_bytes
        self._server_upload_limit_checked = True
        if not self.client:
            return None
        try:
            response = await self.client.content_repository_config()
        except Exception:
            self.logger.error("Failed to fetch server upload limit", exc_info=True)
            return None
        upload_size = getattr(response, "upload_size", None)
        if isinstance(upload_size, int) and upload_size > 0:
            self._server_upload_limit_bytes = upload_size
            return upload_size
        return None

    async def _effective_media_limit_bytes(self) -> int:
        """min(local config, server advertised) — 0 blocks all uploads."""
        local_limit = max(int(self.config.max_media_bytes), 0)
        server_limit = await self._resolve_server_upload_limit_bytes()
        if server_limit is None:
            return local_limit
        return min(local_limit, server_limit) if local_limit else 0

    async def _upload_and_send_attachment(
        self, room_id: str, path: Path, limit_bytes: int,
        relates_to: dict[str, Any] | None = None,
    ) -> str | None:
        """Upload one local file to Matrix and send it as a media message. Returns failure marker or None."""
        if not self.client:
            return _ATTACH_UPLOAD_FAILED.format(path.name or _DEFAULT_ATTACH_NAME)

        resolved = path.expanduser().resolve(strict=False)
        filename = safe_filename(resolved.name) or _DEFAULT_ATTACH_NAME
        fail = _ATTACH_UPLOAD_FAILED.format(filename)

        if not resolved.is_file() or not self._is_workspace_path_allowed(resolved):
            return fail
        try:
            size_bytes = resolved.stat().st_size
        except OSError:
            return fail
        if limit_bytes <= 0 or size_bytes > limit_bytes:
            return _ATTACH_TOO_LARGE.format(filename)

        mime = mimetypes.guess_type(filename, strict=False)[0] or "application/octet-stream"
        try:
            with resolved.open("rb") as f:
                upload_result = await self.client.upload(
                    f, content_type=mime, filename=filename,
                    encrypt=self.config.e2ee_enabled and self._is_encrypted_room(room_id),
                    filesize=size_bytes,
                )
        except Exception:
            self.logger.error("Matrix media upload failed for %s", filename, exc_info=True)
            return fail

        is_tuple_result = isinstance(cast(object, upload_result), tuple)
        upload_response = upload_result[0] if is_tuple_result else upload_result
        encryption_info = (
            upload_result[1]
            if is_tuple_result and isinstance(cast(object, upload_result[1]), dict)
            else None
        )
        if isinstance(upload_response, UploadError):
            return fail
        mxc_url = getattr(upload_response, "content_uri", None)
        if not isinstance(mxc_url, str) or not mxc_url.startswith("mxc://"):
            return fail

        content = self._build_outbound_attachment_content(
            filename=filename, mime=mime, size_bytes=size_bytes,
            mxc_url=mxc_url, encryption_info=encryption_info,
        )
        if relates_to:
            content["m.relates_to"] = relates_to
        try:
            await self._send_room_content(room_id, content)
        except Exception:
            self.logger.error("Matrix room content send failed for room_id=%s", room_id, exc_info=True)
            return fail
        return None

    async def send(self, msg: OutboundMessage) -> None:
        """Send outbound content; clear typing for non-progress messages."""
        if not self.client:
            raise RuntimeError("Matrix client not initialized")
        text = msg.content or ""
        candidates = self._collect_outbound_media_candidates(msg.media)
        relates_to = self._build_thread_relates_to(msg.metadata)
        is_progress = isinstance(msg.event, ProgressEvent)
        try:
            failures: list[str] = []
            if candidates:
                limit_bytes = await self._effective_media_limit_bytes()
                for path in candidates:
                    if fail := await self._upload_and_send_attachment(
                        room_id=msg.chat_id,
                        path=path,
                        limit_bytes=limit_bytes,
                        relates_to=relates_to,
                    ):
                        failures.append(fail)
            if failures:
                text = f"{text.rstrip()}\n{chr(10).join(failures)}" if text.strip() else "\n".join(failures)
            if text.strip():
                content = _build_matrix_text_content(text)
                if relates_to:
                    content["m.relates_to"] = relates_to
                response = await self._send_room_content(msg.chat_id, content)
                if isinstance(response, RoomSendError):
                    raise RuntimeError(f"Matrix message was not delivered: {response}")
        finally:
            if not is_progress:
                await self._stop_typing_keepalive(msg.chat_id, clear_typing=True)

    async def send_delta(
        self,
        chat_id: str,
        delta: str,
        metadata: dict[str, Any] | None = None,
        *,
        stream_id: str | None = None,
        stream_end: bool = False,
        resuming: bool = False,
        merge_next: bool = False,
    ) -> None:
        relates_to = self._build_thread_relates_to(metadata)

        if stream_end and merge_next:
            if not delta:
                return
            stream_end = False
        if stream_end:
            stream_key = _matrix_stream_key(chat_id, stream_id)
            buf = self._stream_bufs.pop(stream_key, None)
            if not buf or not buf.event_id or not buf.text:
                return

            await self._stop_typing_keepalive(chat_id, clear_typing=True)

            content = _build_matrix_text_content(
                buf.text,
                buf.event_id,
                thread_relates_to=relates_to,
            )
            await self._send_room_content(chat_id, content)
            return

        stream_key = _matrix_stream_key(chat_id, stream_id)
        buf = self._stream_bufs.get(stream_key)
        if buf is None:
            buf = _StreamBuf()
            self._stream_bufs[stream_key] = buf
        buf.text += delta

        if not buf.text.strip():
            return

        now = self.monotonic_time()

        if not buf.last_edit or (now - buf.last_edit) >= self._STREAM_EDIT_INTERVAL:
            try:
                content = _build_matrix_text_content(
                    buf.text,
                    buf.event_id,
                    thread_relates_to=relates_to,
                )
                response = await self._send_room_content(chat_id, content)
                buf.last_edit = now
                if not buf.event_id:
                    # we are editing the same message all the time, so only the first time the event id needs to be set
                    buf.event_id = cast(RoomSendResponse, response).event_id
            except Exception:
                self.logger.error("Stream send/edit failed for chat_id=%s", chat_id, exc_info=True)
                await self._stop_typing_keepalive(chat_id, clear_typing=True)


    def _register_event_callbacks(self) -> None:
        client = self._callback_registrar()
        client.add_event_callback(self._on_message, RoomMessageText)
        client.add_event_callback(self._on_media_message, MATRIX_MEDIA_EVENT_FILTER)
        client.add_event_callback(self._on_room_invite, InviteEvent)

    def _register_to_device_callbacks(self) -> None:
        if self.config.e2ee_enabled and self.config.sas_verification:
            client = self._callback_registrar()
            client.add_to_device_callback(
                self._on_key_verification_event,
                (KeyVerificationEvent,),
            )

    def _register_response_callbacks(self) -> None:
        client = self._callback_registrar()
        client.add_response_callback(self._on_sync_error, SyncError)
        client.add_response_callback(self._on_join_error, JoinError)
        client.add_response_callback(self._on_send_error, RoomSendError)
        client.add_response_callback(self._on_sync_invite_fallback, SyncResponse)

    def _is_sas_sender_allowed(self, sender: str) -> bool:
        return bool(sender and self.is_allowed(sender))

    async def _on_key_verification_event(self, event: KeyVerificationEvent) -> None:
        try:
            await self._handle_key_verification_event(event)
        except asyncio.CancelledError:
            raise
        except Exception:
            self.logger.exception("Matrix SAS verification handling failed")

    async def _handle_key_verification_event(self, event: KeyVerificationEvent) -> None:
        if not (self.config.e2ee_enabled and self.config.sas_verification):
            return
        if not self.client:
            return

        sender = str(getattr(event, "sender", "") or "")
        transaction_id = str(getattr(event, "transaction_id", "") or "")
        if not transaction_id or not self._is_sas_sender_allowed(sender):
            return

        if isinstance(event, KeyVerificationStart):
            if "emoji" not in (getattr(event, "short_authentication_string", None) or []):
                self.logger.info(
                    "Ignoring Matrix SAS verification from {} without emoji support",
                    sender,
                )
                return

            response = await self.client.accept_key_verification(transaction_id)
            if isinstance(response, ToDeviceError):
                self.logger.warning("Matrix SAS accept failed for {}: {}", sender, response)
            return

        if isinstance(event, KeyVerificationKey):
            responses = await self.client.send_to_device_messages()
            if any(isinstance(response, ToDeviceError) for response in responses):
                self.logger.warning("Matrix SAS key share failed for {}", sender)
                return

            response = await self.client.confirm_short_auth_string(transaction_id)
            if isinstance(response, ToDeviceError):
                self.logger.warning("Matrix SAS confirm failed for {}: {}", sender, response)
            return

        if isinstance(event, KeyVerificationMac):
            sas = getattr(self.client, "key_verifications", {}).get(transaction_id)
            if sas is not None and getattr(sas, "verified", False):
                self.logger.info("Matrix SAS verification completed for {}", sender)
            return

        if isinstance(event, KeyVerificationCancel):
            self.logger.info(
                "Matrix SAS verification cancelled by {}: {}",
                sender,
                getattr(event, "reason", ""),
            )

    def _is_fatal_auth_response(self, response: Any) -> bool:
        code = getattr(response, "status_code", None)
        is_auth = code in {"M_UNKNOWN_TOKEN", "M_FORBIDDEN", "M_UNAUTHORIZED"}
        return is_auth or bool(getattr(response, "soft_logout", False))

    def _log_response_error(self, label: str, response: Any) -> None:
        """Log Matrix response errors — auth errors at ERROR level, rest at WARNING."""
        is_fatal = self._is_fatal_auth_response(response)
        (self.logger.error if is_fatal else self.logger.warning)("{} failed: {}", label, response)

    async def _on_sync_error(self, response: SyncError) -> None:
        self._log_response_error("sync", response)
        if self._is_fatal_auth_response(response):
            # Auth errors won't recover by retry; stop the sync loop instead of
            # spamming the homeserver every 2s (#1851).
            self.logger.error("Authentication failed irrecoverably; stopping sync loop")
            self._running = False
            if self.client:
                with suppress(Exception):
                    self.client.stop_sync_forever()

    async def _join_room_safe(self, room_id: str) -> bool:
        """Join a room, sending a non-empty POST body.

        nio's ``Api.join()`` produces a POST with no body. Some homeservers
        (notably Continuwuity) reject empty bodies with ``M_BAD_JSON``.
        Sending ``"{}"`` satisfies both strict and lenient servers.
        """
        client = self._require_client()
        method, path = Api.join(client.access_token, room_id)
        try:
            resp = cast(
                JoinResponse | JoinError,
                await client._send(  # type: ignore[reportPrivateUsage, reportUnknownMemberType]
                    JoinResponse, method, path, data="{}"
                ),
            )
        except Exception:
            self.logger.error("Matrix join request exception for room={}", room_id, exc_info=True)
            return False
        if isinstance(resp, JoinError):
            self.logger.error("Matrix auto-join failed for room={}: {}", room_id, resp)
            return False
        self.logger.info("Matrix auto-join succeeded: {}", room_id)
        return True

    async def _on_sync_invite_fallback(self, response: SyncResponse) -> None:
        """Safety net: join pending invites that the event callback may have missed.

        Some homeservers (e.g. Continuwuity) deliver each invite only once.
        If ``_on_room_invite`` fires but the join fails, the sync token
        advances and the invite is never re-delivered. This callback inspects
        the same ``SyncResponse`` for pending invites and joins them, acting
        as a fallback alongside the event-based callback.
        """
        if not response.rooms or not response.rooms.invite:
            return
        for room_id, invite_info in response.rooms.invite.items():
            for event in cast(list[Any], invite_info.invite_state):
                sender = getattr(event, "sender", None)
                if sender and self.is_allowed(cast(str, sender)):
                    await self._join_room_safe(room_id)
                    break

    async def _on_join_error(self, response: JoinError) -> None:
        self._log_response_error("join", response)

    async def _on_send_error(self, response: RoomSendError) -> None:
        self._log_response_error("send", response)

    async def _set_typing(self, room_id: str, typing: bool) -> None:
        """Best-effort typing indicator update."""
        if not self.client:
            return
        with suppress(Exception):
            response = await self.client.room_typing(room_id=room_id, typing_state=typing,
                                                     timeout=TYPING_NOTICE_TIMEOUT_MS)
            if isinstance(response, RoomTypingError):
                self.logger.debug("typing failed for {}: {}", room_id, response)

    async def _start_typing_keepalive(self, room_id: str) -> None:
        """Start periodic typing refresh (spec-recommended keepalive)."""
        await self._stop_typing_keepalive(room_id, clear_typing=False)
        await self._set_typing(room_id, True)
        if not self._running:
            return

        async def loop() -> None:
            with suppress(asyncio.CancelledError):
                while self._running:
                    await asyncio.sleep(TYPING_KEEPALIVE_INTERVAL_MS / 1000)
                    await self._set_typing(room_id, True)

        self._typing_tasks[room_id] = asyncio.create_task(loop())

    async def _stop_typing_keepalive(self, room_id: str, *, clear_typing: bool) -> None:
        if task := self._typing_tasks.pop(room_id, None):
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task
        if clear_typing:
            await self._set_typing(room_id, False)

    async def _sync_loop(self) -> None:
        backoff = 2.0
        while self._running:
            try:
                client = self._require_client()
                await client.sync_forever(timeout=30000, full_state=True)
                backoff = 2.0
            except asyncio.CancelledError:
                break
            except Exception:
                if not self._running:
                    break
                await asyncio.sleep(backoff)
                backoff = min(backoff * 2, 60.0)

    async def _on_room_invite(self, room: MatrixRoom, event: InviteEvent) -> None:
        if self.is_allowed(event.sender):
            await self._join_room_safe(room.room_id)

    def _is_direct_room(self, room: MatrixRoom) -> bool:
        count = getattr(room, "member_count", None)
        return isinstance(count, int) and count <= 2

    def _is_bot_mentioned(self, event: RoomMessage) -> bool:
        """Check m.mentions payload for bot mention."""
        source = getattr(event, "source", None)
        if not isinstance(source, dict):
            return False
        source_data = cast(dict[str, Any], source)
        content = cast(dict[str, Any], source_data.get("content") or {})
        mentions = cast(object, content.get("m.mentions"))
        if not isinstance(mentions, dict):
            return False
        mentions_data = cast(dict[str, Any], mentions)
        user_ids = cast(object, mentions_data.get("user_ids"))
        if isinstance(user_ids, list) and self.config.user_id in user_ids:
            return True
        return bool(
            self.config.allow_room_mentions
            and mentions_data.get("room") is True
        )

    def _is_pre_startup_event(self, event: RoomMessage) -> bool:
        """Skip events that landed in the timeline before this process started.

        Matrix sync replays the room timeline on each startup/restart; without
        this filter old messages would be re-handled as if they were fresh
        (#3553).
        """
        ts = getattr(event, "server_timestamp", None)
        return isinstance(ts, int) and ts < self._started_at_ms

    def _should_process_message(self, room: MatrixRoom, event: RoomMessage) -> bool:
        """Apply sender and room policy checks."""
        if not self.is_allowed(event.sender):
            return False
        if self._is_direct_room(room):
            return True
        policy = self.config.group_policy
        if policy == "open":
            return True
        if policy == "allowlist":
            return room.room_id in (self.config.group_allow_from or [])
        if policy == "mention":
            return self._is_bot_mentioned(event)
        return False

    def _media_dir(self) -> Path:
        return get_media_dir("matrix")

    @staticmethod
    def _event_source_content(event: RoomMessage) -> dict[str, Any]:
        source = getattr(event, "source", None)
        if not isinstance(source, dict):
            return {}
        source_data = cast(dict[str, Any], source)
        content = cast(object, source_data.get("content"))
        return cast(dict[str, Any], content) if isinstance(content, dict) else {}

    def _event_thread_root_id(self, event: RoomMessage) -> str | None:
        relates_to = cast(
            object,
            self._event_source_content(event).get("m.relates_to"),
        )
        if not isinstance(relates_to, dict):
            return None
        relation = cast(dict[str, Any], relates_to)
        if relation.get("rel_type") != "m.thread":
            return None
        root_id = cast(object, relation.get("event_id"))
        return root_id if isinstance(root_id, str) and root_id else None

    def _thread_metadata(self, event: RoomMessage) -> dict[str, str] | None:
        if not (root_id := self._event_thread_root_id(event)):
            return None
        meta: dict[str, str] = {"thread_root_event_id": root_id}
        if isinstance(reply_to := getattr(event, "event_id", None), str) and reply_to:
            meta["thread_reply_to_event_id"] = reply_to
        return meta

    def _thread_session_key(self, room_id: str, event: RoomMessage) -> str | None:
        if not (root_id := self._event_thread_root_id(event)):
            return None
        return f"{self.name}:{room_id}:thread:{root_id}"

    @staticmethod
    def _build_thread_relates_to(metadata: dict[str, Any] | None) -> dict[str, Any] | None:
        if not metadata:
            return None
        root_id = metadata.get("thread_root_event_id")
        if not isinstance(root_id, str) or not root_id:
            return None
        reply_to = metadata.get("thread_reply_to_event_id") or metadata.get("event_id")
        if not isinstance(reply_to, str) or not reply_to:
            return None
        return {"rel_type": "m.thread", "event_id": root_id,
                "m.in_reply_to": {"event_id": reply_to}, "is_falling_back": True}

    def _event_attachment_type(self, event: MatrixMediaEvent) -> str:
        msgtype = self._event_source_content(event).get("msgtype")
        return _MSGTYPE_MAP.get(cast(str, msgtype), "file")

    @staticmethod
    def _is_encrypted_media_event(event: MatrixMediaEvent) -> bool:
        return (isinstance(getattr(event, "key", None), dict)
                and isinstance(getattr(event, "hashes", None), dict)
                and isinstance(getattr(event, "iv", None), str))

    def _event_declared_size_bytes(self, event: MatrixMediaEvent) -> int | None:
        info = cast(object, self._event_source_content(event).get("info"))
        size = (
            cast(dict[str, Any], info).get("size")
            if isinstance(info, dict)
            else None
        )
        return size if type(size) is int and size >= 0 else None  # noqa: E721

    def _event_mime(self, event: MatrixMediaEvent) -> str | None:
        info = cast(object, self._event_source_content(event).get("info"))
        if (
            isinstance(info, dict)
            and isinstance(
                mime := cast(dict[str, Any], info).get("mimetype"),
                str,
            )
            and mime
        ):
            return mime
        mime = getattr(event, "mimetype", None)
        return mime if isinstance(mime, str) and mime else None

    def _event_filename(self, event: MatrixMediaEvent, attachment_type: str) -> str:
        body = getattr(event, "body", None)
        if isinstance(body, str) and body.strip():
            if candidate := safe_filename(Path(body).name):
                return candidate
        return _DEFAULT_ATTACH_NAME if attachment_type == "file" else attachment_type

    def _build_attachment_path(self, event: MatrixMediaEvent, attachment_type: str,
                               filename: str, mime: str | None) -> Path:
        safe_name = safe_filename(Path(filename).name) or _DEFAULT_ATTACH_NAME
        suffix = Path(safe_name).suffix
        if not suffix and mime:
            if guessed := mimetypes.guess_extension(mime, strict=False):
                safe_name, suffix = f"{safe_name}{guessed}", guessed
        stem = (Path(safe_name).stem or attachment_type)[:72]
        suffix = suffix[:16]
        event_id = safe_filename(str(getattr(event, "event_id", "") or "evt").lstrip("$"))
        event_prefix = (event_id[:24] or "evt").strip("_")
        return self._media_dir() / f"{event_prefix}_{stem}{suffix}"

    async def _download_media_bytes(self, mxc_url: str, limit_bytes: int) -> bytes | None:
        if not self.client or limit_bytes <= 0:
            raise _MediaTooLargeError

        parsed = urlparse(mxc_url)
        if parsed.scheme != "mxc" or not parsed.netloc or not parsed.path.strip("/"):
            return None

        homeserver = str(getattr(self.client, "homeserver", "") or self.config.homeserver).rstrip("/")
        media_url = (
            f"{homeserver}/_matrix/client/v1/media/download/"
            f"{quote(parsed.netloc, safe='')}/{quote(parsed.path.strip('/'), safe='')}"
        )
        token = getattr(self.client, "access_token", None) or self.config.access_token
        headers = {"Authorization": f"Bearer {token}"} if token else None
        timeout = aiohttp.ClientTimeout(total=None)

        try:
            async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session:
                async with session.get(media_url, params={"allow_remote": "true"}) as response:
                    if response.status >= 400:
                        self.logger.warning("download failed for {}: HTTP {}", mxc_url, response.status)
                        return None
                    content_length = response.headers.get("Content-Length")
                    if content_length is not None:
                        try:
                            if int(content_length) > limit_bytes:
                                raise _MediaTooLargeError
                        except ValueError:
                            pass

                    chunks = bytearray()
                    async for chunk in response.content.iter_chunked(64 * 1024):
                        chunks.extend(chunk)
                        if len(chunks) > limit_bytes:
                            raise _MediaTooLargeError
                    return bytes(chunks)
        except _MediaTooLargeError:
            raise
        except (aiohttp.ClientError, asyncio.TimeoutError, OSError):
            self.logger.warning("download failed for {}", mxc_url, exc_info=True)
            return None

    def _decrypt_media_bytes(self, event: MatrixMediaEvent, ciphertext: bytes) -> bytes | None:
        key_obj, hashes, iv = getattr(event, "key", None), getattr(event, "hashes", None), getattr(event, "iv", None)
        key = (
            cast(dict[str, Any], key_obj).get("k")
            if isinstance(key_obj, dict)
            else None
        )
        sha256 = (
            cast(dict[str, Any], hashes).get("sha256")
            if isinstance(hashes, dict)
            else None
        )
        if (
            not isinstance(key, str)
            or not isinstance(sha256, str)
            or not isinstance(iv, str)
        ):
            return None
        try:
            return decrypt_attachment(ciphertext, key, sha256, iv)
        except (EncryptionError, ValueError, TypeError):
            self.logger.warning("decrypt failed for event {}", getattr(event, "event_id", ""))
            return None

    async def _fetch_media_attachment(
        self, room: MatrixRoom, event: MatrixMediaEvent,
    ) -> tuple[dict[str, Any] | None, str]:
        """Download, decrypt if needed, and persist a Matrix attachment."""
        atype = self._event_attachment_type(event)
        mime = self._event_mime(event)
        filename = self._event_filename(event, atype)
        mxc_url = getattr(event, "url", None)
        fail = _ATTACH_FAILED.format(filename)

        if not isinstance(mxc_url, str) or not mxc_url.startswith("mxc://"):
            return None, fail

        limit_bytes = await self._effective_media_limit_bytes()
        declared = self._event_declared_size_bytes(event)
        if declared is None or declared > limit_bytes:
            return None, _ATTACH_TOO_LARGE.format(filename)

        try:
            async with self._media_download_semaphore:
                downloaded = await self._download_media_bytes(mxc_url, limit_bytes)
        except _MediaTooLargeError:
            return None, _ATTACH_TOO_LARGE.format(filename)
        if downloaded is None:
            return None, fail

        encrypted = self._is_encrypted_media_event(event)
        data = downloaded
        if encrypted:
            if (data := self._decrypt_media_bytes(event, downloaded)) is None:
                return None, fail

        if len(data) > limit_bytes:
            return None, _ATTACH_TOO_LARGE.format(filename)

        path = self._build_attachment_path(event, atype, filename, mime)
        try:
            path.write_bytes(data)
        except OSError:
            return None, fail

        attachment = {
            "type": atype, "mime": mime, "filename": filename,
            "event_id": str(getattr(event, "event_id", "") or ""),
            "encrypted": encrypted, "size_bytes": len(data),
            "path": str(path), "mxc_url": mxc_url,
        }
        return attachment, _ATTACH_MARKER.format(path)

    def _base_metadata(self, room: MatrixRoom, event: RoomMessage) -> dict[str, Any]:
        """Build common metadata for text and media handlers."""
        meta: dict[str, Any] = {"room": getattr(room, "display_name", room.room_id)}
        if isinstance(eid := getattr(event, "event_id", None), str) and eid:
            meta["event_id"] = eid
        if thread := self._thread_metadata(event):
            meta.update(thread)
        return meta

    async def _on_message(self, room: MatrixRoom, event: RoomMessageText) -> None:
        if (
            event.sender == self.config.user_id
            or self._is_pre_startup_event(event)
            or not self._should_process_message(room, event)
        ):
            return
        await self._start_typing_keepalive(room.room_id)
        try:
            await self._handle_message(
                sender_id=event.sender, chat_id=room.room_id,
                content=event.body, metadata=self._base_metadata(room, event),
                session_key=self._thread_session_key(room.room_id, event),
                is_dm=self._is_direct_room(room),
            )
        except Exception:
            await self._stop_typing_keepalive(room.room_id, clear_typing=True)
            raise

    async def _on_media_message(self, room: MatrixRoom, event: MatrixMediaEvent) -> None:
        if (
            event.sender == self.config.user_id
            or self._is_pre_startup_event(event)
            or not self._should_process_message(room, event)
        ):
            return
        attachment, marker = await self._fetch_media_attachment(room, event)
        parts: list[str] = []
        if isinstance(body := getattr(event, "body", None), str) and body.strip():
            parts.append(body.strip())

        if attachment and attachment.get("type") == "audio":
            transcription = await self.transcribe_audio(attachment["path"])
            if transcription:
                parts.append(f"[transcription: {transcription}]")
            else:
                parts.append(marker)
        elif marker:
            parts.append(marker)

        await self._start_typing_keepalive(room.room_id)
        try:
            meta = self._base_metadata(room, event)
            meta["attachments"] = []
            if attachment:
                meta["attachments"] = [attachment]
            await self._handle_message(
                sender_id=event.sender, chat_id=room.room_id,
                content="\n".join(parts),
                media=[attachment["path"]] if attachment else [],
                metadata=meta,
                session_key=self._thread_session_key(room.room_id, event),
                is_dm=self._is_direct_room(room),
            )
        except Exception:
            await self._stop_typing_keepalive(room.room_id, clear_typing=True)
            raise