ForcePilot/backend/package/yuxi/channel/extensions/zalouser/gateway.py

from __future__ import annotations

import asyncio
import logging
import time

from yuxi.channel.extensions.zalouser.errors import (
    ZaloUserAuthError,
    ZaloUserConnectionError,
)
from yuxi.channel.extensions.zalouser.monitor import ZaloUserMonitor
from yuxi.channel.extensions.zalouser.sidecar_client import ZcaSidecarClient
from yuxi.channel.extensions.zalouser.types import SidecarGatewayContext

logger = logging.getLogger(__name__)

WATCHDOG_INTERVAL_MS = 30_000
WATCHDOG_MAX_GAP_MS = 35_000


class ZaloUserGatewayAdapter:
    def __init__(self):
        self._monitor = ZaloUserMonitor()
        self._client: ZcaSidecarClient | None = None
        self._ctx: SidecarGatewayContext | None = None
        self._receive_task: asyncio.Task | None = None
        self._ping_task: asyncio.Task | None = None
        self._watchdog_task: asyncio.Task | None = None
        self._running = False
        self._last_activity: float = 0
        self._queue: asyncio.Queue | None = None

    @property
    def client(self) -> ZcaSidecarClient | None:
        return self._client

    async def start(self, ctx) -> object:
        account = getattr(ctx, "account", {}) or {}
        sidecar_url = account.get("sidecar_url", "")
        account_id = account.get("account_id", "default")
        profile = account.get("profile", "default")

        self._ctx = SidecarGatewayContext(
            sidecar_url=sidecar_url,
            account_id=account_id,
            profile=profile,
            ping_interval=account.get("ping_interval", 30),
            ping_timeout=account.get("ping_timeout", 10),
        )

        if not sidecar_url:
            raise ZaloUserAuthError("Sidecar URL is required")

        self._client = ZcaSidecarClient(sidecar_url, profile)
        self._running = True
        self._queue = getattr(ctx, "queue", asyncio.Queue())

        authenticated = await self._client.check_auth()
        if not authenticated:
            logger.warning("ZaloUser sidecar not authenticated for profile '%s'", profile)
            return self._queue

        self._last_activity = time.time()

        self._receive_task = asyncio.create_task(self._receive_loop())
        self._ping_task = asyncio.create_task(self._ping_loop())
        self._watchdog_task = asyncio.create_task(self._watchdog_loop())

        return self._queue

    async def stop(self, ctx) -> None:
        self._running = False
        for task in [self._receive_task, self._ping_task, self._watchdog_task]:
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        if self._client:
            await self._client.close()
            self._client = None

    async def login_with_qr_start(self, account_id: str | None = None, *, force: bool = False, timeout_ms: int | None = None) -> dict:
        if not self._client:
            return {"qr_data_url": "", "message": "Sidecar client not initialized", "connected": False}
        try:
            result = await self._client.login_qr_start(force=force, timeout_ms=timeout_ms or 30_000)
            return {
                "qr_data_url": result.get("qrDataUrl", ""),
                "message": result.get("message", "QR code generated"),
                "connected": result.get("connected", False),
            }
        except ZaloUserConnectionError as e:
            return {"qr_data_url": "", "message": str(e), "connected": False}

    async def login_with_qr_wait(self, account_id: str | None = None, *, timeout_ms: int | None = None, current_qr_data_url: str | None = None) -> dict:
        if not self._client:
            return {"connected": False, "message": "Sidecar client not initialized", "qr_data_url": None}
        try:
            result = await self._client.login_qr_wait(timeout_ms=timeout_ms or 120_000)
            return {
                "connected": True,
                "message": "Login successful",
                "qr_data_url": None,
                "user_id": result.get("user_id", ""),
                "display_name": result.get("display_name", ""),
            }
        except ZaloUserAuthError as e:
            return {"connected": False, "message": str(e), "qr_data_url": None}

    async def logout_account(self, ctx: object) -> dict:
        if not self._client:
            return {"cleared": False, "logged_out": None}
        try:
            await self._client.logout()
            return {"cleared": True, "logged_out": True}
        except Exception as e:
            logger.warning("Logout failed: %s", e)
            return {"cleared": False, "logged_out": None}

    async def probe(self, account: dict) -> bool:
        sidecar_url = account.get("sidecar_url", "")
        if not sidecar_url:
            return False
        client = ZcaSidecarClient(sidecar_url, account.get("profile", "default"))
        try:
            return await client.check_auth()
        except Exception:
            return False
        finally:
            await client.close()

    async def _receive_loop(self) -> None:
        while self._running:
            if self._client and self._ctx:
                try:
                    events = await self._client.poll_events(timeout=30.0)
                    if events:
                        self._last_activity = time.time()
                        for event in events:
                            await self._process_event(event)
                except Exception as e:
                    logger.warning("Receive loop error: %s", e)
                    await asyncio.sleep(1)
            else:
                await asyncio.sleep(1)

    async def _process_event(self, event: dict) -> None:
        if not self._ctx or not self._queue:
            return
        unified_msg = self._monitor.parse_event(event, self._ctx.account_id)
        if unified_msg:
            await self._queue.put(unified_msg)

    async def _ping_loop(self) -> None:
        while self._running and self._ctx and self._client:
            interval = self._ctx.ping_interval
            await asyncio.sleep(interval)
            if not self._running:
                break
            try:
                status = await self._client.get_status()
                if status.get("authenticated") or status.get("connected"):
                    self._last_activity = time.time()
            except Exception:
                pass

    async def _watchdog_loop(self) -> None:
        while self._running:
            await asyncio.sleep(WATCHDOG_INTERVAL_MS / 1000.0)
            if not self._running:
                break
            gap = (time.time() - self._last_activity) * 1000
            if gap > WATCHDOG_MAX_GAP_MS:
                logger.warning(
                    "ZaloUser watchdog timeout: %.0fms since last activity",
                    gap,
                )
                self._running = False
                self._last_activity = time.time()
                try:
                    self._running = True
                    self._receive_task = asyncio.create_task(self._receive_loop())
                except Exception as e:
                    logger.error("Watchdog reconnect failed: %s", e)