PhoneWork/bot/handler.py

"""Feishu event handler using lark-oapi long-connection (WebSocket) mode."""

from __future__ import annotations

import asyncio
import json
import logging
import threading
import time

import lark_oapi as lark
from lark_oapi.api.im.v1 import P2ImMessageReceiveV1

from bot.commands import handle_command
from bot.feishu import send_text
from config import FEISHU_APP_ID, FEISHU_APP_SECRET, is_user_allowed
from orchestrator.agent import agent
from orchestrator.tools import set_current_chat

logger = logging.getLogger(__name__)

_main_loop: asyncio.AbstractEventLoop | None = None
_ws_connected: bool = False
_last_message_time: float = 0.0
_reconnect_count: int = 0


def get_ws_status() -> dict:
    """Return WebSocket connection status."""
    return {
        "connected": _ws_connected,
        "last_message_time": _last_message_time,
        "reconnect_count": _reconnect_count,
    }


def _handle_message(data: P2ImMessageReceiveV1) -> None:
    global _last_message_time
    _last_message_time = time.time()

    try:
        message = data.event.message
        sender = data.event.sender

        logger.debug(
            "event type=%r chat_type=%r content=%r",
            getattr(message, "message_type", None),
            getattr(message, "chat_type", None),
            (getattr(message, "content", None) or "")[:100],
        )

        if message.message_type != "text":
            logger.info("Skipping non-text message_type=%r", message.message_type)
            return

        chat_id: str = message.chat_id
        raw_content: str = message.content or "{}"
        content_obj = json.loads(raw_content)
        text: str = content_obj.get("text", "").strip()

        import re
        text = re.sub(r"@\S+\s*", "", text).strip()

        open_id: str = ""
        if sender and sender.sender_id:
            open_id = sender.sender_id.open_id or ""

        if not text:
            logger.info("Empty text after stripping, ignoring")
            return

        logger.info("✉  ...%s → %r", open_id[-8:], text[:80])

        user_id = open_id or chat_id

        if _main_loop is None:
            logger.error("Main event loop not set; cannot process message")
            return

        asyncio.run_coroutine_threadsafe(
            _process_message(user_id, chat_id, text),
            _main_loop,
        )
    except Exception:
        logger.exception("Error in _handle_message")


async def _process_message(user_id: str, chat_id: str, text: str) -> None:
    """Process message: check allowlist, then commands, then agent."""
    try:
        set_current_chat(chat_id)

        if not is_user_allowed(user_id):
            logger.warning("Rejected message from unauthorized user: ...%s", user_id[-8:])
            await send_text(chat_id, "chat_id", "Sorry, you are not authorized to use this bot.")
            return

        reply = await handle_command(user_id, text)
        if reply is None:
            reply = await agent.run(user_id, text)
        if reply:
            await send_text(chat_id, "chat_id", reply)
    except Exception:
        logger.exception("Error processing message for user %s", user_id)


def _handle_any(data: lark.CustomizedEvent) -> None:
    """Catch-all handler to log any event the SDK receives."""
    logger.info("RAW CustomizedEvent: %s", lark.JSON.marshal(data)[:500])


def build_event_handler() -> lark.EventDispatcherHandler:
    """Construct the EventDispatcherHandler with all registered callbacks."""
    handler = (
        lark.EventDispatcherHandler.builder("", "")
        .register_p2_im_message_receive_v1(_handle_message)
        .register_p1_customized_event("im.message.receive_v1", _handle_any)
        .build()
    )
    return handler


def start_websocket_client(loop: asyncio.AbstractEventLoop) -> None:
    """
    Start the lark-oapi WebSocket long-connection client in a background thread.
    Must be called after the asyncio event loop is running.
    """
    global _main_loop
    _main_loop = loop

    def _run_with_reconnect() -> None:
        global _ws_connected, _reconnect_count
        backoff = 1.0
        max_backoff = 60.0

        while True:
            try:
                _ws_connected = False
                event_handler = build_event_handler()
                ws_client = lark.ws.Client(
                    FEISHU_APP_ID,
                    FEISHU_APP_SECRET,
                    event_handler=event_handler,
                    log_level=lark.LogLevel.INFO,
                )

                logger.info("Starting Feishu long-connection client...")
                _ws_connected = True
                _reconnect_count += 1
                ws_client.start()
                logger.warning("WebSocket disconnected, will reconnect...")

            except Exception as e:
                logger.error("WebSocket error: %s", e)

            finally:
                _ws_connected = False

            logger.info("Reconnecting in %.1fs (attempt %d)...", backoff, _reconnect_count)
            time.sleep(backoff)
            backoff = min(backoff * 2, max_backoff)

    thread = threading.Thread(target=_run_with_reconnect, daemon=True, name="feishu-ws")
    thread.start()
    logger.info("Feishu WebSocket thread started")