ForcePilot/backend/package/yuxi/channels/adapters/slack/send.py

from __future__ import annotations

import asyncio
from collections import defaultdict
from typing import Any

from slack_sdk.errors import SlackApiError

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

_SCOPE_CUSTOMIZE_ERRORS = frozenset({"missing_scope", "invalid_arguments", "not_allowed_token_type"})


def _is_customize_scope_error(error: str) -> bool:
    return error in _SCOPE_CUSTOMIZE_ERRORS or "chat:write.customize" in error


class SendQueue:
    def __init__(self):
        self._locks: dict[str, asyncio.Lock] = {}
        self._queues: dict[str, list[asyncio.Task]] = defaultdict(list)

    def _make_key(
        self,
        account_id: str = "",
        token: str = "",
        recipient: str = "",
        thread_ts: str | None = None,
    ) -> str:
        return f"{account_id}:{token[:12] if token else ''}:{recipient}:{thread_ts or ''}"

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

    async def enqueue(
        self,
        account_id: str = "",
        token: str = "",
        recipient: str = "",
        thread_ts: str | None = None,
    ) -> asyncio.Lock:
        key = self._make_key(account_id, token, recipient, thread_ts)
        lock = self._get_lock(key)
        await lock.acquire()
        return lock

    def release(self, lock: asyncio.Lock) -> None:
        try:
            lock.release()
        except RuntimeError:
            pass

    def clear(self) -> None:
        self._locks.clear()
        self._queues.clear()


_send_queue: SendQueue | None = None


def get_send_queue() -> SendQueue:
    global _send_queue
    if _send_queue is None:
        _send_queue = SendQueue()
    return _send_queue


async def send_message(
    client,
    channel: str,
    text: str,
    *,
    thread_ts: str | None = None,
    blocks: list | None = None,
    mrkdwn: bool = True,
    unfurl_links: bool = True,
    username: str | None = None,
    icon_url: str | None = None,
    icon_emoji: str | None = None,
    account_id: str = "",
    token: str = "",
    use_queue: bool = False,
) -> DeliveryResult:
    if use_queue:
        queue = get_send_queue()
        lock = await queue.enqueue(account_id, token, channel, thread_ts)
        try:
            return await _send_message_impl(
                client,
                channel,
                text,
                thread_ts=thread_ts,
                blocks=blocks,
                mrkdwn=mrkdwn,
                unfurl_links=unfurl_links,
                username=username,
                icon_url=icon_url,
                icon_emoji=icon_emoji,
            )
        finally:
            queue.release(lock)

    return await _send_message_impl(
        client,
        channel,
        text,
        thread_ts=thread_ts,
        blocks=blocks,
        mrkdwn=mrkdwn,
        unfurl_links=unfurl_links,
        username=username,
        icon_url=icon_url,
        icon_emoji=icon_emoji,
    )


async def _send_message_impl(
    client,
    channel: str,
    text: str,
    *,
    thread_ts: str | None = None,
    blocks: list | None = None,
    mrkdwn: bool = True,
    unfurl_links: bool = True,
    username: str | None = None,
    icon_url: str | None = None,
    icon_emoji: str | None = None,
) -> DeliveryResult:
    try:
        params: dict[str, Any] = {
            "channel": channel,
            "text": text,
            "mrkdwn": mrkdwn,
            "unfurl_links": unfurl_links,
        }
        if thread_ts:
            params["thread_ts"] = thread_ts
        if blocks:
            params["blocks"] = blocks

        has_customize = False
        if username:
            params["username"] = username
            has_customize = True
        if icon_url:
            params["icon_url"] = icon_url
            has_customize = True
        if icon_emoji:
            params["icon_emoji"] = icon_emoji
            has_customize = True

        result = await client.chat_postMessage(**params)
        ok = result.get("ok", False)
        if not ok and has_customize and _is_customize_scope_error(result.get("error", "")):
            return await _retry_without_customize(client, channel, text, thread_ts, blocks, mrkdwn, unfurl_links)

        return DeliveryResult(
            success=ok,
            message_id=result.get("ts"),
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        err = e.response.get("error", str(e))
        if _is_customize_scope_error(err):
            return await _retry_without_customize(client, channel, text, thread_ts, blocks, mrkdwn, unfurl_links)
        logger.error(f"Slack send_message failed: {err}")
        return DeliveryResult(
            success=False,
            error=err,
        )


async def _retry_without_customize(
    client,
    channel: str,
    text: str,
    thread_ts: str | None,
    blocks: list | None,
    mrkdwn: bool,
    unfurl_links: bool,
) -> DeliveryResult:
    logger.info("chat:write.customize scope missing, retrying without customize params")
    try:
        params: dict[str, Any] = {
            "channel": channel,
            "text": text,
            "mrkdwn": mrkdwn,
            "unfurl_links": unfurl_links,
        }
        if thread_ts:
            params["thread_ts"] = thread_ts
        if blocks:
            params["blocks"] = blocks

        result = await client.chat_postMessage(**params)
        ok = result.get("ok", False)
        return DeliveryResult(
            success=ok,
            message_id=result.get("ts"),
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        err = e.response.get("error", str(e))
        logger.error(f"Slack send_message retry failed: {err}")
        return DeliveryResult(
            success=False,
            error=err,
        )


async def send_blocks(
    client,
    channel: str,
    blocks: list,
    *,
    text: str = "",
    thread_ts: str | None = None,
    unfurl_links: bool = True,
) -> DeliveryResult:
    try:
        params: dict[str, Any] = {
            "channel": channel,
            "blocks": blocks,
            "text": text or "(Block Kit message)",
            "mrkdwn": True,
            "unfurl_links": unfurl_links,
        }
        if thread_ts:
            params["thread_ts"] = thread_ts

        result = await client.chat_postMessage(**params)
        ok = result.get("ok", False)
        return DeliveryResult(
            success=ok,
            message_id=result.get("ts"),
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        logger.error(f"Slack send_blocks failed: {e.response.get('error', e)}")
        return DeliveryResult(
            success=False,
            error=e.response.get("error", str(e)),
        )


async def send_ephemeral(
    client,
    channel: str,
    user: str,
    text: str,
    *,
    blocks: list | None = None,
    thread_ts: str | None = None,
) -> DeliveryResult:
    try:
        params: dict[str, Any] = {
            "channel": channel,
            "user": user,
            "text": text,
        }
        if blocks:
            params["blocks"] = blocks
        if thread_ts:
            params["thread_ts"] = thread_ts

        result = await client.chat_postEphemeral(**params)
        ok = result.get("ok", False)
        return DeliveryResult(
            success=ok,
            message_id=result.get("message_ts"),
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        logger.error(f"Slack send_ephemeral failed: {e.response.get('error', e)}")
        return DeliveryResult(
            success=False,
            error=e.response.get("error", str(e)),
        )


async def update_message(
    client,
    channel: str,
    ts: str,
    text: str,
    *,
    blocks: list | None = None,
    mrkdwn: bool = True,
) -> DeliveryResult:
    try:
        params: dict[str, Any] = {
            "channel": channel,
            "ts": ts,
            "text": text,
            "mrkdwn": mrkdwn,
        }
        if blocks:
            params["blocks"] = blocks

        result = await client.chat_update(**params)
        ok = result.get("ok", False)
        return DeliveryResult(
            success=ok,
            message_id=result.get("ts"),
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        logger.error(f"Slack update_message failed: {e.response.get('error', e)}")
        return DeliveryResult(
            success=False,
            error=e.response.get("error", str(e)),
        )


async def delete_message(
    client,
    channel: str,
    ts: str,
) -> DeliveryResult:
    try:
        result = await client.chat_delete(channel=channel, ts=ts)
        ok = result.get("ok", False)
        return DeliveryResult(
            success=ok,
            error=result.get("error") if not ok else None,
        )
    except SlackApiError as e:
        logger.error(f"Slack delete_message failed: {e.response.get('error', e)}")
        return DeliveryResult(
            success=False,
            error=e.response.get("error", str(e)),
        )