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

from __future__ import annotations

import asyncio
import json
import time
from pathlib import Path
from typing import Any

from yuxi.utils.logging_config import logger


class CredentialQueue:
    def __init__(self, auth_dir: Path, debounce_ms: int = 500):
        self._auth_dir = auth_dir
        self._lock = asyncio.Lock()
        self._write_count = 0
        self._debounce_ms = debounce_ms
        self._pending_creds: dict[str, Any] | None = None
        self._last_write_time = 0.0
        self._debounce_task: asyncio.Task | None = None

    async def write_auth(self, creds: dict[str, Any]) -> int:
        loop = asyncio.get_running_loop()
        now = time.monotonic()
        async with self._lock:
            self._write_count += 1
            count = self._write_count

            if now - self._last_write_time < self._debounce_ms / 1000.0:
                self._pending_creds = creds
                self._schedule_debounced_write(loop)
                return count

            self._last_write_time = now
            self._pending_creds = None

        return await self._do_write(creds, count)

    def _schedule_debounced_write(self, loop: asyncio.AbstractEventLoop) -> None:
        if self._debounce_task is not None:
            self._debounce_task.cancel()

        async def _delayed_flush():
            await asyncio.sleep(self._debounce_ms / 1000.0)
            await self._flush_debounced()

        self._debounce_task = loop.create_task(_delayed_flush())

    async def _flush_debounced(self) -> None:
        async with self._lock:
            creds = self._pending_creds
            self._pending_creds = None
            self._debounce_task = None
            if creds is None:
                return
            count = self._write_count

        await self._do_write(creds, count)

    async def _do_write(self, creds: dict[str, Any], count: int) -> int:
        path = self._auth_dir / "creds.json"
        backup = self._auth_dir / "creds.json.bak"
        tmp_path = self._auth_dir / "creds.json.tmp"
        try:
            data = json.dumps(creds, indent=2, ensure_ascii=False)

            def _sync_write() -> int:
                tmp_path.write_text(data, encoding="utf-8")
                if path.exists():
                    try:
                        path.replace(backup)
                    except OSError:
                        pass
                tmp_path.replace(path)
                return count

            result = await asyncio.to_thread(_sync_write)
            logger.debug(f"CredentialQueue: wrote creds (#{result}) to {path}")
            return result
        except (OSError, TypeError) as e:
            logger.error(f"CredentialQueue: write failed (#{count}): {e}")
            raise

    async def read_auth(self) -> dict[str, Any] | None:
        path = self._auth_dir / "creds.json"
        backup = self._auth_dir / "creds.json.bak"

        def _sync_read():
            if not path.exists():
                return None
            try:
                return json.loads(path.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError) as e:
                logger.error(f"CredentialQueue: read failed: {e}")
                if backup.exists():
                    try:
                        data = json.loads(backup.read_text(encoding="utf-8"))
                        backup.replace(path)
                        logger.info("CredentialQueue: restored creds from backup")
                        return data
                    except (OSError, json.JSONDecodeError) as e2:
                        logger.error(f"CredentialQueue: backup recovery failed: {e2}")
                return None

        return await asyncio.to_thread(_sync_read)

    async def clear_auth(self) -> bool:
        async with self._lock:
            self._pending_creds = None
            if self._debounce_task is not None:
                self._debounce_task.cancel()
                self._debounce_task = None

        path = self._auth_dir / "creds.json"
        backup = self._auth_dir / "creds.json.bak"

        def _sync_clear():
            if path.exists():
                if backup.exists():
                    backup.unlink()
                path.rename(backup)
                logger.info(f"CredentialQueue: cleared creds, backup at {backup}")
            return True

        try:
            result = await asyncio.to_thread(_sync_clear)
            return result
        except OSError as e:
            logger.error(f"CredentialQueue: clear failed: {e}")
            return False

    @property
    def write_count(self) -> int:
        return self._write_count