ForcePilot/backend/package/yuxi/channel/extensions/mattermost/client.py

from __future__ import annotations

import asyncio
import logging
from typing import Any

import httpx

from yuxi.channel.extensions.mattermost.config import normalize_mattermost_base_url
from yuxi.channel.extensions.mattermost.errors import is_retryable_error, parse_mattermost_error
from yuxi.channel.extensions.mattermost.reconnect import with_jitter

logger = logging.getLogger(__name__)


class MattermostClient:
    def __init__(
        self,
        base_url: str,
        bot_token: str,
        allow_private_network: bool = False,
        timeout: float = 30.0,
    ):
        self.base_url = normalize_mattermost_base_url(base_url)
        self.api_url = f"{self.base_url}/api/v4"
        self.bot_token = bot_token

        limits = httpx.Limits(max_keepalive_connections=5, max_connections=20)
        transport = httpx.AsyncHTTPTransport(
            limits=limits,
            retries=0,
        )

        self._client = httpx.AsyncClient(
            base_url=self.api_url,
            headers={
                "Authorization": f"Bearer {bot_token}",
                "Content-Type": "application/json",
            },
            timeout=httpx.Timeout(timeout),
            transport=transport,
        )

    async def close(self) -> None:
        await self._client.aclose()

    async def _request(
        self,
        method: str,
        path: str,
        body: dict | None = None,
        params: dict | None = None,
    ) -> dict:
        response = await self._client.request(method, path, json=body, params=params)
        if response.status_code >= 400:
            raise parse_mattermost_error(response)
        return response.json() if response.text else {}

    async def _request_raw(
        self,
        method: str,
        path: str,
        body: dict | None = None,
        params: dict | None = None,
    ) -> httpx.Response:
        response = await self._client.request(method, path, json=body, params=params)
        return response

    # ── Users ────────────────────────────────────────────

    async def fetch_me(self) -> dict:
        return await self._request("GET", "/users/me")

    async def fetch_user(self, user_id: str) -> dict:
        return await self._request("GET", f"/users/{user_id}")

    async def fetch_user_by_username(self, username: str) -> dict:
        return await self._request("GET", f"/users/username/{username}")

    async def fetch_users(self, page: int = 0, per_page: int = 200, in_team: str = "") -> list[dict]:
        params: dict[str, Any] = {"page": page, "per_page": per_page}
        if in_team:
            params["in_team"] = in_team
        return await self._request("GET", "/users", params=params)

    async def fetch_user_teams(self, user_id: str) -> list[dict]:
        return await self._request("GET", f"/users/{user_id}/teams")

    async def update_user_status(self, user_id: str, status: str) -> dict:
        return await self._request("PUT", f"/users/{user_id}/status", body={"status": status})

    async def search_users(
        self,
        term: str,
        *,
        team_id: str = "",
        not_in_channel: str = "",
        not_in_team: str = "",
        page: int = 0,
        per_page: int = 60,
    ) -> list[dict]:
        body: dict[str, Any] = {"term": term}
        if team_id:
            body["team_id"] = team_id
        if not_in_channel:
            body["not_in_channel"] = not_in_channel
        if not_in_team:
            body["not_in_team"] = not_in_team
        params = {"page": page, "per_page": per_page}
        return await self._request("POST", "/users/search", body=body, params=params)

    async def fetch_user_preferences(self, user_id: str, category: str = "") -> list[dict]:
        params = {}
        if category:
            params["category"] = category
        return await self._request("GET", f"/users/{user_id}/preferences", params=params)

    async def update_user_preferences(
        self,
        user_id: str,
        preferences: list[dict],
    ) -> dict:
        return await self._request("PUT", f"/users/{user_id}/preferences", body=preferences)

    # ── Channels ─────────────────────────────────────────

    async def fetch_channel(self, channel_id: str) -> dict:
        return await self._request("GET", f"/channels/{channel_id}")

    async def fetch_channel_by_name(self, team_id: str, name: str) -> dict:
        return await self._request("GET", f"/teams/{team_id}/channels/name/{name}")

    async def fetch_team_channels(self, team_id: str, page: int = 0, per_page: int = 200) -> list[dict]:
        return await self._request(
            "GET",
            f"/teams/{team_id}/channels",
            params={"page": page, "per_page": per_page},
        )

    async def create_direct_channel(self, user_ids: list[str]) -> dict:
        return await self._request("POST", "/channels/direct", body=user_ids)

    async def create_group_channel(self, user_ids: list[str]) -> dict:
        return await self._request("POST", "/channels/group", body=user_ids)

    async def fetch_channels_for_user(self, user_id: str) -> list[dict]:
        return await self._request("GET", f"/users/{user_id}/channels")

    async def create_channel(
        self,
        team_id: str,
        name: str,
        display_name: str,
        *,
        channel_type: str = "O",
        purpose: str = "",
        header: str = "",
    ) -> dict:
        body = {
            "team_id": team_id,
            "name": name,
            "display_name": display_name,
            "type": channel_type,
            "purpose": purpose,
            "header": header,
        }
        return await self._request("POST", "/channels", body=body)

    async def update_channel(self, channel_id: str, payload: dict) -> dict:
        return await self._request("PUT", f"/channels/{channel_id}", body=payload)

    async def archive_channel(self, channel_id: str) -> dict:
        return await self._request("DELETE", f"/channels/{channel_id}")

    async def add_channel_member(self, channel_id: str, user_id: str) -> dict:
        body = {"user_id": user_id}
        return await self._request("POST", f"/channels/{channel_id}/members", body=body)

    async def remove_channel_member(self, channel_id: str, user_id: str) -> dict:
        return await self._request("DELETE", f"/channels/{channel_id}/members/{user_id}")

    async def get_channel_members(
        self,
        channel_id: str,
        page: int = 0,
        per_page: int = 60,
    ) -> list[dict]:
        return await self._request(
            "GET",
            f"/channels/{channel_id}/members",
            params={"page": page, "per_page": per_page},
        )

    async def mark_channel_viewed(self, channel_id: str) -> dict:
        body = {"channel_id": channel_id}
        return await self._request("POST", "/channels/members/me/view", body=body)

    async def set_channel_notify_props(
        self,
        channel_id: str,
        *,
        desktop: str = "default",
        email: str = "default",
        push: str = "default",
        mark_unread: str = "all",
        ignore_channel_mentions: str = "default",
    ) -> dict:
        body = {
            "channel_id": channel_id,
            "user_id": "me",
            "notify_props": {
                "desktop": desktop,
                "email": email,
                "push": push,
                "mark_unread": mark_unread,
                "ignore_channel_mentions": ignore_channel_mentions,
            },
        }
        return await self._request("PUT", f"/channels/{channel_id}/members/me/notify_props", body=body)

    async def search_channels(self, team_id: str, term: str) -> list[dict]:
        body: dict[str, Any] = {"term": term}
        return await self._request("POST", f"/teams/{team_id}/channels/search", body=body)

    async def search_all_channels(
        self,
        term: str,
        *,
        page: int = 0,
        per_page: int = 60,
    ) -> list[dict]:
        body: dict[str, Any] = {"term": term}
        params = {"page": page, "per_page": per_page}
        return await self._request("POST", "/channels/search", body=body, params=params)

    # ── Posts ────────────────────────────────────────────

    async def create_post(self, payload: dict) -> dict:
        return await self._request("POST", "/posts", body=payload)

    async def update_post(self, post_id: str, payload: dict) -> dict:
        return await self._request("PUT", f"/posts/{post_id}", body=payload)

    async def delete_post(self, post_id: str) -> dict:
        return await self._request("DELETE", f"/posts/{post_id}")

    async def fetch_post(self, post_id: str) -> dict:
        return await self._request("GET", f"/posts/{post_id}")

    async def fetch_post_thread(self, post_id: str, per_page: int = 60) -> dict:
        return await self._request("GET", f"/posts/{post_id}/thread", params={"perPage": per_page})

    async def search_posts(
        self,
        team_id: str,
        terms: str,
        *,
        is_or_search: bool = False,
        page: int = 0,
        per_page: int = 60,
    ) -> dict:
        body: dict[str, Any] = {
            "terms": terms,
            "is_or_search": is_or_search,
        }
        params = {"page": page, "per_page": per_page}
        return await self._request("POST", f"/teams/{team_id}/posts/search", body=body, params=params)

    async def create_ephemeral_post(
        self,
        user_id: str,
        channel_id: str,
        message: str,
        *,
        root_id: str = "",
        props: dict | None = None,
    ) -> dict:
        body: dict[str, Any] = {
            "user_id": user_id,
            "post": {
                "channel_id": channel_id,
                "message": message,
            },
        }
        if root_id:
            body["post"]["root_id"] = root_id
        if props:
            body["post"]["props"] = props
        return await self._request("POST", "/posts/ephemeral", body=body)

    async def pin_post(self, post_id: str) -> dict:
        return await self._request("POST", f"/posts/{post_id}/pin")

    async def unpin_post(self, post_id: str) -> dict:
        return await self._request("POST", f"/posts/{post_id}/unpin")

    # ── Files ────────────────────────────────────────────

    async def upload_file(
        self,
        channel_id: str,
        file_data: bytes,
        filename: str,
        mime_type: str = "application/octet-stream",
    ) -> dict:
        files = {"files": (filename, file_data, mime_type)}
        data = {"channel_id": channel_id}
        response = await self._client.post(
            "/files",
            data=data,
            files=files,
        )
        if response.status_code >= 400:
            raise parse_mattermost_error(response)
        return response.json()

    async def get_file(self, file_id: str) -> bytes:
        response = await self._client.get(f"/files/{file_id}")
        if response.status_code >= 400:
            raise parse_mattermost_error(response)
        return response.content

    # ── Reactions ────────────────────────────────────────

    async def add_reaction(self, user_id: str, post_id: str, emoji_name: str) -> dict:
        return await self._request(
            "POST",
            "/reactions",
            body={
                "user_id": user_id,
                "post_id": post_id,
                "emoji_name": emoji_name,
            },
        )

    async def remove_reaction(self, user_id: str, post_id: str, emoji_name: str) -> dict:
        return await self._request(
            "DELETE",
            f"/users/{user_id}/posts/{post_id}/reactions/{emoji_name}",
        )

    async def fetch_post_reactions(self, post_id: str) -> list[dict]:
        return await self._request("GET", f"/posts/{post_id}/reactions")

    # ── Teams ────────────────────────────────────────────

    async def fetch_team(self, team_id: str) -> dict:
        return await self._request("GET", f"/teams/{team_id}")

    # ── Typing ───────────────────────────────────────────

    async def send_typing(self, channel_id: str) -> None:
        await self._request("POST", "/users/me/typing", body={"channel_id": channel_id})

    # ── Commands ─────────────────────────────────────────

    async def list_commands(self, team_id: str, custom_only: bool = True) -> list[dict]:
        params: dict[str, Any] = {"custom_only": custom_only}
        return await self._request("GET", "/commands", params={**params, "team_id": team_id})

    async def create_command(self, payload: dict) -> dict:
        return await self._request("POST", "/commands", body=payload)

    async def update_command(self, command_id: str, payload: dict) -> dict:
        return await self._request("PUT", f"/commands/{command_id}", body=payload)

    async def delete_command(self, command_id: str) -> dict:
        return await self._request("DELETE", f"/commands/{command_id}")

    async def execute_command(self, channel_id: str, command: str) -> dict:
        return await self._request(
            "POST",
            "/commands/execute",
            body={"channel_id": channel_id, "command": command},
        )


async def create_direct_channel_with_retry(
    client: MattermostClient,
    user_ids: list[str],
    max_retries: int = 3,
    initial_delay_ms: int = 1000,
    max_delay_ms: int = 10000,
) -> dict:
    for attempt in range(max_retries + 1):
        try:
            return await client.create_direct_channel(user_ids)
        except Exception as e:
            if attempt >= max_retries:
                raise
            status_code = getattr(e, "status_code", 0)
            if not is_retryable_error(status_code):
                raise
            delay = min(initial_delay_ms * (2 ** attempt), max_delay_ms)
            delay = with_jitter(delay, 0.2)
            logger.warning(
                "DM channel creation retry %d/%d after error: %s. Waiting %d ms",
                attempt + 1, max_retries, e, delay,
            )
            await asyncio.sleep(delay / 1000)


async def send_webhook(hook_url: str, text: str, *, username: str = "", channel: str = "") -> dict:
    async with httpx.AsyncClient() as client:
        body: dict[str, Any] = {"text": text}
        if username:
            body["username"] = username
        if channel:
            body["channel"] = channel
        response = await client.post(hook_url, json=body)
        if response.status_code >= 400:
            raise parse_mattermost_error(response)
        return response.json() if response.text else {}