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

from __future__ import annotations

import asyncio
import random
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any, TYPE_CHECKING
from collections.abc import Callable, Awaitable

if TYPE_CHECKING:
    from .send import MessageSender

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

STREAM_UPDATE_MIN_INTERVAL_MS = 1500
STREAM_MIN_INITIAL_CHARS = 20
STREAM_MAX_AGE_MS = 45000
STREAM_FALLBACK_CHUNK_SIZE = 2000
DEFAULT_STREAM_CHUNK_LIMIT = 4000

_STREAM_INFORMATIVE_TEXTS = [
    "Thinking about this...",
    "Working on it...",
    "Checking some details...",
    "Putting things together...",
]

_SIMPLE_INFORMATIVE_TEXTS = [
    "Let me think about that...",
    "One moment please...",
    "Looking into this...",
]


class StreamPhase(Enum):
    INIT = "init"
    INFORMATIVE = "informative"
    STREAMING = "streaming"
    FINAL = "final"
    CANCELLED = "cancelled"


@dataclass
class StreamState:
    phase: StreamPhase = StreamPhase.INIT
    stream_id: str = ""
    message_id: str = ""
    chat_id: str = ""
    accumulated_text: str = ""
    last_update_time: float = 0.0
    created_at: float = 0.0
    chunk_count: int = 0
    total_send_count: int = 0
    has_fallback: bool = False
    cancelled: bool = False


class ReplyStreamController:
    """流与 ReplyDispatcher 协调层。

    在 Agent 生成回复过程中，管理流的生命周期回调。
    注册 on_reply_start / on_partial_reply / prepare_payload / finalize。
    """

    def __init__(self):
        self._on_reply_start: Callable[[str], Awaitable[DeliveryResult]] | None = None
        self._on_partial_reply: Callable[[str, str], Awaitable[DeliveryResult]] | None = None
        self._on_finalize: Callable[[str, str, bool], Awaitable[DeliveryResult]] | None = None
        self._prepare_payload: Callable[[str], dict[str, Any]] | None = None
        self._active_chats: dict[str, str] = {}

    def set_callbacks(
        self,
        on_reply_start: Callable[[str], Awaitable[DeliveryResult]] | None = None,
        on_partial_reply: Callable[[str, str], Awaitable[DeliveryResult]] | None = None,
        on_finalize: Callable[[str, str, bool], Awaitable[DeliveryResult]] | None = None,
        prepare_payload: Callable[[str], dict[str, Any]] | None = None,
    ) -> None:
        self._on_reply_start = on_reply_start
        self._on_partial_reply = on_partial_reply
        self._on_finalize = on_finalize
        self._prepare_payload = prepare_payload

    async def start_reply(self, chat_id: str, initial_text: str) -> str | None:
        if self._on_reply_start:
            result = await self._on_reply_start(initial_text)
            if result.success and result.message_id:
                self._active_chats[chat_id] = result.message_id
                return result.message_id
        return None

    async def partial_reply(self, chat_id: str, text: str) -> DeliveryResult | None:
        if self._on_partial_reply:
            return await self._on_partial_reply(chat_id, text)
        return None

    async def finalize(self, chat_id: str, text: str, cancelled: bool = False) -> DeliveryResult | None:
        if self._on_finalize:
            result = await self._on_finalize(chat_id, text, cancelled)
            self._active_chats.pop(chat_id, None)
            return result
        return None

    def get_message_id(self, chat_id: str) -> str | None:
        return self._active_chats.get(chat_id)

    def prepare_activity(self, text: str) -> dict[str, Any]:
        if self._prepare_payload:
            return self._prepare_payload(text)
        return {"type": "message", "text": text, "textFormat": "markdown"}


class TeamsHttpStream:
    """TeamsHttpStream 三阶段流式传输。

    三阶段协议:
    1. informative — 蓝色进度条阶段，发送状态文本（随机选取）
    2. streaming — 流式编辑阶段，持续 update_activity
    3. final — 最终消息，流结束

    限制: 仅 personal (direct) 聊天支持流式传输，
    群组/频道聊天自动降级为块发送。

    失败回退: 流中断/超时自动降级为块发送。
    """

    def __init__(
        self,
        sender: MessageSender,
        throttle_ms: int = STREAM_UPDATE_MIN_INTERVAL_MS,
        min_initial_chars: int = STREAM_MIN_INITIAL_CHARS,
        max_stream_age_ms: int = STREAM_MAX_AGE_MS,
        chunk_limit: int = DEFAULT_STREAM_CHUNK_LIMIT,
        informative_texts: list[str] | None = None,
    ):
        self._sender = sender
        self._throttle_ms = throttle_ms
        self._min_initial_chars = min_initial_chars
        self._max_stream_age_ms = max_stream_age_ms
        self._chunk_limit = chunk_limit
        self._informative_texts = informative_texts or _STREAM_INFORMATIVE_TEXTS
        self._states: dict[str, StreamState] = {}
        self._reply_controller = ReplyStreamController()
        self._chat_types: dict[str, str] = {}

    @property
    def reply_controller(self) -> ReplyStreamController:
        return self._reply_controller

    def register_chat(self, chat_id: str, chat_type: str) -> None:
        self._chat_types[chat_id] = chat_type

    def supports_streaming(self, chat_id: str) -> bool:
        chat_type = self._chat_types.get(chat_id, "group")
        return chat_type == "direct"

    def _random_informative_text(self) -> str:
        return random.choice(self._informative_texts)

    async def send_informative(self, chat_id: str) -> DeliveryResult:
        state = self._states.get(chat_id)
        if state and state.phase == StreamPhase.STREAMING:
            return DeliveryResult(success=True, message_id=state.message_id)

        if not self.supports_streaming(chat_id):
            return DeliveryResult(success=False, error="Streaming not supported for this chat type, fallback to block")

        text = self._random_informative_text()
        activity = {
            "type": "message",
            "text": text,
            "textFormat": "markdown",
            "channelData": {"streamId": str(time.monotonic()).replace(".", ""), "streamType": "informative"},
            "entities": [
                {
                    "type": "streaminfo",
                    "streamId": str(time.monotonic()).replace(".", ""),
                    "streamType": "informative",
                }
            ],
        }

        result = await self._sender.send_activity(chat_id, activity)
        if result.success and result.message_id:
            now = time.monotonic()
            self._states[chat_id] = StreamState(
                phase=StreamPhase.INFORMATIVE,
                stream_id=str(int(now * 1000)),
                message_id=result.message_id,
                chat_id=chat_id,
                created_at=now,
                last_update_time=now,
            )
        return result

    async def stream_chunk(self, chat_id: str, chunk: str, finished: bool = False) -> DeliveryResult:
        state = self._states.get(chat_id)

        if state is None or state.cancelled:
            return DeliveryResult(success=False, error="No active stream")

        now = time.monotonic()
        stream_age_ms = (now - state.created_at) * 1000

        if stream_age_ms > self._max_stream_age_ms:
            logger.warning(f"MSTeams stream aged out for {chat_id} ({stream_age_ms:.0f}ms)")
            state.cancelled = True
            error_activity = {
                "type": "message",
                "text": "⏰ 响应生成超时，请稍后重试。",
                "textFormat": "markdown",
            }
            try:
                await self._sender.update_activity(state.chat_id, state.message_id, error_activity)
            except Exception:
                pass
            self._states.pop(chat_id, None)
            return DeliveryResult(success=False, error="Stream aged out, fallback to block")

        state.accumulated_text += chunk
        state.chunk_count += 1

        if state.phase == StreamPhase.INFORMATIVE:
            if len(state.accumulated_text) >= self._min_initial_chars:
                state.phase = StreamPhase.STREAMING
            else:
                return DeliveryResult(success=True, message_id=state.message_id)

        elapsed_since_update = (now - state.last_update_time) * 1000

        if elapsed_since_update < self._throttle_ms and not finished:
            state.total_send_count += 1
            return DeliveryResult(success=True, message_id=state.message_id)

        result = await self._send_stream_update(state, finished)
        state.last_update_time = time.monotonic()
        state.total_send_count += 1
        return result

    async def _send_stream_update(self, state: StreamState, finished: bool) -> DeliveryResult:
        text = state.accumulated_text[: self._chunk_limit]
        activity = {
            "type": "message",
            "text": text,
            "textFormat": "markdown",
        }

        if state.phase == StreamPhase.STREAMING:
            activity["entities"] = [
                {
                    "type": "streaminfo",
                    "streamId": state.stream_id,
                    "streamType": "streaming",
                }
            ]
        elif finished:
            activity["entities"] = [
                {
                    "type": "streaminfo",
                    "streamId": state.stream_id,
                    "streamType": "final",
                }
            ]

        result = await self._sender.update_activity(state.chat_id, state.message_id, activity)

        if finished:
            state.phase = StreamPhase.FINAL
            self._states.pop(state.chat_id, None)
        else:
            state.last_update_time = time.monotonic()

        return result

    async def finalize(self, chat_id: str) -> DeliveryResult:
        state = self._states.get(chat_id)
        if not state:
            return DeliveryResult(success=False, error="No active stream")

        result = await self._send_stream_update(state, finished=True)
        self._states.pop(chat_id, None)
        return result

    async def cancel(self, chat_id: str) -> None:
        state = self._states.get(chat_id)
        if state:
            state.cancelled = True
            cancel_activity = {
                "type": "message",
                "text": "🔄 响应已取消。",
                "textFormat": "markdown",
            }
            try:
                await self._sender.update_activity(state.chat_id, state.message_id, cancel_activity)
            except Exception:
                pass
            self._states.pop(chat_id, None)

    def get_message_id(self, chat_id: str) -> str | None:
        state = self._states.get(chat_id)
        return state.message_id if state else None

    def is_active(self, chat_id: str) -> bool:
        return chat_id in self._states and not self._states[chat_id].cancelled

    def clear(self) -> None:
        self._states.clear()

    @property
    def active_count(self) -> int:
        return len([s for s in self._states.values() if not s.cancelled])


class StreamManager:
    """兼容层：保持旧 StreamManager 接口，内部委托给 TeamsHttpStream。

    逐步迁移到 TeamsHttpStream。
    """

    def __init__(self, sender: MessageSender | None = None):
        self._messages: dict[str, str] = {}
        self._texts: dict[str, str] = {}
        self._last_update: dict[str, float] = {}
        self._locks: dict[str, asyncio.Lock] = {}
        self._stream: TeamsHttpStream | None = None
        self._sender: MessageSender | None = sender

    def _get_lock(self, chat_id: str) -> asyncio.Lock:
        if chat_id not in self._locks:
            self._locks[chat_id] = asyncio.Lock()
        return self._locks[chat_id]

    def set_sender(self, sender: MessageSender) -> None:
        self._sender = sender

    @property
    def http_stream(self) -> TeamsHttpStream | None:
        return self._stream

    def ensure_http_stream(self) -> TeamsHttpStream:
        if self._stream is None:
            if self._sender is None:
                raise RuntimeError("StreamManager: sender not set")
            self._stream = TeamsHttpStream(self._sender)
        return self._stream

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

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

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

    async def append_text(self, chat_id: str, chunk: str) -> str:
        async with self._get_lock(chat_id):
            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
        return elapsed_ms >= STREAM_UPDATE_MIN_INTERVAL_MS

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

    async def send_update(
        self,
        sender: MessageSender,
        chat_id: str,
        finished: bool = False,
        chunk_limit: int = DEFAULT_STREAM_CHUNK_LIMIT,
    ) -> DeliveryResult | None:
        async with self._get_lock(chat_id):
            message_id = self._messages.get(chat_id)
            text = self._texts.get(chat_id)
            if not message_id or text is None:
                return None

            activity = {
                "type": "message",
                "text": text[:chunk_limit],
                "textFormat": "markdown",
            }
            if finished:
                self._messages.pop(chat_id, None)
                self._texts.pop(chat_id, None)
                self._last_update.pop(chat_id, None)
                self._locks.pop(chat_id, None)
            else:
                self._last_update[chat_id] = time.monotonic()

        return await sender.update_activity(chat_id, message_id, activity)

    def clear(self) -> None:
        self._messages.clear()
        self._texts.clear()
        self._last_update.clear()
        self._locks.clear()
        if self._stream:
            self._stream.clear()

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

    def register_chat_type(self, chat_id: str, chat_type: str) -> None:
        stream = self.ensure_http_stream()
        stream.register_chat(chat_id, chat_type)

    def supports_streaming(self, chat_id: str) -> bool:
        if self._stream:
            return self._stream.supports_streaming(chat_id)
        return True


_DEFAULT_BLOCK_STREAMING_COALESCE = {
    "minChars": 1500,
    "idleMs": 1000,
}