ForcePilot/backend/package/yuxi/channels/adapters/whatsapp/bridge.py

from __future__ import annotations

import asyncio
import os
import secrets
import signal
from typing import Any

import aiohttp

from yuxi.channels.infra.external_process import ExternalProcessManager
from yuxi.channels.models import DeliveryResult, HealthStatus
from yuxi.utils.logging_config import logger

DEFAULT_BRIDGE_PORT = 9080
DEFAULT_AUTH_DIR = "~/.forcepilot/whatsapp/auth"
DEFAULT_NODE_PATH = "node"
DEFAULT_CONNECT_TIMEOUT_MS = 20000
DEFAULT_QUERY_TIMEOUT_MS = 30000


class BaileysBridge(ExternalProcessManager):
    RETRY_MAX = 3
    RETRY_BASE_DELAY = 1.0

    def __init__(self, config: dict[str, Any] | None = None):
        self._config = config or {}
        self._port = self._config.get("bridge_port", DEFAULT_BRIDGE_PORT)
        self._auth_dir = os.path.expanduser(self._config.get("auth_dir", DEFAULT_AUTH_DIR))
        self._node_path = self._config.get("node_path", DEFAULT_NODE_PATH)
        self._connect_timeout_ms = self._config.get("connect_timeout_ms", DEFAULT_CONNECT_TIMEOUT_MS)
        self._query_timeout_ms = self._config.get("query_timeout_ms", DEFAULT_QUERY_TIMEOUT_MS)
        self._bridge_dir = os.path.join(os.path.dirname(__file__), "baileys-bridge")
        self._process: asyncio.subprocess.Process | None = None
        self._http_session: aiohttp.ClientSession | None = None
        self._bridge_token: str = secrets.token_hex(32)
        self._retry_statuses: set[int] = {429, 500, 502, 503, 504}

    async def _request_with_retry(
        self, method: str, path: str, json_data: dict[str, Any] | None = None
    ) -> tuple[bool, Any]:
        if self._http_session is None:
            return False, {"error": "Bridge not started"}

        url = f"{self.base_url}{path}"
        last_error = None
        for attempt in range(self.RETRY_MAX):
            try:
                kwargs = {"headers": self._auth_headers}
                if json_data is not None:
                    kwargs["json"] = json_data
                async with getattr(self._http_session, method)(url, **kwargs) as resp:
                    if resp.status in self._retry_statuses and attempt < self.RETRY_MAX - 1:
                        await asyncio.sleep(self.RETRY_BASE_DELAY * (2**attempt))
                        continue
                    data = await resp.json()
                    return resp.status == 200, data
            except aiohttp.ClientConnectorError as e:
                last_error = str(e)
                if attempt < self.RETRY_MAX - 1:
                    await asyncio.sleep(self.RETRY_BASE_DELAY * (2**attempt))
            except Exception as e:
                last_error = str(e)
                if attempt < self.RETRY_MAX - 1:
                    await asyncio.sleep(self.RETRY_BASE_DELAY * (2**attempt))
        return False, {"error": last_error or "request failed after retries"}

    @property
    def base_url(self) -> str:
        return f"http://localhost:{self._port}"

    @property
    def _auth_headers(self) -> dict[str, str]:
        return {"X-Bridge-Token": self._bridge_token}

    async def start(self) -> None:
        os.makedirs(self._auth_dir, exist_ok=True)

        env = os.environ.copy()
        env["WHATSAPP_AUTH_DIR"] = self._auth_dir
        env["WHATSAPP_BRIDGE_PORT"] = str(self._port)
        env["WHATSAPP_BRIDGE_TOKEN"] = self._bridge_token
        env["WHATSAPP_CONNECT_TIMEOUT_MS"] = str(self._connect_timeout_ms)
        env["WHATSAPP_QUERY_TIMEOUT_MS"] = str(self._query_timeout_ms)

        if proxy := os.environ.get("HTTPS_PROXY", os.environ.get("https_proxy", "")):
            env["HTTPS_PROXY"] = proxy
        if proxy := os.environ.get("HTTP_PROXY", os.environ.get("http_proxy", "")):
            env["HTTP_PROXY"] = proxy
        if kw := self._config.get("keepAliveIntervalMs"):
            env["WHATSAPP_KEEPALIVE_INTERVAL_MS"] = str(kw)
        if sfh := self._config.get("syncFullHistory"):
            env["WHATSAPP_SYNC_FULL_HISTORY"] = "true" if sfh else "false"
        if mooc := self._config.get("markOnlineOnConnect"):
            env["WHATSAPP_MARK_ONLINE_ON_CONNECT"] = "true" if mooc else "false"

        self._process = await asyncio.create_subprocess_exec(
            self._node_path,
            os.path.join(self._bridge_dir, "index.js"),
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        self._http_session = aiohttp.ClientSession()
        await self._wait_ready(timeout=30.0)
        logger.info(f"Baileys bridge started on port {self._port}")

    async def stop(self) -> None:
        if self._http_session:
            await self._http_session.close()
            self._http_session = None
        if self._process and self._process.returncode is None:
            self._process.send_signal(signal.SIGTERM)
            try:
                await asyncio.wait_for(self._process.wait(), timeout=10.0)
            except TimeoutError:
                self._process.kill()
                await self._process.wait()
        if self._process:
            if self._process.stdout:
                self._process.stdout.close()
            if self._process.stderr:
                self._process.stderr.close()
            self._process = None
        logger.info("Baileys bridge stopped")

    async def health_check(self) -> HealthStatus:
        if self._http_session is None:
            return HealthStatus(status="unhealthy", last_error="Bridge not started")
        try:
            async with self._http_session.get(f"{self.base_url}/health") as resp:
                data = await resp.json()
                return HealthStatus(
                    status="healthy" if data.get("connected") else "degraded",
                    metadata={
                        "jid": data.get("jid", ""),
                        "connected": data.get("connected", False),
                    },
                )
        except Exception as e:
            return HealthStatus(status="unhealthy", last_error=str(e))

    async def _wait_ready(self, timeout: float) -> None:
        deadline = asyncio.get_event_loop().time() + timeout
        last_error = None
        while asyncio.get_event_loop().time() < deadline:
            try:
                async with self._http_session.get(f"{self.base_url}/health") as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        if data.get("status") == "running":
                            return
            except aiohttp.ClientConnectorError as e:
                last_error = str(e)
            await asyncio.sleep(0.5)
        detail = f"Last error: {last_error}" if last_error else "No response from bridge"
        raise TimeoutError(
            f"Baileys bridge not ready after {timeout}s. {detail}. "
            f"Check: node_path='{self._node_path}', bridge_dir='{self._bridge_dir}'"
        )

    async def send_message(
        self,
        jid: str,
        content: str,
        reply_to: str | None = None,
        silent: bool = False,
    ) -> DeliveryResult:
        payload: dict[str, Any] = {"jid": jid, "content": content}
        if reply_to:
            payload["reply_to"] = reply_to
        if silent:
            payload["silent"] = True
        ok, data = await self._request_with_retry("post", "/api/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_media(
        self,
        jid: str,
        media_type: str,
        media_path: str,
        caption: str = "",
        reply_to: str | None = None,
    ) -> DeliveryResult:
        payload: dict[str, Any] = {
            "jid": jid,
            "media_type": media_type,
            "media_path": media_path,
            "caption": caption,
        }
        if reply_to:
            payload["reply_to"] = reply_to
        ok, data = await self._request_with_retry("post", "/api/send-media", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_reaction(
        self,
        jid: str,
        message_id: str,
        emoji: str,
    ) -> DeliveryResult:
        payload = {"jid": jid, "message_id": message_id, "emoji": emoji}
        ok, data = await self._request_with_retry("post", "/api/react", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_presence(self, jid: str, presence: str) -> DeliveryResult:
        payload = {"jid": jid, "presence": presence}
        ok, data = await self._request_with_retry("post", "/api/presence", payload)
        return DeliveryResult(
            success=ok,
            error=data.get("error"),
        )

    async def send_read_receipt(self, jid: str, message_ids: list[str]) -> DeliveryResult:
        payload = {"jid": jid, "message_ids": message_ids}
        ok, data = await self._request_with_retry("post", "/api/read-receipt", payload)
        return DeliveryResult(
            success=ok,
            error=data.get("error"),
        )

    async def delete_message(self, jid: str, message_id: str) -> DeliveryResult:
        payload = {"jid": jid, "message_id": message_id}
        ok, data = await self._request_with_retry("post", "/api/message/delete", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def create_poll(
        self,
        jid: str,
        name: str,
        options: list[str],
        selectable_count: int = 1,
    ) -> DeliveryResult:
        payload = {
            "jid": jid,
            "name": name,
            "options": options,
            "selectable_count": selectable_count,
        }
        ok, data = await self._request_with_retry("post", "/api/poll/create", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_location(
        self,
        jid: str,
        latitude: float,
        longitude: float,
        name: str = "",
        address: str = "",
    ) -> DeliveryResult:
        payload = {
            "jid": jid,
            "latitude": latitude,
            "longitude": longitude,
        }
        if name:
            payload["name"] = name
        if address:
            payload["address"] = address
        ok, data = await self._request_with_retry("post", "/api/location/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_contact(
        self,
        jid: str,
        contacts: list[dict[str, str]],
    ) -> DeliveryResult:
        payload = {"jid": jid, "contacts": contacts}
        ok, data = await self._request_with_retry("post", "/api/contact/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_sticker(
        self,
        jid: str,
        sticker_path: str,
        reply_to: str | None = None,
    ) -> DeliveryResult:
        payload: dict[str, Any] = {"jid": jid, "sticker_path": sticker_path}
        if reply_to:
            payload["reply_to"] = reply_to
        ok, data = await self._request_with_retry("post", "/api/sticker/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_buttons(
        self,
        jid: str,
        text: str,
        buttons: list[dict[str, str]],
        title: str = "",
        footer: str = "",
    ) -> DeliveryResult:
        payload: dict[str, Any] = {"jid": jid, "text": text, "buttons": buttons}
        if title:
            payload["title"] = title
        if footer:
            payload["footer"] = footer
        ok, data = await self._request_with_retry("post", "/api/buttons/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def send_list_message(
        self,
        jid: str,
        text: str,
        sections: list[dict[str, Any]],
        title: str = "",
        footer: str = "",
        button_text: str = "Select",
    ) -> DeliveryResult:
        payload: dict[str, Any] = {
            "jid": jid,
            "text": text,
            "sections": sections,
        }
        if title:
            payload["title"] = title
        if footer:
            payload["footer"] = footer
        if button_text:
            payload["button_text"] = button_text
        ok, data = await self._request_with_retry("post", "/api/list/send", payload)
        return DeliveryResult(
            success=ok,
            message_id=data.get("message_id"),
            error=data.get("error"),
        )

    async def get_qr(self) -> dict[str, Any]:
        async with self._http_session.post(
            f"{self.base_url}/api/qr/login",
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def wait_scan(self, timeout: float = 120.0) -> dict[str, Any]:
        async with self._http_session.post(
            f"{self.base_url}/api/qr/wait",
            json={"timeout": timeout},
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def get_qr_status(self) -> dict[str, Any]:
        async with self._http_session.get(
            f"{self.base_url}/api/qr/status",
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def logout(self) -> dict[str, Any]:
        async with self._http_session.post(
            f"{self.base_url}/api/logout",
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def get_groups(self) -> dict[str, Any]:
        async with self._http_session.post(
            f"{self.base_url}/api/group/list",
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def get_group_info(self, group_jid: str) -> dict[str, Any]:
        async with self._http_session.post(
            f"{self.base_url}/api/group/info",
            json={"jid": group_jid},
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def get_message_history(self, jid: str, limit: int = 50, before: str | None = None) -> dict[str, Any]:
        payload: dict[str, Any] = {"jid": jid, "limit": limit}
        if before:
            payload["before"] = before
        async with self._http_session.post(
            f"{self.base_url}/api/message/history",
            json=payload,
            headers=self._auth_headers,
        ) as resp:
            return await resp.json()

    async def download_media(self, remote_jid: str, message_id: str, message: dict[str, Any]) -> bytes:
        payload = {
            "remoteJid": remote_jid,
            "id": message_id,
            "message": message,
        }
        async with self._http_session.post(
            f"{self.base_url}/api/media/download",
            json=payload,
            headers=self._auth_headers,
        ) as resp:
            if resp.status == 200:
                return await resp.read()
            data = await resp.json()
            raise RuntimeError(data.get("error", "Media download failed"))

    async def get_profile_picture(self, jid: str) -> str | None:
        try:
            async with self._http_session.post(
                f"{self.base_url}/api/user/picture",
                json={"jid": jid},
                headers=self._auth_headers,
            ) as resp:
                data = await resp.json()
                return data.get("url")
        except Exception:
            return None