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)