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

from __future__ import annotations

import asyncio
import json
import os
import random
from typing import Any
from collections.abc import Callable

from yuxi.channels.adapters.zalo_oa.cache import SentMessageCache
from yuxi.channels.adapters.zalo_oa.client import ZaloOAClient
from yuxi.channels.adapters.zalo_oa.formatter import ZaloOAMessageFormatter
from yuxi.channels.exceptions import ChannelRateLimitError, TokenExpiredError
from yuxi.channels.models import ChannelResponse, DeliveryResult
from yuxi.utils.logging_config import logger

MAX_TOKEN_REFRESH_RETRIES = 2
BROADCAST_STATE_FILE = "zalo_oa_broadcast_state.json"


class BroadcastState:
    def __init__(self, state_file: str = BROADCAST_STATE_FILE):
        self._state_file = state_file
        self._state: dict[str, Any] = {}

    def load(self) -> dict[str, Any] | None:
        try:
            if os.path.exists(self._state_file):
                with open(self._state_file) as f:
                    self._state = json.load(f)
                return self._state
        except (OSError, json.JSONDecodeError):
            pass
        return None

    def save(self, broadcast_id: str, last_index: int, total: int, text: str):
        self._state = {
            "broadcast_id": broadcast_id,
            "last_index": last_index,
            "total": total,
            "text": text,
        }
        try:
            with open(self._state_file, "w") as f:
                json.dump(self._state, f)
        except OSError:
            pass

    def clear(self):
        self._state = {}
        try:
            if os.path.exists(self._state_file):
                os.remove(self._state_file)
        except OSError:
            pass

    @property
    def has_state(self) -> bool:
        return bool(self._state.get("broadcast_id"))


class ZaloOASender:
    def __init__(self, client: ZaloOAClient, formatter: ZaloOAMessageFormatter, config: dict[str, Any] | None = None):
        self._client = client
        self._formatter = formatter
        self._config = config or {}
        self._response_prefix = config.get("response_prefix", config.get("responsePrefix", "")) if config else ""
        self._message_cache = SentMessageCache()
        self._send_progress_callback: Callable[[int, int, str], None] | None = None
        self._broadcast_state = BroadcastState()

    async def send(self, response: ChannelResponse) -> DeliveryResult:
        try:
            template = self._formatter.format(response)
            return await self._do_send(template)
        except Exception as e:
            logger.error(f"[ZaloOA] Send failed: {e}")
            return DeliveryResult(success=False, error=str(e))

    async def _do_send(self, template: dict[str, Any]) -> DeliveryResult:
        retry_config = self._config.get("retry", {})
        max_attempts = retry_config.get("attempts", 3)
        min_delay = retry_config.get("min_delay_ms", 400) / 1000
        max_delay = retry_config.get("max_delay_ms", 30000) / 1000

        last_error = None
        token_refresh_attempts = 0

        for attempt in range(max_attempts):
            try:
                data = await self._client.send_message(template)
                msg_id = data.get("message_id", "")
                recipient = template.get("recipient", {}).get("user_id", "")
                content = template.get("message", {}).get("text", "")
                if msg_id:
                    self._message_cache.put(msg_id, recipient, content)
                return DeliveryResult(
                    success=True,
                    message_id=msg_id,
                )
            except ChannelRateLimitError:
                retry_after = 5
                logger.warning(f"[ZaloOA] Rate limited, waiting {retry_after}s")
                await asyncio.sleep(retry_after)
                continue
            except TokenExpiredError:
                if token_refresh_attempts < MAX_TOKEN_REFRESH_RETRIES:
                    token_refresh_attempts += 1
                    await asyncio.sleep(1)
                    continue
                last_error = "Token expired, max refresh retries exceeded"
            except Exception as e:
                last_error = str(e)
                if attempt < max_attempts - 1:
                    delay = random.uniform(0, min(max_delay, min_delay * (2**attempt)))
                    logger.warning(f"[ZaloOA] Retry {attempt + 1}/{max_attempts} after {delay:.1f}s: {last_error}")
                    await asyncio.sleep(delay)
                else:
                    return DeliveryResult(success=False, error=last_error)

        return DeliveryResult(
            success=False,
            error=f"Max retries exceeded: {last_error}",
        )

    async def send_with_fallback(
        self,
        recipient_id: str,
        content: str,
        attachments: list[dict[str, Any]] | None = None,
    ) -> DeliveryResult:
        if attachments:
            try:
                attachment_id = attachments[0].get("attachment_id")
                if attachment_id:
                    media_type = attachments[0].get("type", "image")
                    if media_type == "image":
                        await self.send_upload_photo_indicator(recipient_id)
                    template = self._formatter._build_media(recipient_id, media_type, attachment_id)
                    result = await self._do_send(template)
                    if result.success:
                        return result
            except Exception as e:
                logger.warning(f"[ZaloOA] Rich message failed, falling back to text: {e}")

        text = content
        if attachments:
            urls = [att.get("url", "") for att in attachments if att.get("url")]
            if urls:
                text += "\n\n" + "\n".join(urls)

        template = self._formatter._build_text(recipient_id, text)
        return await self._do_send(template)

    async def send_typing_indicator(self, recipient_id: str) -> bool:
        return await self._client.send_chat_action(recipient_id, "typing")

    async def send_upload_photo_indicator(self, recipient_id: str) -> bool:
        return await self._client.send_chat_action(recipient_id, "upload_photo")

    async def send_payload_with_chunked_text_and_media(
        self,
        recipient_id: str,
        text: str,
        media_attachments: list[dict[str, Any]] | None = None,
    ) -> DeliveryResult:
        from yuxi.channels.adapters.zalo_oa.chunking import chunk_text

        text_limit = self._formatter.max_text_length
        results: list[DeliveryResult] = []

        if media_attachments:
            for att in media_attachments:
                if att.get("attachment_id"):
                    media_type = att.get("type", "image")
                    if media_type == "image":
                        await self.send_upload_photo_indicator(recipient_id)
                    template = self._formatter._build_media(recipient_id, media_type, att["attachment_id"])
                    result = await self._do_send(template)
                    results.append(result)

        if text:
            if self._response_prefix and not text.startswith(self._response_prefix):
                text = self._response_prefix + text

            chunks = chunk_text(text, text_limit)
            for chunk in chunks:
                template = self._formatter._build_text(recipient_id, chunk)
                result = await self._do_send(template)
                results.append(result)

        if not results:
            return DeliveryResult(success=False, error="No content to send")

        success = all(r.success for r in results)
        last_msg_id = results[-1].message_id if results else None
        return DeliveryResult(success=success, message_id=last_msg_id)

    def set_progress_callback(self, callback: Callable[[int, int, str], None]):
        self._send_progress_callback = callback

    async def broadcast(
        self,
        text: str,
        media_attachments: list[dict[str, Any]] | None = None,
        disable_notification: bool = False,
        start_index: int = 0,
    ) -> DeliveryResult:
        followers = []
        offset = 0
        while True:
            result = await self._client.get_followers(offset=offset, count=50)
            batch = result.get("followers", [])
            followers.extend(batch)
            if len(batch) < 50 or len(followers) >= result.get("total", 0):
                break
            offset += len(batch)

        if not followers:
            return DeliveryResult(success=False, error="No followers to broadcast to")

        total = len(followers)
        success_count = 0
        fail_count = 0
        last_msg_id = None

        for idx, follower in enumerate(followers):
            if idx < start_index:
                success_count += 1
                continue
            user_id = str(follower.get("user_id", ""))
            if not user_id:
                continue

            try:
                result = await self.send_payload_with_chunked_text_and_media(user_id, text, media_attachments)
                if result.success:
                    success_count += 1
                    last_msg_id = result.message_id
                else:
                    fail_count += 1
            except Exception as e:
                fail_count += 1
                logger.warning(f"[ZaloOA] Broadcast to {user_id} failed: {e}")

            if self._send_progress_callback:
                self._send_progress_callback(idx + 1, total, user_id)

            if (idx + 1) % 10 == 0:
                self._broadcast_state.save("zalo_oa_broadcast", idx + 1, total, text)

            if fail_count > total * 0.5:
                return DeliveryResult(
                    success=False,
                    error=f"Broadcast aborted: {fail_count}/{total} failures",
                )

        return DeliveryResult(
            success=fail_count == 0,
            message_id=last_msg_id,
            metadata={"sent": success_count, "failed": fail_count, "total": total},
        )

    async def broadcast_resume(self, text: str | None = None) -> DeliveryResult:
        state = self._broadcast_state.load()
        if not state:
            return DeliveryResult(success=False, error="No broadcast state to resume")
        last_index = state.get("last_index", 0)
        total = state.get("total", 0)
        broadcast_text = state.get("text", text or "")
        logger.info(f"[ZaloOA] Resuming broadcast from index {last_index}/{total}")
        return await self.broadcast(broadcast_text, start_index=last_index)

    async def silent_send(self, recipient_id: str, text: str) -> DeliveryResult:
        template = self._formatter._build_text(recipient_id, text)
        return await self._do_send(template)

    async def send_audio(self, recipient_id: str, audio_data: bytes, filename: str = "audio.mp3") -> DeliveryResult:
        try:
            attachment_id = await self._client.upload_file(audio_data, filename)
        except Exception as e:
            return DeliveryResult(success=False, error=f"Audio upload failed: {e}")

        template = self._formatter._build_media(recipient_id, "file", attachment_id)
        return await self._do_send(template)