data-query-agent/src/data_query_agent/feishu.py

from __future__ import annotations

import asyncio
import json
import logging
import threading
import time
from pathlib import Path
from typing import Any, Callable

import httpx
import lark_oapi as lark

from .config import Settings
from .models import IncomingMessage

logger = logging.getLogger(__name__)


class FeishuApiError(RuntimeError):
    pass


class FeishuApi:
    def __init__(self, settings: Settings) -> None:
        self.app_id = settings.feishu_app_id
        self.app_secret = settings.feishu_app_secret
        self.base_url = settings.feishu_domain + "/open-apis"
        self._client = httpx.AsyncClient(timeout=30)
        self._token = ""
        self._token_expires_at = 0.0
        self._token_lock = asyncio.Lock()

    @staticmethod
    def _payload(response: httpx.Response, action: str) -> dict[str, Any]:
        try:
            response.raise_for_status()
            payload = response.json()
        except (httpx.HTTPError, ValueError) as exc:
            raise FeishuApiError(f"{action}失败：HTTP 响应无效") from exc
        if int(payload.get("code") or 0) != 0:
            raise FeishuApiError(f"{action}失败：code={payload.get('code')}, msg={payload.get('msg')}")
        return payload

    async def token(self) -> str:
        if self._token and time.monotonic() < self._token_expires_at:
            return self._token
        async with self._token_lock:
            if self._token and time.monotonic() < self._token_expires_at:
                return self._token
            response = await self._client.post(
                f"{self.base_url}/auth/v3/tenant_access_token/internal",
                json={"app_id": self.app_id, "app_secret": self.app_secret},
            )
            payload = self._payload(response, "获取 tenant access token")
            self._token = str(payload["tenant_access_token"])
            self._token_expires_at = time.monotonic() + max(60, int(payload.get("expire") or 7200) - 300)
            return self._token

    async def _headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {await self.token()}"}

    async def bot_open_id(self) -> str:
        response = await self._client.get(f"{self.base_url}/bot/v3/info", headers=await self._headers())
        payload = self._payload(response, "获取机器人信息")
        open_id = str((payload.get("bot") or {}).get("open_id") or "")
        if not open_id:
            raise FeishuApiError("机器人信息缺少 open_id")
        return open_id

    async def reply_text(self, message_id: str, text: str) -> None:
        response = await self._client.post(
            f"{self.base_url}/im/v1/messages/{message_id}/reply",
            headers={**await self._headers(), "Content-Type": "application/json"},
            json={"msg_type": "text", "content": json.dumps({"text": text}, ensure_ascii=False)},
        )
        self._payload(response, "回复飞书消息")

    async def reply_card(self, message_id: str, title: str, markdown: str, url: str | None = None) -> None:
        elements: list[dict[str, Any]] = [{"tag": "markdown", "content": markdown}]
        if url:
            elements.extend([
                {"tag": "hr"},
                {"tag": "action", "actions": [{
                    "tag": "button",
                    "type": "primary",
                    "text": {"tag": "plain_text", "content": "打开在线表格"},
                    "url": url,
                }]},
            ])
        card = {
            "config": {"wide_screen_mode": True},
            "header": {"template": "blue", "title": {"tag": "plain_text", "content": title[:80]}},
            "elements": elements,
        }
        response = await self._client.post(
            f"{self.base_url}/im/v1/messages/{message_id}/reply",
            headers={**await self._headers(), "Content-Type": "application/json"},
            json={"msg_type": "interactive", "content": json.dumps(card, ensure_ascii=False)},
        )
        self._payload(response, "回复飞书卡片")

    async def publish_sheet(self, file_path: Path, title: str) -> str:
        extension = file_path.suffix.lower().lstrip(".")
        if extension not in {"csv", "xls", "xlsx"} or not file_path.is_file():
            raise ValueError("待发布文件必须是存在的 CSV/XLS/XLSX")
        headers = await self._headers()
        with file_path.open("rb") as handle:
            response = await self._client.post(
                f"{self.base_url}/drive/v1/medias/upload_all",
                headers=headers,
                data={
                    "file_name": file_path.name,
                    "parent_type": "ccm_import_open",
                    "size": str(file_path.stat().st_size),
                    "extra": json.dumps({"obj_type": "sheet", "file_extension": extension}),
                },
                files={"file": (file_path.name, handle, "application/octet-stream")},
            )
        file_token = str(self._payload(response, "上传待导入文件")["data"]["file_token"])
        response = await self._client.post(
            f"{self.base_url}/drive/v1/import_tasks",
            headers={**headers, "Content-Type": "application/json"},
            json={
                "file_extension": extension,
                "file_token": file_token,
                "type": "sheet",
                "file_name": title[:100],
                "point": {"mount_type": 1, "mount_key": ""},
            },
        )
        ticket = str(self._payload(response, "创建表格导入任务")["data"]["ticket"])
        result: dict[str, Any] | None = None
        for _ in range(45):
            await asyncio.sleep(2)
            response = await self._client.get(
                f"{self.base_url}/drive/v1/import_tasks/{ticket}", headers=headers
            )
            current = self._payload(response, "查询表格导入结果").get("data", {}).get("result", {})
            if current.get("job_status") == 0:
                result = current
                break
            if current.get("job_status") == 3:
                raise FeishuApiError(f"导入在线表格失败：{current.get('job_error_msg', 'unknown error')}")
        if result is None:
            raise TimeoutError("导入在线表格超过 90 秒仍未完成")
        sheet_token = str(result.get("token") or "")
        url = str(result.get("url") or "")
        file_type = str(result.get("type") or "sheet")
        if not sheet_token or not url:
            raise FeishuApiError("导入响应缺少表格 token 或 URL")
        response = await self._client.patch(
            f"{self.base_url}/drive/v2/permissions/{sheet_token}/public",
            headers={**headers, "Content-Type": "application/json"},
            params={"type": file_type},
            json={"external_access": False, "link_share_entity": "tenant_readable"},
        )
        self._payload(response, "设置企业内链接可读权限")
        return url

    async def close(self) -> None:
        await self._client.aclose()


class FeishuLongConnection:
    def __init__(self, settings: Settings, bot_open_id: str, callback: Callable[[IncomingMessage], None]) -> None:
        self.settings = settings
        self.bot_open_id = bot_open_id
        self.callback = callback
        self._thread: threading.Thread | None = None
        self._client: Any = None

    @staticmethod
    def _parse_content(content: str, message_type: str) -> str:
        try:
            payload = json.loads(content)
        except (TypeError, json.JSONDecodeError):
            return str(content)
        if message_type == "text":
            return str(payload.get("text") or "")
        if message_type == "post":
            parts: list[str] = []
            body = payload.get("zh_cn") or payload
            if body.get("title"):
                parts.append(str(body["title"]))
            for paragraph in body.get("content", []) or []:
                for item in paragraph if isinstance(paragraph, list) else []:
                    if item.get("tag") in {"text", "a"}:
                        parts.append(str(item.get("text") or item.get("href") or ""))
            return "\n".join(parts).strip()
        return ""

    def _handle(self, data: Any) -> None:
        try:
            event = data.event
            message = event.message
            sender_id = event.sender.sender_id if event.sender else None
            mentions = list(message.mentions or [])
            mentioned_bot = any(
                getattr(getattr(mention, "id", None), "open_id", None) == self.bot_open_id
                for mention in mentions
            )
            text = self._parse_content(message.content, message.message_type)
            for mention in mentions:
                key = str(getattr(mention, "key", "") or "")
                name = str(getattr(mention, "name", "") or "")
                if key:
                    text = text.replace(key, "")
                if name:
                    text = text.replace(f"@{name}", "")
            incoming = IncomingMessage(
                message_id=str(message.message_id or ""),
                chat_id=str(message.chat_id or ""),
                chat_type=str(message.chat_type or ""),
                sender_open_id=str(getattr(sender_id, "open_id", "") or ""),
                text=text.strip(),
                message_type=str(message.message_type or ""),
                mentioned_bot=mentioned_bot,
            )
            self.callback(incoming)
        except Exception:
            logger.exception("Failed to parse Feishu event")

    def start(self) -> threading.Thread:
        dispatcher = (
            lark.EventDispatcherHandler.builder("", "")
            .register_p2_im_message_receive_v1(self._handle)
            .build()
        )
        domain = lark.LARK_DOMAIN if "larksuite" in self.settings.feishu_domain else lark.FEISHU_DOMAIN
        self._client = lark.ws.Client(
            self.settings.feishu_app_id,
            self.settings.feishu_app_secret,
            event_handler=dispatcher,
            domain=domain,
            # INFO includes the full WebSocket URL and its ephemeral access key.
            log_level=lark.LogLevel.CRITICAL,
        )
        self._thread = threading.Thread(target=self._client.start, name="feishu-websocket", daemon=True)
        self._thread.start()
        logger.info("Feishu WebSocket thread started")
        return self._thread