huihuiSquare/digital-avatar-app/backend/services/takeover_service.py

"""Takeover service — message listening, decision, reply execution."""
import json
import logging
import os
from typing import Optional

import httpx
from sqlalchemy.orm import Session

from models import Avatar, Authorization
from services.boxim_client import BoxIMClient

logger = logging.getLogger(__name__)


class TakeoverService:
    """Service for handling avatar takeover — generating replies and sending them via IM."""

    def __init__(
        self,
        db: Session,
        boxim_client: BoxIMClient,
        redis_client=None,
    ):
        self.db = db
        self.boxim = boxim_client
        self.redis = redis_client
        self._chat_api_base = os.getenv(
            "TAKEOVER_CHAT_API_BASE", "http://localhost:8000/api"
        )

    def check_takeover_enabled(
        self, owner_huihui_id: str, from_user_id: str
    ) -> Optional[Authorization]:
        """Check whether takeover is enabled for the given target user."""
        avatar = self.db.query(Avatar).filter(Avatar.owner_id == owner_huihui_id).first()
        if not avatar:
            return None

        auth = (
            self.db.query(Authorization)
            .filter(Authorization.avatar_id == avatar.id)
            .filter(Authorization.target_id == from_user_id)
            .filter(Authorization.takeover_enabled == True)
            .first()
        )
        return auth if auth and auth.takeover_enabled else None

    async def generate_reply(self, avatar_id: str, message: str) -> str:
        """Call the avatar chat endpoint to generate a reply."""
        try:
            async with httpx.AsyncClient(timeout=30) as client:
                r = await client.post(
                    f"{self._chat_api_base}/avatar/{avatar_id}/chat",
                    json={"message": message, "history": []},
                )
            data = r.json()
            if data.get("code") in (0, 200):
                return data.get("data", {}).get("answer", "")
            logger.warning(f"Avatar chat API returned error code: {data}")
            return ""
        except Exception as e:
            logger.error(f"Failed to call avatar chat API: {e}")
            return ""

    async def execute_takeover(self, auth: Authorization, message: dict) -> bool:
        """Execute takeover: generate a reply and send it as the owner via IM."""
        try:
            # Resolve owner through Avatar model
            avatar = self.db.query(Avatar).filter(Avatar.id == auth.avatar_id).first()
            if not avatar:
                logger.warning(f"Avatar not found: {auth.avatar_id}")
                return False

            owner_huihui_id = avatar.owner_id
            credentials = await self.boxim.get_credentials(owner_huihui_id)
            if not credentials:
                logger.warning(f"Cannot obtain IM credentials for owner: {owner_huihui_id}")
                return False

            reply = await self.generate_reply(auth.avatar_id, message.get("content", ""))
            if not reply:
                logger.warning("Avatar did not generate a reply")
                return False

            success = await self.boxim.send_p2p_message(
                from_accid=credentials["accid"],
                to_accid=message.get("from_accid", ""),
                content=reply,
            )
            if success:
                logger.info(f"Takeover reply sent successfully: {reply[:50]}...")
            return success
        except Exception as e:
            logger.error(f"Takeover execution failed: {e}")
            return False

    def enqueue_delayed_message(self, auth: Authorization, message: dict):
        """Write a message into the Redis delayed queue (TTL = delay + 10s buffer)."""
        if not self.redis:
            logger.warning("Redis not configured, degrading to immediate takeover")
            return

        avatar = self.db.query(Avatar).filter(Avatar.id == auth.avatar_id).first()
        owner_huihui_id = avatar.owner_id if avatar else ""
        key = f"takeover:delayed:{auth.target_id}:{message.get('msg_id', '')}"
        value = json.dumps({
            "avatar_id": auth.avatar_id,
            "from_accid": message.get("from_accid", ""),
            "content": message.get("content", ""),
            "owner_huihui_id": owner_huihui_id,
        })
        self.redis.setex(key, auth.takeover_delay_seconds + 10, value)
        logger.info(f"Message enqueued to delayed queue: {key}")

    async def process_delayed_queue(self):
        """Process expired messages from the delayed queue.

        Scans Redis keys matching the takeover:delayed: pattern and dispatches
        each to execute_takeover after resolving the Authorization.
        """
        if not self.redis:
            return
        try:
            pattern = "takeover:delayed:*"
            keys = self.redis.keys(pattern)
            for key in keys:
                raw = self.redis.get(key)
                if not raw:
                    continue
                data = json.loads(raw)
                auth = (
                    self.db.query(Authorization)
                    .filter(Authorization.target_id == key.split(":")[2])
                    .first()
                )
                if auth:
                    message = {
                        "msg_id": key.split(":")[-1],
                        "from_accid": data.get("from_accid", ""),
                        "content": data.get("content", ""),
                    }
                    await self.execute_takeover(auth, message)
                self.redis.delete(key)
        except Exception as e:
            logger.error(f"Failed to process delayed queue: {e}")

    async def poll_and_process_messages(self):
        """Periodic polling job: fetch unread messages and process each."""
        try:
            messages = await self.fetch_unread_messages()
            for msg in messages:
                await self.process_message(msg)
        except Exception as e:
            logger.error(f"poll_and_process_messages failed: {e}")

    async def fetch_unread_messages(self) -> list:
        """Fetch unread messages from Box IM. Stub — replace with real API call."""
        logger.debug("fetch_unread_messages: no real API wired yet")
        return []

    async def process_message(self, message: dict):
        """Process a single message: check takeover, dispatch immediate or delayed."""
        owner_id = message.get("owner_huihui_id", "")
        from_id = message.get("from_accid", "")

        auth = self.check_takeover_enabled(owner_id, from_id)
        if not auth:
            return

        if auth.takeover_mode == "immediate":
            await self.execute_takeover(auth, message)
        else:
            self.enqueue_delayed_message(auth, message)