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

"""Microsoft Teams 连接模式多样化。

支持 Webhook (默认)、WebSocket (Bot Framework Streaming Extensions)、
Polling 三种连接模式，可通过配置项 connection_mode 切换。
"""

from __future__ import annotations

import asyncio
import json
from typing import Any

import aiohttp

from yuxi.utils.logging_config import logger

CONNECTION_MODE_WEBHOOK = "webhook"
CONNECTION_MODE_WEBSOCKET = "websocket"
CONNECTION_MODE_POLLING = "polling"

VALID_CONNECTION_MODES = {CONNECTION_MODE_WEBHOOK, CONNECTION_MODE_WEBSOCKET, CONNECTION_MODE_POLLING}

BOT_STREAMING_URL_TEMPLATE = "https://directline.botframework.com/v3/directline/conversations/{conversation_id}/stream"
BOT_ACTIVITIES_URL_TEMPLATE = "https://smba.trafficmanager.net/emea/v3/conversations/{conversation_id}/activities"

POLLING_DEFAULT_INTERVAL_S = 2.0
POLLING_MAX_INTERVAL_S = 10.0
WEBSOCKET_PING_INTERVAL_S = 30.0
WEBSOCKET_PONG_TIMEOUT_S = 10.0


class WebSocketClient:
    """Bot Framework Streaming Extensions WebSocket 客户端。

    提供实时双向连接，延迟低于 Webhook 模式。
    """

    def __init__(
        self,
        app_id: str,
        app_password: str,
        stream_url: str = "",
        ping_interval: float = WEBSOCKET_PING_INTERVAL_S,
        pong_timeout: float = WEBSOCKET_PONG_TIMEOUT_S,
    ):
        self._app_id = app_id
        self._app_password = app_password
        self._stream_url = stream_url
        self._ping_interval = ping_interval
        self._pong_timeout = pong_timeout
        self._ws: aiohttp.ClientWebSocketResponse | None = None
        self._session: aiohttp.ClientSession | None = None
        self._running = False
        self._on_message: Any = None

    def set_message_handler(self, handler: Any) -> None:
        self._on_message = handler

    async def connect(self, conversation_id: str = "") -> None:
        if self._running:
            return

        self._session = aiohttp.ClientSession()
        stream_url = self._stream_url or BOT_STREAMING_URL_TEMPLATE.format(conversation_id=conversation_id)

        headers = {
            "Authorization": f"Bearer {self._app_password}",
            "Content-Type": "application/json",
        }

        self._ws = await self._session.ws_connect(stream_url, headers=headers, heartbeat=self._ping_interval)
        self._running = True
        logger.info(f"MSTeams WebSocket connected: {stream_url}")

        asyncio.create_task(self._read_loop())

    async def _read_loop(self) -> None:
        while self._running and self._ws is not None:
            try:
                msg = await self._ws.receive(timeout=self._pong_timeout)
                if msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        activities = data.get("activities", [])
                        for activity in activities:
                            if self._on_message:
                                await self._on_message(activity)
                    except json.JSONDecodeError:
                        logger.warning("MSTeams WebSocket: invalid JSON received")
                elif msg.type == aiohttp.WSMsgType.CLOSED:
                    logger.info("MSTeams WebSocket closed by server")
                    break
                elif msg.type == aiohttp.WSMsgType.ERROR:
                    logger.error(f"MSTeams WebSocket error: {self._ws.exception()}")
                    break
            except TimeoutError:
                continue
            except Exception as e:
                logger.error(f"MSTeams WebSocket read error: {e}")
                await asyncio.sleep(1)

        if self._running:
            logger.info("MSTeams WebSocket disconnected, attempting reconnect...")
            await asyncio.sleep(2)
            if self._running:
                asyncio.create_task(self._reconnect())

    async def _reconnect(self) -> None:
        try:
            await self.close()
            await asyncio.sleep(3)
            await self.connect()
        except Exception as e:
            logger.error(f"MSTeams WebSocket reconnect failed: {e}")

    async def send_activity(self, activity: dict[str, Any]) -> bool:
        if not self._ws or self._ws.closed:
            return False
        try:
            data = json.dumps(activity)
            await self._ws.send_str(data)
            return True
        except Exception as e:
            logger.error(f"MSTeams WebSocket send error: {e}")
            return False

    async def close(self) -> None:
        self._running = False
        if self._ws and not self._ws.closed:
            await self._ws.close()
            self._ws = None
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None


class PollingClient:
    """Bot Framework Connector API 轮询拉取客户端。

    适用于无法使用 Webhook 或 WebSocket 的部署场景。
    """

    def __init__(
        self,
        app_id: str,
        app_password: str,
        service_url: str = "",
        poll_interval: float = POLLING_DEFAULT_INTERVAL_S,
        max_interval: float = POLLING_MAX_INTERVAL_S,
    ):
        self._app_id = app_id
        self._app_password = app_password
        self._service_url = (service_url or "https://smba.trafficmanager.net/emea").rstrip("/")
        self._poll_interval = poll_interval
        self._max_interval = max_interval
        self._session: aiohttp.ClientSession | None = None
        self._running = False
        self._last_watermark: str = ""
        self._on_message: Any = None

    def set_message_handler(self, handler: Any) -> None:
        self._on_message = handler

    async def connect(self) -> None:
        if self._running:
            return
        self._session = aiohttp.ClientSession()
        self._running = True
        logger.info(f"MSTeams Polling started: interval={self._poll_interval}s")
        asyncio.create_task(self._poll_loop())

    async def _poll_loop(self) -> None:
        backoff = self._poll_interval

        while self._running:
            try:
                activities = await self._fetch_activities()
                if activities:
                    backoff = self._poll_interval
                    for activity in activities:
                        if self._on_message:
                            await self._on_message(activity)
                else:
                    backoff = min(backoff * 1.5, self._max_interval)
            except Exception as e:
                logger.warning(f"MSTeams Polling fetch error: {e}")
                backoff = min(backoff * 2, self._max_interval)

            await asyncio.sleep(backoff)

    async def _fetch_activities(self) -> list[dict[str, Any]]:
        if not self._session:
            return []

        url = BOT_ACTIVITIES_URL_TEMPLATE.format(conversation_id="all")
        headers = {
            "Authorization": f"Bearer {self._app_password}",
            "Content-Type": "application/json",
        }

        async with self._session.get(url, headers=headers) as resp:
            if resp.status == 200:
                data = await resp.json()
                activities = data.get("activities", [])
                self._last_watermark = data.get("watermark", self._last_watermark)
                return activities
            elif resp.status == 429:
                logger.warning("MSTeams Polling rate limited")
                return []
            else:
                body = await resp.text()
                logger.warning(f"MSTeams Polling HTTP {resp.status}: {body[:200]}")
                return []

    async def close(self) -> None:
        self._running = False
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None


class ConnectionModeManager:
    def __init__(self, config: dict[str, Any]):
        self._mode = config.get("connection_mode", CONNECTION_MODE_WEBHOOK)
        if self._mode not in VALID_CONNECTION_MODES:
            logger.warning(f"Invalid connection_mode '{self._mode}', falling back to 'webhook'")
            self._mode = CONNECTION_MODE_WEBHOOK

        self._ws_client: WebSocketClient | None = None
        self._poll_client: PollingClient | None = None
        self._message_router: Any = None

    @property
    def mode(self) -> str:
        return self._mode

    @property
    def is_webhook(self) -> bool:
        return self._mode == CONNECTION_MODE_WEBHOOK

    @property
    def is_websocket(self) -> bool:
        return self._mode == CONNECTION_MODE_WEBSOCKET

    @property
    def is_polling(self) -> bool:
        return self._mode == CONNECTION_MODE_POLLING

    def set_message_router(self, router: Any) -> None:
        self._message_router = router

    async def start(
        self,
        app_id: str,
        app_password: str,
        service_url: str = "",
        conversation_id: str = "",
    ) -> None:
        if self._mode == CONNECTION_MODE_WEBSOCKET:
            self._ws_client = WebSocketClient(app_id, app_password, stream_url=conversation_id)
            if self._message_router:
                self._ws_client.set_message_handler(self._message_router)
            await self._ws_client.connect(conversation_id)
            logger.info("MSTeams connection mode: WebSocket")
        elif self._mode == CONNECTION_MODE_POLLING:
            self._poll_client = PollingClient(app_id, app_password, service_url)
            if self._message_router:
                self._poll_client.set_message_handler(self._message_router)
            await self._poll_client.connect()
            logger.info("MSTeams connection mode: Polling")
        else:
            logger.info("MSTeams connection mode: Webhook")

    async def stop(self) -> None:
        if self._ws_client:
            await self._ws_client.close()
            self._ws_client = None
        if self._poll_client:
            await self._poll_client.close()
            self._poll_client = None

    async def send_activity(self, activity: dict[str, Any]) -> bool:
        if self._mode == CONNECTION_MODE_WEBSOCKET and self._ws_client:
            return await self._ws_client.send_activity(activity)
        return False