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

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

import httpx

from yuxi.channel.extensions.zulip.errors import (
    ZulipAPIError,
    ZulipAuthError,
    ZulipQueueExpiredError,
    ZulipRateLimitError,
)

logger = logging.getLogger(__name__)


class ZulipAsyncClient:
    def __init__(
        self,
        realm_url: str,
        bot_email: str,
        bot_api_key: str,
        timeout: float = 90.0,
        connect_timeout: float = 10.0,
    ):
        self.realm_url = realm_url.rstrip("/")
        self.api_url = f"{self.realm_url}/api/v1"
        self.bot_email = bot_email

        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,
            auth=(bot_email, bot_api_key),
            timeout=httpx.Timeout(timeout, connect=connect_timeout),
            transport=transport,
        )

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

    async def _request(
        self,
        method: str,
        path: str,
        data: dict | None = None,
        params: dict | None = None,
        files: dict | None = None,
        timeout: httpx.Timeout | None = None,
    ) -> dict:
        kwargs: dict[str, Any] = {"method": method, "url": path}
        if data:
            kwargs["data"] = data
        if params:
            kwargs["params"] = params
        if files:
            kwargs["files"] = files
        if timeout:
            kwargs["timeout"] = timeout

        resp = await self._client.request(**kwargs)
        if resp.status_code >= 400:
            try:
                body = resp.json()
            except Exception:
                body = {"code": "UNKNOWN", "msg": resp.text}

            if resp.status_code == 401:
                raise ZulipAuthError(resp.status_code, body)
            if resp.status_code == 429:
                raise ZulipRateLimitError(resp.status_code, body)
            if body.get("code") == "BAD_EVENT_QUEUE_ID":
                raise ZulipQueueExpiredError(resp.status_code, body)

            raise ZulipAPIError(resp.status_code, body)
        return resp.json()

    async def register_queue(
        self,
        event_types: list[str] | None = None,
        narrow: list | None = None,
        fetch_event_types: list[str] | None = None,
    ) -> dict:
        data: dict[str, str] = {}
        if event_types:
            data["event_types"] = json.dumps(event_types)
        if narrow:
            data["narrow"] = json.dumps(narrow)
        if fetch_event_types:
            data["fetch_event_types"] = json.dumps(fetch_event_types)
        return await self._request("POST", "/register", data=data)

    async def get_events(
        self,
        queue_id: str,
        last_event_id: int,
        dont_block: bool = False,
    ) -> dict:
        params: dict[str, Any] = {
            "queue_id": queue_id,
            "last_event_id": last_event_id,
        }
        if dont_block:
            params["dont_block"] = "true"
        return await self._request(
            "GET",
            "/events",
            params=params,
            timeout=httpx.Timeout(90.0),
        )

    async def delete_queue(self, queue_id: str) -> dict:
        return await self._request("DELETE", "/events", data={"queue_id": queue_id})

    async def send_message(self, payload: dict) -> dict:
        return await self._request("POST", "/messages", data=payload)

    async def update_message(
        self,
        message_id: int,
        content: str,
        propagate_mode: str = "change_one",
    ) -> dict:
        data: dict[str, Any] = {
            "content": content,
            "propagate_mode": propagate_mode,
        }
        return await self._request("PATCH", f"/messages/{message_id}", data=data)

    async def delete_message(self, message_id: int) -> dict:
        return await self._request("DELETE", f"/messages/{message_id}")

    async def upload_file(self, file_path: str) -> dict:
        path = Path(file_path)
        with open(path, "rb") as f:
            files = {"file": (path.name, f)}
            return await self._request("POST", "/user_uploads", files=files)

    async def add_reaction(self, message_id: int, emoji_name: str) -> dict:
        data = {
            "emoji_name": emoji_name,
            "reaction_type": "unicode_emoji",
        }
        return await self._request("POST", f"/messages/{message_id}/reactions", data=data)

    async def remove_reaction(self, message_id: int, emoji_name: str) -> dict:
        params = {
            "emoji_name": emoji_name,
            "reaction_type": "unicode_emoji",
        }
        return await self._request("DELETE", f"/messages/{message_id}/reactions", params=params)

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

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

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

    async def get_streams(self, include_public: bool = True) -> dict:
        params: dict[str, Any] = {}
        if include_public:
            params["include_public"] = "true"
        return await self._request("GET", "/streams", params=params)

    async def get_stream_id(self, stream_name: str) -> dict:
        return await self._request("GET", "/get_stream_id", params={"stream": stream_name})

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

    async def get_messages(
        self,
        narrow: list[dict] | None = None,
        anchor: str | None = None,
        num_before: int = 0,
        num_after: int = 20,
        include_anchor: bool = True,
        apply_markdown: bool = True,
        client_gravatar: bool = True,
    ) -> dict:
        params: dict[str, Any] = {
            "num_before": num_before,
            "num_after": num_after,
            "include_anchor": str(include_anchor).lower(),
            "apply_markdown": str(apply_markdown).lower(),
            "client_gravatar": str(client_gravatar).lower(),
        }
        if anchor:
            params["anchor"] = anchor
        if narrow:
            params["narrow"] = json.dumps(narrow)
        return await self._request("GET", "/messages", params=params)

    async def send_typing(
        self,
        op: str,
        to: list[str] | str,
        topic: str | None = None,
    ) -> dict:
        data: dict[str, Any] = {"op": op, "to": to}
        if topic:
            data["topic"] = topic
        return await self._request("POST", "/typing", data=data)

    async def get_bot_storage(self, key: str | None = None) -> dict:
        params: dict[str, str] = {}
        if key:
            params["key"] = key
        return await self._request("GET", "/bot_storage", params=params)

    async def put_bot_storage(self, key: str, value: str) -> dict:
        return await self._request("PUT", "/bot_storage", data={"key": key, "value": value})

    async def delete_bot_storage(self, key: str) -> dict:
        return await self._request("DELETE", "/bot_storage", data={"key": key})

    async def get_message(
        self,
        message_id: int,
        apply_markdown: bool = True,
    ) -> dict:
        params = {"apply_markdown": str(apply_markdown).lower()}
        return await self._request("GET", f"/messages/{message_id}", params=params)

    async def get_message_history(self, message_id: int) -> dict:
        return await self._request("GET", f"/messages/{message_id}/history")

    async def subscribe(self, subscriptions: list[dict]) -> dict:
        data = {"subscriptions": json.dumps(subscriptions)}
        return await self._request("POST", "/users/me/subscriptions", data=data)

    async def unsubscribe(self, stream_names: list[str]) -> dict:
        return await self._request(
            "DELETE",
            "/users/me/subscriptions",
            data={"subscriptions": json.dumps(stream_names)},
        )

    async def get_topics(self, stream_id: int) -> dict:
        return await self._request("GET", f"/users/me/{stream_id}/topics")

    async def get_stream_members(self, stream_id: int) -> dict:
        return await self._request("GET", f"/streams/{stream_id}/members")

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

    async def update_message_flags(
        self,
        messages: list[int],
        op: str,
        flag: str,
    ) -> dict:
        data = {
            "messages": json.dumps(messages),
            "op": op,
            "flag": flag,
        }
        return await self._request("POST", "/messages/flags", data=data)

    async def mark_all_as_read(self) -> dict:
        return await self._request("POST", "/mark_all_as_read", data={})

    async def mark_stream_as_read(self, stream_id: int) -> dict:
        return await self._request("POST", "/mark_stream_as_read", data={"stream_id": stream_id})

    async def mark_topic_as_read(self, stream_id: int, topic_name: str) -> dict:
        return await self._request(
            "POST",
            "/mark_topic_as_read",
            data={"stream_id": stream_id, "topic_name": topic_name},
        )

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

    async def update_bot_status(
        self,
        status_text: str = "",
        away: bool = False,
        emoji_name: str = "",
        emoji_code: str = "",
    ) -> dict:
        data: dict[str, Any] = {}
        if status_text:
            data["status_text"] = status_text
        if away:
            data["away"] = str(away).lower()
        if emoji_name:
            data["emoji_name"] = emoji_name
        if emoji_code:
            data["emoji_code"] = emoji_code
        return await self._request("POST", "/users/me/status", data=data)

    async def get_realm_emoji(self) -> dict:
        return await self._request("GET", "/realm/emoji")

    async def get_server_settings(self) -> dict:
        return await self._request("GET", "/server_settings")

    async def get_stream_by_id(self, stream_id: int) -> dict:
        return await self._request("GET", f"/streams/{stream_id}")

    def get_file_url(self, realm_id_str: str, filename: str) -> str:
        return f"{self.realm_url}/user_uploads/{realm_id_str}/{filename}"

    async def report_message(self, message_id: int, reason: str) -> dict:
        return await self._request(
            "POST",
            "/report/message",
            data={"message_id": message_id, "reason": reason},
        )

    async def create_stream(
        self,
        name: str,
        description: str = "",
        invite_only: bool = False,
        history_public_to_subscribers: bool = False,
    ) -> dict:
        data = {
            "subscriptions": json.dumps([{"name": name, "description": description}]),
            "invite_only": str(invite_only).lower(),
            "history_public_to_subscribers": str(history_public_to_subscribers).lower(),
        }
        return await self._request("POST", "/users/me/subscriptions", data=data)

    async def update_stream(
        self,
        stream_id: int,
        *,
        name: str | None = None,
        description: str | None = None,
        is_private: bool | None = None,
    ) -> dict:
        data: dict[str, Any] = {}
        if name is not None:
            data["new_name"] = name
        if description is not None:
            data["description"] = description
        if is_private is not None:
            data["is_private"] = str(is_private).lower()
        return await self._request("PATCH", f"/streams/{stream_id}", data=data)

    async def archive_stream(self, stream_id: int) -> dict:
        return await self._request("DELETE", f"/streams/{stream_id}")

    async def get_stream_email(self, stream_id: int) -> dict:
        return await self._request("GET", f"/streams/{stream_id}/email_address")

    async def mute_topic(
        self,
        stream_name: str,
        topic: str,
        op: str = "add",
    ) -> dict:
        return await self._request(
            "PATCH",
            "/users/me/subscriptions/muted_topics",
            data={"stream": stream_name, "topic": topic, "op": op},
        )

    async def update_user_topic(
        self,
        stream_id: int,
        topic: str,
        visibility_policy: int,
    ) -> dict:
        return await self._request(
            "POST",
            "/user_topics",
            data={
                "stream_id": stream_id,
                "topic": topic,
                "visibility_policy": visibility_policy,
            },
        )

    async def delete_topic(self, stream_id: int, topic_name: str) -> dict:
        return await self._request(
            "DELETE",
            f"/streams/{stream_id}/delete_topic",
            data={"topic_name": topic_name},
        )

    async def get_subscription_status(self, user_id: int, stream_id: int) -> dict:
        return await self._request("GET", f"/users/{user_id}/subscriptions/{stream_id}")