ForcePilot/backend/package/yuxi/channel/extensions/msteams/gateway.py

from __future__ import annotations

import asyncio
import logging
import os
from pathlib import Path

import uvicorn

from .auth import verify_botframework_jwt
from .config import MSTeamsConfigAdapter
from .conversation_store import ConversationStore
from .dedupe import MSTeamsDedupeStore
from .file_consent import FileConsentStore, FileConsentRequest
from .graph import MSTeamsGraphClient
from .monitor import activity_to_unified_message, is_echo_activity
from .polls import PollStore
from .proactive import save_conversation_reference
from .sdk import BotFrameworkAdapter
from .sent_message_cache import SentMessageCache
from .sso import SSOTokenStore, exchange_token_on_behalf_of
from .types import BotFrameworkActivity

logger = logging.getLogger(__name__)

DEFAULT_WEBHOOK_HOST = "0.0.0.0"
DEFAULT_WEBHOOK_PORT = 3978


class MSTeamsGateway:
    def __init__(
        self,
        config_adapter: MSTeamsConfigAdapter,
        account: dict,
        global_config: dict,
    ):
        self._config_adapter = config_adapter
        self._account = account
        self._global_config = global_config
        self._account_id = account.get("account_id", "default")
        self._app_id = account.get("app_id", "")
        self._tenant_id = account.get("tenant_id", "")
        self._running = False
        self._server_task: asyncio.Task | None = None
        self._stop_event = asyncio.Event()
        self._ctx = None

        channel_cfg = account.get("config", {}) or {}
        webhook_cfg = channel_cfg.get("webhook", {}) or {}
        self._webhook_host = webhook_cfg.get("host", DEFAULT_WEBHOOK_HOST)
        self._webhook_port = webhook_cfg.get("port", DEFAULT_WEBHOOK_PORT)

        state_dir = channel_cfg.get("stateDir") or global_config.get("stateDir") or "./data/msteams"
        os.makedirs(state_dir, exist_ok=True)
        store_path = Path(state_dir) / f"conversations-{self._account_id}.json"
        self._conversation_store = ConversationStore(str(store_path))
        self._conversation_store.load()

        poll_path = Path(state_dir) / f"polls-{self._account_id}.json"
        self._poll_store = PollStore(str(poll_path))
        self._poll_store.load()

        self._file_consent_store = FileConsentStore()
        self._sso_token_store = SSOTokenStore()
        self._sent_message_cache = SentMessageCache()
        self._dedupe_store = MSTeamsDedupeStore()

        graph_cfg = channel_cfg.get("graph", {}) or {}
        self._graph_client: MSTeamsGraphClient | None = None
        if graph_cfg.get("enabled") and graph_cfg.get("client_id"):
            self._graph_client = MSTeamsGraphClient(
                tenant_id=self._tenant_id,
                client_id=graph_cfg.get("client_id", ""),
                client_secret=graph_cfg.get("client_secret", ""),
            )

        self._adapter = BotFrameworkAdapter(
            app_id=self._app_id,
            app_password=account.get("app_password", ""),
            tenant_id=self._tenant_id,
        )

    @property
    def adapter(self) -> BotFrameworkAdapter:
        return self._adapter

    @property
    def conversation_store(self) -> ConversationStore:
        return self._conversation_store

    @property
    def poll_store(self) -> PollStore:
        return self._poll_store

    @property
    def file_consent_store(self) -> FileConsentStore:
        return self._file_consent_store

    @property
    def sso_token_store(self) -> SSOTokenStore:
        return self._sso_token_store

    @property
    def sent_message_cache(self) -> SentMessageCache:
        return self._sent_message_cache

    @property
    def dedupe_store(self) -> MSTeamsDedupeStore:
        return self._dedupe_store

    @property
    def graph_client(self) -> MSTeamsGraphClient | None:
        return self._graph_client

    @property
    def app_id(self) -> str:
        return self._app_id

    @property
    def account_id(self) -> str:
        return self._account_id

    async def start(self, ctx=None) -> asyncio.Task:
        if self._running:
            raise RuntimeError("MSTeams gateway already running")

        self._ctx = ctx
        self._running = True
        self._stop_event.clear()

        self._server_task = asyncio.create_task(self._run_server())
        logger.info(
            "MS Teams gateway started for account '%s' on %s:%d",
            self._account_id,
            self._webhook_host,
            self._webhook_port,
        )
        return self._server_task

    async def stop(self) -> None:
        if not self._running:
            return

        self._running = False
        self._stop_event.set()

        if self._server_task:
            self._server_task.cancel()
            try:
                await self._server_task
            except asyncio.CancelledError:
                pass

        self._conversation_store.save()
        self._poll_store.save()
        self._sent_message_cache.cleanup()
        logger.info("MS Teams gateway stopped for account '%s'", self._account_id)

    async def _run_server(self):
        async def jwt_verifier(token: str) -> dict:
            return await verify_botframework_jwt(token, self._app_id, self._tenant_id)

        async def handle_activity(activity_data: dict) -> dict:
            activity = BotFrameworkActivity.from_dict(activity_data)

            if activity.type == "invoke":
                return await self._handle_invoke(activity)

            if is_echo_activity(activity, self._app_id):
                return {"status": 200}

            try:
                await save_conversation_reference(self._conversation_store, activity, self._app_id)
            except Exception:
                logger.exception("Failed to save conversation reference")

            unified_msg = activity_to_unified_message(activity, self._app_id, self._account_id)

            if unified_msg is not None:
                if self._ctx and self._ctx.queue:
                    try:
                        self._ctx.queue.put_nowait(unified_msg)
                    except Exception:
                        logger.exception("Failed to enqueue MS Teams message")

            return {"status": 200}

        from .webhook import create_webhook_app

        app = create_webhook_app(
            bot_app_id=self._app_id,
            adapter=self._adapter,
            conversation_store=self._conversation_store,
            message_handler=handle_activity,
            jwt_verifier=jwt_verifier,
        )

        config = uvicorn.Config(
            app=app,
            host=self._webhook_host,
            port=self._webhook_port,
            log_level="warning",
            access_log=False,
        )
        server = uvicorn.Server(config)

        try:
            await server.serve()
        except asyncio.CancelledError:
            pass
        finally:
            self._conversation_store.save()
            self._poll_store.save()
            self._sent_message_cache.cleanup()

    async def _handle_invoke(self, activity: BotFrameworkActivity) -> dict:
        name = activity.name or ""

        if name == "adaptiveCard/action":
            return await self._handle_adaptive_card_action(activity)

        if name == "task/fetch":
            from .task_modules import handle_task_fetch
            return await handle_task_fetch(activity, self._ctx)

        if name == "task/submit":
            from .task_modules import handle_task_submit
            return await handle_task_submit(activity, self._ctx)

        if name == "dialog/fetch":
            from .task_modules import handle_dialog_fetch
            return await handle_dialog_fetch(activity, self._ctx)

        if name == "dialog/submit":
            from .task_modules import handle_dialog_submit
            return await handle_dialog_submit(activity, self._ctx)

        if name.startswith("composeExtension/"):
            from .message_extension import handle_compose_extension
            return await handle_compose_extension(name, activity, self._ctx)

        if name == "signin/tokenExchange":
            return await self._handle_token_exchange(activity)

        if name == "signin/verifyState":
            return await self._handle_signin_verify_state(activity)

        return {"status": 200}

    async def _handle_meeting_event(self, activity: BotFrameworkActivity) -> None:
        from .meetings import parse_meeting_event

        event_data = parse_meeting_event(activity.raw)
        if not event_data:
            return

        if self._ctx and hasattr(self._ctx, "queue") and self._ctx.queue:
            try:
                from yuxi.channel.message.models import MessageType, PeerInfo, UnifiedMessage
                from yuxi.channel.routing.models import PeerKind

                unified_msg = UnifiedMessage(
                    msg_id=f"msteams:meeting:{activity.id}",
                    channel_type="msteams",
                    account_id=getattr(self._ctx, "account_id", ""),
                    content="",
                    sender=PeerInfo(id=activity.from_id, kind=PeerKind.DIRECT),
                    message_type=MessageType.EVENT,
                    raw_payload=activity.raw,
                    metadata={
                        "event_type": "meeting_event",
                        "meeting_id": event_data.get("meeting_id"),
                        "participant_count": event_data.get("participant_count"),
                    },
                    conversation_label=f"msteams:{activity.conversation_id}",
                    native_channel_id=activity.conversation_id,
                )
                self._ctx.queue.put_nowait(unified_msg)
            except Exception:
                logger.exception("Failed to enqueue meeting event")

    async def _handle_token_exchange(self, activity: BotFrameworkActivity) -> dict:
        value = activity.value or {}
        token = value.get("token", "")
        if not token:
            return {"status": 400, "body": {"message": "No token provided"}}

        try:
            result = await exchange_token_on_behalf_of(
                self._tenant_id,
                self._app_id,
                self._account.get("app_password", ""),
                token,
                "https://graph.microsoft.com/.default",
            )
            if result:
                return {"status": 200, "body": {"id": activity.id, "token": result.get("access_token")}}
            return {"status": 403, "body": {"message": "Token exchange failed"}}
        except Exception as e:
            logger.exception("SSO token exchange failed")
            return {"status": 500, "body": {"message": str(e)}}

    async def _handle_signin_verify_state(self, activity: BotFrameworkActivity) -> dict:
        value = activity.value or {}
        state = value.get("state", "")

        stored = self._sso_token_store.verify_state(state)
        if stored:
            return {"status": 200, "body": {"valid": True}}
        return {"status": 403, "body": {"valid": False, "message": "State expired or invalid"}}

    async def _handle_adaptive_card_action(self, activity: BotFrameworkActivity) -> dict:
        value = activity.value or {}
        action = value.get("action", "")
        data = value.get("data", value)

        if action == "file_consent_accept":
            return await self._handle_file_consent(action="accept", data=data)
        elif action == "file_consent_decline":
            return await self._handle_file_consent(action="decline", data=data)
        elif action == "sso_signin":
            return await self._handle_sso_signin_consent(data=data, activity=activity)
        elif isinstance(data, dict) and data.get("action") in ("poll_vote",):
            return await self._handle_poll_vote(data=data, activity=activity)

        return {"status": 200}

    async def _handle_file_consent(self, action: str, data: dict) -> dict:
        consent_id = data.get("consent_id", "")
        if not consent_id:
            return {"status": 400}

        req = self._file_consent_store.get(consent_id)
        if not req:
            return {"status": 200, "body": {"message": "Consent request expired"}}

        if action == "accept":
            self._file_consent_store.accept(consent_id)
            if req.file_url and self._graph_client:
                asyncio.create_task(self._upload_consented_file(req))
        elif action == "decline":
            self._file_consent_store.decline(consent_id)

        return {"status": 200}

    async def _upload_consented_file(self, req: FileConsentRequest) -> None:
        from .files import download_public_file
        from .graph_upload import upload_to_onedrive
        from .sdk import build_message_activity

        try:
            file_bytes = await download_public_file(req.file_url)
            if not file_bytes:
                raise RuntimeError("Failed to download file")

            if req.upload_path:
                result = await upload_to_onedrive(
                    self._graph_client,
                    file_bytes,
                    req.file_name,
                    folder=req.upload_path,
                )
            else:
                result = await upload_to_onedrive(
                    self._graph_client,
                    file_bytes,
                    req.file_name,
                )

            if result and result.get("web_url"):
                await self._adapter.send_to_conversation(
                    req.service_url,
                    req.conversation_id,
                    build_message_activity(
                        f"✅ 文件 **{req.file_name}** ({_format_size(len(file_bytes))}) 上传完成\n"
                        f"[查看文件]({result['web_url']})",
                        reply_to_id=req.reply_to_id,
                        tenant_id=req.tenant_id,
                    ),
                )
            else:
                await self._adapter.send_to_conversation(
                    req.service_url,
                    req.conversation_id,
                    build_message_activity(
                        f"✅ 文件 **{req.file_name}** ({_format_size(len(file_bytes))}) 上传完成",
                        reply_to_id=req.reply_to_id,
                        tenant_id=req.tenant_id,
                    ),
                )
        except Exception as e:
            logger.exception("File consent upload failed")
            await self._adapter.send_to_conversation(
                req.service_url,
                req.conversation_id,
                build_message_activity(
                    f"❌ 文件上传失败: {e}",
                    reply_to_id=req.reply_to_id,
                    tenant_id=req.tenant_id,
                ),
            )

    async def _handle_sso_signin_consent(self, data: dict, activity: BotFrameworkActivity) -> dict:
        state = data.get("state", "")
        if not state:
            return {"status": 400}

        stored = self._sso_token_store.verify_state(state)
        if not stored:
            return {"status": 200, "body": {"message": "SSO state expired or invalid"}}

        return {"status": 200}

    async def _handle_poll_vote(self, data: dict, activity: BotFrameworkActivity) -> dict:
        poll_id = data.get("poll_id", "")
        selected = data.get("selected", [])
        if isinstance(selected, str):
            selected = [selected]

        sender_id = activity.from_id
        from .polls import handle_poll_vote

        result = await handle_poll_vote(
            self._adapter,
            self._poll_store,
            {"poll_id": poll_id, "selected": selected},
            sender_id=sender_id,
        )

        if not result.get("success"):
            return {"status": 200, "body": {"message": result.get("error", "Vote failed")}}

        poll = self._poll_store.get(poll_id)
        if poll:
            from .polls import build_poll_results_text

            text = build_poll_results_text(poll)
            activity_data: dict = {
                "type": "message",
                "text": f"✅ 投票成功！当前结果:\n\n{text}",
            }
            if poll.tenant_id:
                activity_data.setdefault("channelData", {})
                activity_data["channelData"]["tenant"] = {"id": poll.tenant_id}

            await self._adapter.send_to_conversation(
                poll.service_url,
                poll.conversation_id,
                activity_data,
            )

        return {"status": 200}

    async def probe(self) -> bool:
        try:
            token_ok = await self._adapter.probe()
            return token_ok
        except Exception as e:
            logger.warning("MS Teams probe failed: %s", e)
            return False