ForcePilot/backend/package/yuxi/channel/extensions/qqbot/streaming.py

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any

from yuxi.channel.extensions.qqbot.api_client import QQBotApiClient

logger = logging.getLogger(__name__)

STREAM_THROTTLE_MS = 500
STREAM_MIN_THROTTLE_MS = 300
STREAM_LONG_INTERVAL_MS = 2000
STREAM_LONG_BATCH_WINDOW_MS = 300


class FlushController:
    def __init__(self, throttle_ms: int = STREAM_THROTTLE_MS):
        self._ready = False
        self._flush_lock = asyncio.Lock()
        self._needs_reflush = False
        self._is_completed = False
        self._last_update_time = 0.0
        self._throttle_ms = throttle_ms

    @property
    def ready(self) -> bool:
        return self._ready

    @ready.setter
    def ready(self, value: bool) -> None:
        self._ready = value

    @property
    def is_completed(self) -> bool:
        return self._is_completed

    @is_completed.setter
    def is_completed(self, value: bool) -> None:
        self._is_completed = value

    @property
    def needs_reflush(self) -> bool:
        return self._needs_reflush

    @needs_reflush.setter
    def needs_reflush(self, value: bool) -> None:
        self._needs_reflush = value

    @property
    def last_update_time(self) -> float:
        return self._last_update_time

    @last_update_time.setter
    def last_update_time(self, value: float) -> None:
        self._last_update_time = value


class C2CStreamingController:
    def __init__(self, api_client: QQBotApiClient):
        self._api_client = api_client
        self._sessions: dict[str, FlushController] = {}
        self._stream_ids: dict[str, str] = {}
        self._pending_text: dict[str, str] = {}
        self._flush_tasks: dict[str, asyncio.Task] = {}
        self._msg_seqs: dict[str, int] = {}

    async def start_streaming_session(self, openid: str) -> str:
        if openid in self._stream_ids:
            return self._stream_ids[openid]

        data = await self._api_client.start_c2c_stream(openid)
        stream_id = data.get("stream_id", data.get("id", ""))
        self._stream_ids[openid] = stream_id
        self._pending_text[openid] = ""
        self._sessions[openid] = FlushController()
        logger.debug("Streaming session started: openid=%s, stream_id=%s", openid, stream_id)
        return stream_id

    async def feed(self, openid: str, text: str) -> None:
        if openid not in self._pending_text:
            return
        self._pending_text[openid] += text

        controller = self._sessions.get(openid)
        if controller:
            controller.last_update_time = time.time()
            controller.needs_reflush = True

        await self._schedule_flush(openid)

    async def flush(self, openid: str) -> None:
        text = self._pending_text.get(openid, "")
        stream_id = self._stream_ids.get(openid, "")
        controller = self._sessions.get(openid)

        if not text or not stream_id:
            return

        try:
            if not controller or not controller.ready:
                await self._api_client.send_c2c_stream_chunk(openid, stream_id, text)
                if controller:
                    controller.ready = True
            else:
                await self._api_client.send_c2c_stream_chunk(openid, stream_id, text)

            self._pending_text[openid] = ""

            if controller:
                controller.last_update_time = time.time()
                controller.needs_reflush = False
        except Exception as e:
            logger.warning("Stream flush failed for openid=%s: %s", openid, e)
            raise

    async def complete(self, openid: str) -> None:
        stream_id = self._stream_ids.pop(openid, "")
        if not stream_id:
            return

        controller = self._sessions.pop(openid, None)
        if controller:
            controller.is_completed = True

        pending = self._pending_text.pop(openid, "")
        if pending:
            try:
                await self._api_client.send_c2c_stream_chunk(openid, stream_id, pending)
            except Exception:
                logger.exception("Final stream flush failed for openid=%s", openid)

        await self._cancel_flush_task(openid)

        try:
            await self._api_client.complete_c2c_stream(openid, stream_id)
        except Exception:
            logger.exception("Stream completion failed for openid=%s", openid)

    async def abort(self, openid: str, reason: str | None = None) -> None:
        stream_id = self._stream_ids.pop(openid, "")
        if not stream_id:
            return

        self._sessions.pop(openid, None)
        self._pending_text.pop(openid, "")
        await self._cancel_flush_task(openid)

        try:
            await self._api_client.abort_c2c_stream(openid, stream_id)
        except Exception:
            logger.exception("Stream abort failed for openid=%s", openid)

    async def _schedule_flush(self, openid: str) -> None:
        controller = self._sessions.get(openid)
        if not controller or controller.is_completed:
            return

        await self._cancel_flush_task(openid)

        delay_ms = controller._throttle_ms
        now = time.time()
        elapsed = (now - controller.last_update_time) * 1000

        if elapsed > STREAM_LONG_INTERVAL_MS:
            delay_ms = STREAM_LONG_BATCH_WINDOW_MS
        else:
            delay_ms = max(STREAM_MIN_THROTTLE_MS, delay_ms)

        async def _delayed_flush():
            await asyncio.sleep(delay_ms / 1000.0)
            if openid in self._pending_text and controller and not controller.is_completed:
                try:
                    async with controller._flush_lock:
                        await self.flush(openid)
                except Exception:
                    pass

        self._flush_tasks[openid] = asyncio.create_task(_delayed_flush())

    async def _cancel_flush_task(self, openid: str) -> None:
        task = self._flush_tasks.pop(openid, None)
        if task and not task.done():
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

    async def fallback_to_static(self, openid: str, text: str) -> bool:
        await self.abort(openid)
        if text:
            await self._api_client.send_c2c_message(openid, text, msg_type=0)
        return True