bot.py

"""X Chat delivery loop. Keep this stable; put product logic in handler.py."""

from __future__ import annotations

import logging
import queue
import time
from datetime import datetime, timedelta, timezone
from typing import Any

from xchat_bot.chat import (
    ChatCore,
    event_attachments,
    event_key_version,
    message_text,
    prep_to_request,
)
from xchat_bot.handler import Handler, Message, Reply
from xchat_bot.store import Store
from xchat_bot.x_api import RateLimited, XChatClient, _dump

logger = logging.getLogger("xchat_bot.bot")

RECENT_HOURS = 48
SWEEP_WITH_STREAM = 60.0
SWEEP_WITHOUT_STREAM = 8.0


def _created_at_key(item: dict[str, Any]) -> str:
    return str(item.get("created_at") or "")


def _parse_created_at(value: Any) -> datetime | None:
    if value is None:
        return None
    if isinstance(value, datetime):
        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
    text = str(value).strip()
    if not text:
        return None
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        dt = datetime.fromisoformat(text)
    except ValueError:
        return None
    return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)


def activity_to_page(payload: dict[str, Any]) -> dict[str, Any] | None:
    encoded = payload.get("encoded_event")
    key_change = payload.get("conversation_key_change_event")
    if not encoded and not key_change:
        return None
    created = str(payload.get("created_at") or "")
    if not created and payload.get("created_at_msec") is not None:
        try:
            msec = int(payload["created_at_msec"])
            created = datetime.fromtimestamp(msec / 1000, tz=timezone.utc).isoformat()
        except (TypeError, ValueError, OSError):
            created = ""
    item = {
        "id": payload.get("id"),
        "sender_id": payload.get("sender_id"),
        "conversation_id": str(payload.get("conversation_id") or "").replace(":", "-"),
        "encoded_event": encoded,
        "created_at": created,
    }
    meta: dict[str, Any] = {}
    if key_change:
        meta["conversation_key_events"] = (
            key_change if isinstance(key_change, list) else [key_change]
        )
    data = [item] if encoded else []
    return {"data": data, "meta": meta}


def _is_recent(item: dict[str, Any], hours: int = RECENT_HOURS) -> bool:
    dt = _parse_created_at(item.get("created_at"))
    if dt is None:
        return True
    return datetime.now(timezone.utc) - dt <= timedelta(hours=hours)


class Bot:
    def __init__(
        self,
        core: ChatCore,
        api: XChatClient,
        store: Store,
        handler: Handler,
        bot_user_id: str,
        extra_ids: list[str] | None = None,
    ) -> None:
        self.core = core
        self.api = api
        self.store = store
        self.handler = handler
        self.bot_user_id = bot_user_id
        self.extra_ids = extra_ids or []
        self._signing_keys: list[dict[str, str]] = []
        self._known_senders: set[str] = set()
        self._conv_keys: dict[str, dict[str, bytes]] = {}
        self.core.set_identity(bot_user_id)
        self.core.set_cache_keys(True)
        self._watch: set[str] = set()
        self._resolved_peers: set[str] = set()
        self._discovered: queue.Queue[dict[str, Any]] = queue.Queue()
        self._decrypt_skip_logged = False
        self._requests_logged = False
        self.stream_enabled = False
        self._until: dict[str, float] = {}
        self._next_sweep = 0.0
        self._need_poll: set[str] = set()
        for cid in store.watched_ids():
            self._watch.add(cid)

    def _remember_keys(self, conversation_id: str, keys: dict[str, bytes]) -> None:
        if not keys:
            return
        aliases = {conversation_id, conversation_id.replace(":", "-")}
        for alias in aliases:
            if alias:
                self._conv_keys.setdefault(alias, {}).update(keys)

    def _watch_conversation(self, conversation_id: str, peer_user_id: str | None = None) -> None:
        cid = conversation_id.replace(":", "-")
        if not cid or cid == self.bot_user_id:
            return
        if cid not in self._watch:
            logger.info("watching conv=%s peer=%s", cid, peer_user_id or "")
        self._watch.add(cid)
        self.store.watch_conversation(cid, peer_user_id)

    def _register_signing_keys(self, events: list[dict[str, Any]]) -> None:
        senders = {
            str(e.get("sender_id")) for e in events if e.get("sender_id")
        } - self._known_senders
        for sender_id in senders:
            try:
                for pk in self.api.get_public_keys(sender_id):
                    identity = pk.get("public_key") or ""
                    signing = pk.get("signing_public_key") or ""
                    sig = pk.get("identity_public_key_signature") or ""
                    if identity and signing and sig:
                        if not self.core.verify_key_binding(identity, signing, sig):
                            logger.warning("key_binding_failed sender=%s", sender_id)
                            continue
                    elif not signing:
                        continue
                    self._signing_keys.append(
                        {
                            "user_id": sender_id,
                            "public_key_version": str(pk.get("public_key_version") or ""),
                            "public_key": signing,
                            "identity_public_key": identity,
                            "identity_public_key_signature": sig,
                        }
                    )
                if any(k["user_id"] == sender_id and k.get("public_key") for k in self._signing_keys):
                    self._known_senders.add(sender_id)
                else:
                    logger.warning("no_signing_key sender=%s", sender_id)
            except Exception:
                logger.warning("public_keys_fetch_failed sender=%s", sender_id, exc_info=True)
        if senders:
            self.core.set_signing_keys(self._signing_keys)

    def _public_key_input(self, user_id: str) -> dict[str, str]:
        rows = self.api.get_public_keys(user_id)
        if not rows:
            raise RuntimeError(f"no public keys for {user_id}")
        row = rows[0]
        identity = row.get("public_key") or ""
        signing = row.get("signing_public_key") or ""
        sig = row.get("identity_public_key_signature") or ""
        if identity and signing and sig and not self.core.verify_key_binding(identity, signing, sig):
            raise RuntimeError(f"unverified public key for {user_id}")
        return {
            "user_id": user_id,
            "public_key": identity,
            "key_version": str(row.get("public_key_version") or ""),
        }

    def _ensure_keys(self, conversation_id: str, peer_id: str) -> None:
        try:
            prepared = self.core.prepare_conversation_key_change(
                [self._public_key_input(self.bot_user_id), self._public_key_input(peer_id)]
            )
            resp = self.api.add_conversation_keys(peer_id, prep_to_request(prepared))
            canonical = str((resp.get("data") or {}).get("conversation_id") or conversation_id)
            raw = prepared.get("conversation_key")
            version = prepared.get("conversation_key_version")
            if raw and version:
                self._remember_keys(canonical, {str(version): raw})
            logger.info("conversation_keys_added conv=%s", canonical)
        except Exception:
            logger.exception("conversation_keys_failed peer=%s", peer_id)

    def _ingest_page(
        self,
        conversation_id: str,
        page: dict[str, Any],
        *,
        reply: bool,
        only_recent: bool = False,
        reply_last_inbound: bool = False,
    ) -> None:
        raw: list[dict[str, Any]] = []
        for item in page.get("data") or []:
            dumped = item if isinstance(item, dict) else _dump(item)
            if not isinstance(dumped, dict):
                continue
            api_id = str(dumped.get("id") or "")
            event_conv = str(dumped.get("conversation_id") or conversation_id).replace(":", "-")
            if api_id and self.store.seen(event_conv, api_id):
                continue
            raw.append(dumped)
        raw.sort(key=_created_at_key)
        if not raw:
            return
        self._register_signing_keys(raw)
        key_events_b64 = (page.get("meta") or {}).get("conversation_key_events") or []
        events_b64 = list(key_events_b64) + [
            e["encoded_event"] for e in raw if e.get("encoded_event")
        ]
        if events_b64:
            batch = self.core.decrypt_batch(events_b64)
            keys = (batch.get("conversation_keys") or {}).get("keys") or {}
            self._remember_keys(conversation_id, keys)
            for item in raw:
                event_conv = str(item.get("conversation_id") or "").replace(":", "-")
                if event_conv:
                    self._remember_keys(event_conv, keys)

        last_inbound_id = ""
        if reply_last_inbound:
            for candidate in raw:
                sender = str(candidate.get("sender_id") or "")
                if sender and sender != self.bot_user_id and (
                    not only_recent or _is_recent(candidate)
                ):
                    last_inbound_id = str(candidate.get("id") or "")

        for item in raw:
            event_b64 = item.get("encoded_event")
            if not event_b64:
                continue
            event_conv = str(item.get("conversation_id") or conversation_id).replace(":", "-")
            api_id = str(item.get("id") or "")
            keyring = self._conv_keys.get(event_conv) or self._conv_keys.get(conversation_id, {})
            try:
                event = self.core.decrypt_one(event_b64, keyring)
            except Exception as err:
                if api_id:
                    self.store.mark_seen(event_conv, api_id)
                if not self._decrypt_skip_logged:
                    logger.info(
                        "skipping events that fail signature verification "
                        "(permanent; not retried)"
                    )
                    self._decrypt_skip_logged = True
                logger.debug("decrypt_skipped conv=%s id=%s err=%s", event_conv, api_id, err)
                continue
            if event.get("type") == "KeyChange":
                rotated = self.core.decrypt_batch([event_b64])
                rotated_keys = (rotated.get("conversation_keys") or {}).get("keys") or {}
                self._remember_keys(conversation_id, rotated_keys)
                self._remember_keys(event_conv, rotated_keys)
                if api_id:
                    self.store.mark_seen(event_conv, api_id)
                continue
            should_reply = reply and (not only_recent or _is_recent(item))
            if reply_last_inbound:
                should_reply = should_reply and api_id == last_inbound_id
            self._maybe_reply(event_conv or conversation_id, event, reply=should_reply)
            if api_id:
                self.store.mark_seen(event_conv, api_id)

    def pending(self) -> bool:
        return not self._discovered.empty()

    def _cooling(self, key: str = "global") -> bool:
        return time.monotonic() < self._until.get(key, 0.0)

    def _cool(self, key: str, seconds: float) -> None:
        until = time.monotonic() + max(1.0, seconds)
        self._until[key] = max(self._until.get(key, 0.0), until)

    def _on_rate_limited(self, err: RateLimited, conv: str | None = None) -> None:
        self._cool("global", err.retry_after)
        if conv:
            self._cool(conv, err.retry_after)
            self._need_poll.add(conv)
        logger.warning("rate_limited where=%s retry_after=%.0fs", err.where or conv, err.retry_after)

    def bootstrap(self, conversation_id: str) -> None:
        """Load keys. Reply to recent inbound; skip older history."""
        cur = self.store.cursor(conversation_id)
        if cur["bootstrapped"]:
            return
        first = self.api.get_events(conversation_id, max_results=100, all_pages=True)
        inbound = [
            item
            for item in (first.get("data") or [])
            if str((item if isinstance(item, dict) else _dump(item)).get("sender_id") or "")
            not in {"", self.bot_user_id}
        ]
        self._ingest_page(
            conversation_id, first, reply=True, only_recent=True, reply_last_inbound=True
        )
        self.store.set_cursor(conversation_id, None, True)
        logger.info("bootstrapped conv=%s inbound=%d", conversation_id, len(inbound))

    def poll_conversation(self, conversation_id: str) -> None:
        cid = conversation_id.replace(":", "-")
        if self._cooling("global") or self._cooling(cid):
            self._need_poll.add(cid)
            return
        try:
            self.bootstrap(cid)
            page = self.api.get_events(cid, max_results=50)
        except RateLimited as err:
            self._on_rate_limited(err, cid)
            return
        self._need_poll.discard(cid)
        for item in page.get("data") or []:
            dumped = item if isinstance(item, dict) else _dump(item)
            event_conv = str((dumped or {}).get("conversation_id") or cid).replace(":", "-")
            sender = str((dumped or {}).get("sender_id") or "")
            peer = sender if sender and sender != self.bot_user_id else None
            self._watch_conversation(event_conv, peer)
        self._ingest_page(cid, page, reply=True)

    def _maybe_reply(self, conversation_id: str, event: dict[str, Any], *, reply: bool) -> None:
        event_id = str(event.get("id") or event.get("message_id") or "")
        sender_id = str(event.get("sender_id") or "")
        conv = str(event.get("conversation_id") or conversation_id).replace(":", "-")
        if not event_id:
            return
        if self.store.seen(conv, event_id):
            return
        self.store.mark_seen(conv, event_id)
        if sender_id and sender_id != self.bot_user_id:
            self._watch_conversation(conv, sender_id)
        if not reply or sender_id == self.bot_user_id:
            return
        text = message_text(event)
        images = self._inbound_images(conv, event)
        if not text and not images:
            return
        try:
            self.api.indicate_typing(conv)
        except Exception:
            logger.debug("typing_failed conv=%s", conv, exc_info=True)
        inbound = Message(
            user_id=sender_id,
            conversation_id=conv,
            text=text,
            images=tuple(images),
        )
        try:
            answer = self.handler.handle(inbound)
        except Exception:
            logger.exception("handler_failed user=%s", sender_id)
            answer = Reply("Something glitched on my side. Try that again in a second.")
        self._send_reply(conv, sender_id, answer)

    def _inbound_images(self, conv: str, event: dict[str, Any]) -> list[bytes]:
        attachments = event_attachments(event)
        if not attachments:
            return []
        keyring = self._conv_keys.get(conv) or self._conv_keys.get(conv.replace("-", ":"), {})
        version = event_key_version(event)
        raw_key = keyring.get(version) if version else None
        if raw_key is None:
            latest = self._latest_key(conv)
            raw_key = latest[1] if latest else None
        if raw_key is None:
            logger.warning("inbound_image_no_key conv=%s", conv)
            return []
        images: list[bytes] = []
        for att in attachments:
            hash_key = str(att.get("media_hash_key") or att.get("mediaHashKey") or "")
            if not hash_key:
                continue
            try:
                from chat_xdk import detect_mime_type

                cipher = self.api.download_chat_media(conv, hash_key)
                plain = self.core.decrypt_stream(cipher, raw_key)
                mime = str(detect_mime_type(plain) or "")
                if mime.startswith("image/"):
                    images.append(plain)
            except Exception:
                logger.warning("inbound_image_failed conv=%s hash=%s", conv, hash_key, exc_info=True)
        return images

    def _latest_key(self, conversation_id: str) -> tuple[str, bytes] | None:
        keys = self._conv_keys.get(conversation_id) or self._conv_keys.get(
            conversation_id.replace("-", ":"), {}
        )
        if not keys:
            return None

        def _ver(value: str) -> int:
            try:
                return int(value)
            except ValueError:
                return 0

        version = max(keys, key=_ver)
        return str(version), keys[version]

    def _encrypt(
        self,
        conv: str,
        sender_id: str,
        text: str,
        attachments: list[dict[str, Any]] | None = None,
        conversation_key: bytes | None = None,
        conversation_key_version: str | None = None,
    ) -> dict[str, str] | None:
        try:
            return self.core.encrypt_message(
                conv,
                text,
                conversation_key=conversation_key,
                conversation_key_version=conversation_key_version,
                attachments=attachments,
            )
        except ValueError:
            logger.warning("no_conversation_key conv=%s — establishing", conv)
            self._ensure_keys(conv, sender_id)
            try:
                return self.core.encrypt_message(
                    conv,
                    text,
                    conversation_key=conversation_key,
                    conversation_key_version=conversation_key_version,
                    attachments=attachments,
                )
            except Exception:
                logger.exception("encrypt_failed conv=%s", conv)
                return None
        except Exception:
            logger.exception("encrypt_failed conv=%s", conv)
            return None

    def _send_text(self, conv: str, sender_id: str, text: str) -> bool:
        body = self._encrypt(conv, sender_id, text)
        if not body:
            return False
        try:
            self.api.send_message(conv, body)
            logger.info("replied user=%s len=%d", sender_id, len(text))
            return True
        except Exception:
            logger.exception("send_failed conv=%s", conv)
            return False

    def _send_reply(self, conv: str, sender_id: str, reply: Reply) -> bool:
        if not reply.text and not reply.image:
            return False
        if reply.image is None:
            return self._send_text(conv, sender_id, reply.text)
        return self._send_image(conv, sender_id, reply.text or " ", reply.image)

    def _send_image(self, conv: str, sender_id: str, text: str, image: bytes) -> bool:
        key = self._latest_key(conv)
        if key is None:
            self._ensure_keys(conv, sender_id)
            key = self._latest_key(conv)
        if key is None:
            logger.warning("image_no_key conv=%s", conv)
            return self._send_text(conv, sender_id, text)
        version, raw_key = key
        try:
            from chat_xdk import detect_image_dimensions, detect_mime_type

            ciphertext = self.core.encrypt_stream(image, raw_key)
            media_hash_key = self.api.upload_chat_media(conv, ciphertext)
            dims = detect_image_dimensions(image) or (1024, 1024)
            mime = str(detect_mime_type(image) or "image/jpeg")
            ext = "png" if "png" in mime else "jpg"
            attachments = [
                {
                    "attachment_type": "media",
                    "media_hash_key": media_hash_key,
                    "width": dims[0],
                    "height": dims[1],
                    "filesize_bytes": len(image),
                    "filename": f"image.{ext}",
                }
            ]
            body = self.core.encrypt_message(
                conv,
                text,
                conversation_key=raw_key,
                conversation_key_version=version,
                attachments=attachments,
            )
            self.api.send_message(conv, body)
            logger.info("image_sent user=%s bytes=%d", sender_id, len(image))
            return True
        except Exception:
            logger.exception("image_send_failed conv=%s", conv)
            return self._send_text(conv, sender_id, text)

    def _already_watching_peer(self, peer_id: str) -> bool:
        if peer_id in self._watch:
            return True
        return any(peer_id in cid.split("-") for cid in self._watch)

    def _resolve_peer_threads(self) -> None:
        pending = list(dict.fromkeys([*self.extra_ids, *self.store.peer_ids()]))
        for extra in pending:
            if extra == self.bot_user_id or extra in self._resolved_peers:
                continue
            if self._already_watching_peer(extra):
                self._resolved_peers.add(extra)
                continue
            if self._cooling("global"):
                return
            try:
                canonical = self.api.canonical_conversation_id(extra)
            except RateLimited as err:
                self._on_rate_limited(err)
                return
            except Exception:
                logger.warning("peer_conversation_lookup_failed peer=%s", extra, exc_info=True)
                continue
            if canonical:
                self._watch_conversation(canonical, extra)
                self._resolved_peers.add(extra)

    def note_activity(self, event: dict[str, Any]) -> None:
        data = event.get("data") if isinstance(event.get("data"), dict) else event
        payload = data.get("payload") if isinstance(data, dict) else None
        if not isinstance(payload, dict):
            payload = data if isinstance(data, dict) else {}
        conv = str(payload.get("conversation_id") or "").replace(":", "-")
        sender = str(payload.get("sender_id") or "")
        event_type = str((data or {}).get("event_type") or "")
        event_uuid = str((data or {}).get("event_uuid") or "")
        if conv or sender:
            logger.info("activity %s conv=%s sender=%s", event_type or "chat", conv, sender)
            self._discovered.put(
                {
                    "conv": conv,
                    "sender": sender or None,
                    "event_uuid": event_uuid,
                    "payload": payload,
                }
            )

    def _handle_discovered(self, item: dict[str, Any]) -> None:
        conv = str(item.get("conv") or "").replace(":", "-")
        sender = item.get("sender")
        payload = item.get("payload") if isinstance(item.get("payload"), dict) else {}
        event_uuid = str(item.get("event_uuid") or "")
        if conv:
            self._watch_conversation(conv, sender if isinstance(sender, str) else None)
        if sender and not self._already_watching_peer(str(sender)):
            if self._cooling("global"):
                self._need_poll.add(conv or str(sender))
            else:
                try:
                    canonical = self.api.canonical_conversation_id(str(sender))
                    if canonical:
                        self._watch_conversation(canonical, str(sender))
                        conv = conv or canonical
                except RateLimited as err:
                    self._on_rate_limited(err, conv or str(sender))
                    return
                except Exception:
                    logger.warning("discovered_peer_lookup_failed peer=%s", sender, exc_info=True)
        if not conv:
            return
        if event_uuid and self.store.seen(conv, f"uuid:{event_uuid}"):
            return
        bootstrapped = bool(self.store.cursor(conv)["bootstrapped"])
        page = activity_to_page(payload)
        if bootstrapped and page:
            if event_uuid:
                self.store.mark_seen(conv, f"uuid:{event_uuid}")
            self._ingest_page(conv, page, reply=True)
            return
        self.poll_conversation(conv)

    def _drain_discovered(self) -> None:
        while True:
            try:
                item = self._discovered.get_nowait()
            except queue.Empty:
                return
            try:
                self._handle_discovered(item)
            except RateLimited as err:
                conv = str(item.get("conv") or "")
                self._on_rate_limited(err, conv or None)
            except Exception:
                logger.exception("discovered_event_failed conv=%s", item.get("conv"))

    def _retry_needed_polls(self) -> None:
        if self._cooling("global"):
            return
        for cid in list(self._need_poll):
            if self._cooling(cid):
                continue
            try:
                self.poll_conversation(cid)
            except RateLimited as err:
                self._on_rate_limited(err, cid)
                return
            except Exception:
                logger.exception("poll_conversation_failed conv=%s", cid)

    def poll_inbox(self) -> None:
        self._drain_discovered()
        self._retry_needed_polls()
        now = time.monotonic()
        if self._cooling("global") or now < self._next_sweep:
            return
        sweep_every = SWEEP_WITH_STREAM if self.stream_enabled else SWEEP_WITHOUT_STREAM
        self._next_sweep = now + sweep_every
        try:
            ids, meta = self.api.conversation_ids()
        except RateLimited as err:
            self._on_rate_limited(err)
            return
        except Exception:
            logger.exception("list_conversations_failed")
            return
        for cid in ids:
            self._watch_conversation(cid)
        self._resolve_peer_threads()
        if meta.get("has_message_requests") and not self._requests_logged:
            logger.info(
                "message_requests_pending inbox=%d watching=%d — discovering "
                "those threads via chat.received on the activity stream",
                len(ids),
                len(self._watch),
            )
            self._requests_logged = True
        logger.info(
            "sweep inbox=%d watching=%d requests=%s extras=%d",
            len(ids),
            len(self._watch),
            meta.get("has_message_requests"),
            len(self.extra_ids),
        )
        for cid in list(self._watch):
            try:
                self.poll_conversation(cid)
            except RateLimited as err:
                self._on_rate_limited(err, cid)
                return
            except Exception:
                logger.exception("poll_conversation_failed conv=%s", cid)
            time.sleep(0.35)