ForcePilot/backend/package/yuxi/channels/adapters/googlechat/streaming.py

from __future__ import annotations

import time
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from googleapiclient.discovery import Resource

from yuxi.channels.models import DeliveryResult
from yuxi.utils.logging_config import logger

from .send import update_message, send_message, delete_message

DEFAULT_UPDATE_INTERVAL_MS = 800
DEFAULT_COALESCE_MIN_CHARS = 1500
DEFAULT_COALESCE_IDLE_MS = 1000


class StreamManager:
    def __init__(
        self,
        update_interval_ms: int = DEFAULT_UPDATE_INTERVAL_MS,
        coalesce_min_chars: int = DEFAULT_COALESCE_MIN_CHARS,
        coalesce_idle_ms: int = DEFAULT_COALESCE_IDLE_MS,
    ):
        self._messages: dict[str, str] = {}
        self._texts: dict[str, str] = {}
        self._last_update: dict[str, float] = {}
        self._update_interval_ms = update_interval_ms
        self._coalesce_min_chars = coalesce_min_chars
        self._coalesce_idle_ms = coalesce_idle_ms

        self._typing_messages: dict[str, str] = {}
        self._typing_replaced: set[str] = set()

    @property
    def update_interval_ms(self) -> int:
        return self._update_interval_ms

    def has_pending(self, chat_id: str) -> bool:
        return chat_id in self._messages

    def get_message_name(self, chat_id: str) -> str | None:
        return self._messages.get(chat_id)

    def get_accumulated_text(self, chat_id: str) -> str:
        return self._texts.get(chat_id, "")

    def register_message(self, chat_id: str, message_name: str, text: str) -> None:
        self._messages[chat_id] = message_name
        self._texts[chat_id] = text
        self._last_update[chat_id] = time.monotonic()

    def register_first_chunk(self, chat_id: str, message_name: str, text: str) -> None:
        self._messages[chat_id] = message_name
        self._texts[chat_id] = text
        self._last_update[chat_id] = time.monotonic()

        typing_id = self._typing_messages.pop(chat_id, None)
        if typing_id:
            self._typing_replaced.add(chat_id)
            logger.debug(f"Stream: replaced typing message {typing_id} with first chunk {message_name}")

    def append_text(self, chat_id: str, chunk: str) -> str:
        current = self._texts.get(chat_id, "")
        current += chunk
        self._texts[chat_id] = current
        return current

    def should_update(self, chat_id: str) -> bool:
        last = self._last_update.get(chat_id, 0)
        elapsed_ms = (time.monotonic() - last) * 1000

        if elapsed_ms >= self._update_interval_ms:
            return True

        if self._coalesce_idle_ms > 0:
            idle_threshold = max(self._coalesce_idle_ms, self._update_interval_ms)
            if elapsed_ms >= idle_threshold:
                new_chars_since_last = 0
                return new_chars_since_last >= self._coalesce_min_chars

        return False

    def mark_update(self, chat_id: str) -> None:
        self._last_update[chat_id] = time.monotonic()

    async def create_typing_message(
        self,
        chat_service: Resource,
        chat_id: str,
        bot_name: str = "Bot",
    ) -> str | None:
        if chat_id in self._typing_messages:
            return self._typing_messages[chat_id]

        try:
            body = {"text": f"*{bot_name} is typing...*"}
            result = await send_message(chat_service, chat_id, body)
            if result.success and result.message_id:
                self._typing_messages[chat_id] = result.message_id
                return result.message_id
        except Exception as e:
            logger.warning(f"Failed to create typing message in {chat_id}: {e}")

        return None

    async def clear_typing_message(
        self,
        chat_service: Resource,
        chat_id: str,
    ) -> None:
        msg_id = self._typing_messages.pop(chat_id, None)
        if msg_id and chat_id not in self._typing_replaced:
            try:
                await delete_message(chat_service, msg_id)
            except Exception as e:
                logger.warning(f"Failed to delete typing message {msg_id}: {e}")
        self._typing_replaced.discard(chat_id)

    async def handle_error(
        self,
        chat_service: Resource,
        chat_id: str,
    ) -> None:
        await self.clear_typing_message(chat_service, chat_id)
        self._messages.pop(chat_id, None)
        self._texts.pop(chat_id, None)
        self._last_update.pop(chat_id, None)

    async def finalize(
        self,
        chat_service: Resource,
        chat_id: str,
    ) -> DeliveryResult:
        await self.clear_typing_message(chat_service, chat_id)

        message_name = self._messages.pop(chat_id, None)
        text = self._texts.pop(chat_id, None)
        self._last_update.pop(chat_id, None)

        if not message_name or text is None:
            return DeliveryResult(success=False, error="No pending stream message")

        return await update_message(chat_service, message_name, {"text": text})

    async def send_update(
        self,
        chat_service: Resource,
        chat_id: str,
        finished: bool = False,
    ) -> DeliveryResult | None:
        message_name = self._messages.get(chat_id)
        text = self._texts.get(chat_id)
        if not message_name or text is None:
            return None

        if chat_id in self._typing_messages and text:
            typing_id = self._typing_messages.pop(chat_id, None)
            if typing_id:
                try:
                    await delete_message(chat_service, typing_id)
                except Exception:
                    pass
            self._typing_replaced.add(chat_id)

        body = {"text": text if finished else text + " ⏳"}
        result = await update_message(chat_service, message_name, body)

        if result.success:
            if finished:
                self._messages.pop(chat_id, None)
                self._texts.pop(chat_id, None)
                self._last_update.pop(chat_id, None)
                await self.clear_typing_message(chat_service, chat_id)
            else:
                self._last_update[chat_id] = time.monotonic()

        return result

    def clear(self) -> None:
        self._messages.clear()
        self._texts.clear()
        self._last_update.clear()
        self._typing_messages.clear()
        self._typing_replaced.clear()

    @property
    def pending_count(self) -> int:
        return len(self._messages)