WechatOnCloud/bridge/woc_bridge/routes/messages.py

from __future__ import annotations

import asyncio
import logging
from typing import Optional

from fastapi import APIRouter, Query, Request
from fastapi.responses import StreamingResponse

from woc_bridge.config import _require_db_reader, _require_message_streamer
from woc_bridge.db.coordinator import with_db_retry, _check_db_readable
from woc_bridge.models import (
    MessagesResponse,
    MessagesBySessionResponse,
    MessageSearchResponse,
    BridgeError,
)
from woc_bridge.messaging import StreamEvent, format_sse

logger = logging.getLogger("woc-bridge")
router = APIRouter()


@router.get("/api/messages/since", response_model=MessagesResponse)
@with_db_retry
async def get_messages_since(
    cursor: int = 0,
    limit: int = 50,
    is_sender: Optional[bool] = Query(
        None,
        description="按发送方向过滤：True=仅本人发送 / False=仅对方发送 / None=不过滤（默认）",
    ),
) -> MessagesResponse:
    """增量消息拉取（按 create_time 游标）。

    channels PullerAdapter 的核心接口：调用方维护本地 cursor（首次为 0），
    每次拉取 create_time > cursor 的消息，把响应中 next_cursor 存下来作为
    下次的 cursor。has_more=true 时应立即继续拉取，否则按 poll_interval_ms 轮询。

    Args:
        cursor: 上次拉取的最大 create_time，首次传 0（拉全量）
        limit: 最多返回条数，1~200，默认 50。建议用 /api/status 返回的
            max_batch_size 作为上限参考
        is_sender: 按发送方向过滤，True=仅本人发送 / False=仅对方发送 /
            None=不过滤（默认）。SQL 层过滤，无法定位本人时返回空

    Returns:
        MessagesResponse：含 messages 列表 / next_cursor / has_more

    Raises:
        BridgeError(INVALID_PARAMS): cursor<0 或 limit 越界（HTTP 400）
        BridgeError(DB_NOT_FOUND): 未找到微信消息 DB（HTTP 500）
        BridgeError(DB_ENCRYPTED): DB 已加密，需 SQLCipher 密钥（HTTP 503）

    Notes:
        - DB 读取通过 cp 快照避免锁定原文件，sqlite3 同步阻塞，故用
          asyncio.to_thread 包装，避免阻塞 event loop
        - 群消息 content 形如 "wxid:\\n正文"，DbReader 会拆出 sender 字段
        - has_more=true 当且仅当返回条数 >= limit（可能还有更多未拉取）
    """
    # 参数校验
    if cursor < 0:
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"cursor 必须 >= 0，收到 {cursor}",
        )
    if limit < 1 or limit > 200:
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"limit 必须在 1~200 之间，收到 {limit}",
        )

    db_reader = _require_db_reader()

    # DB 加密时返回 DB_ENCRYPTED，而非空数据
    await _check_db_readable()

    # DB 读取是同步阻塞操作，用 to_thread 包装
    result = await asyncio.to_thread(
        db_reader.get_messages_since, cursor, limit, is_sender
    )
    msg_count = len(result.get("messages", []))
    logger.info(
        "messages/since: cursor=%d limit=%d is_sender=%s → 返回 %d 条, next_cursor=%s has_more=%s",
        cursor, limit, is_sender, msg_count, result.get("next_cursor"), result.get("has_more"),
    )
    return MessagesResponse(**result)


# ---------------------------------------------------------------------------
# 路由：GET /api/messages/stream（SSE 实时消息推送）
# ---------------------------------------------------------------------------


@router.get("/api/messages/stream")
async def stream_messages(request: Request) -> StreamingResponse:
    """SSE 实时消息推送长连接。

    建立 text/event-stream 连接后，服务端主动推送新消息事件，替代客户端
    轮询 /api/messages/since。内部由 MessageStreamer 周期性轮询 message DB
    （有订阅者 1s、无订阅者 5s），DB mtime 变化时读增量消息并广播。

    事件类型：
    - sync：连接建立时立即发送，data.cursor 为当前全局游标。客户端据此
      判断是否需要用 /api/messages/since 补全断线期间消息。
    - messages：一批新消息，data 含 messages / next_cursor / has_more，
      结构与 /api/messages/since 响应体一致。事件 id 为 next_cursor。
    - status：DB 状态变化时发送，data 含 db_accessible / db_error_code。
      DB 加密/退出/恢复可读时触发，客户端据此决定是否降级到轮询。
    - heartbeat：每 30 秒心跳，保活 NAT/代理连接，data 为空对象。
    - kicked：订阅者被服务端剔除（超过 MAX_SUBSCRIBERS 上限），data.reason
      给出原因。客户端收到后应关闭连接并按需重连。

    断线补偿：
        SSE 不保证 100% 投递。客户端断线重连后应：
        1. 连接本接口，拿 sync 事件中的 cursor
        2. 用 max(本地cursor, sync.cursor) 调 /api/messages/since 补全
        3. 补全后继续消费本接口事件流

    Returns:
        StreamingResponse：media_type=text/event-stream，持续推送 SSE 帧

    Notes:
        - DB 加密时不关闭连接，改发 status 事件通知客户端；key 提取成功后
          自动恢复推送
        - 响应头设置 X-Accel-Buffering: no 关闭 nginx 缓冲，避免事件被攒批
        - 客户端断开时自动取消订阅，释放队列
        - 不抛业务错误：即使 DB 不可读也建立连接，通过 status 事件告知
        - 订阅者上限由 MessageStreamer.MAX_SUBSCRIBERS 控制，超限时最早
          订阅者会被剔除并收到 kicked 事件
    """
    streamer = _require_message_streamer()

    queue = streamer.subscribe()

    async def event_generator():
        try:
            while True:
                if await request.is_disconnected():
                    break
                try:
                    event = await asyncio.wait_for(queue.get(), timeout=30.0)
                except asyncio.TimeoutError:
                    # 队列 30s 无事件，发心跳保活
                    event = StreamEvent(event="heartbeat", data={})
                # kicked 事件：被服务端剔除（订阅者超上限），发送后立即退出
                if event.event == "kicked":
                    yield format_sse(event)
                    break
                yield format_sse(event)
        finally:
            streamer.unsubscribe(queue)

    headers = {
        "Cache-Control": "no-cache",
        "X-Accel-Buffering": "no",  # nginx 关闭缓冲，确保实时推送
        "Connection": "keep-alive",
    }
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers=headers,
    )


# ---------------------------------------------------------------------------
# 路由：GET /api/messages/by_session
# ---------------------------------------------------------------------------
@router.get("/api/messages/by_session", response_model=MessagesBySessionResponse)
@with_db_retry
async def get_messages_by_session(
    talker: str,
    cursor: int = 0,
    limit: int = 50,
    direction: str = "before",
    is_sender: Optional[bool] = Query(
        None,
        description="按发送方向过滤：True=仅本人发送 / False=仅对方发送 / None=不过滤（默认）",
    ),
) -> MessagesBySessionResponse:
    """按会话拉取历史消息（单会话，O(1) 定位分片表）。

    通过 talker 计算 Msg_<MD5(talker)> 表名直接查单表，避免遍历所有分片。
    用于客户端聊天窗口的历史消息展示与上滑加载。

    Args:
        talker: 会话对方 wxid（群消息为 chatroom id）
        cursor: 游标，首次传 0；before 模式取该时间之前的旧消息，
            after 模式取该时间之后的新消息
        limit: 最多返回条数，1~200，默认 50
        direction: before（默认，往前翻历史）/ after（往后拉新消息）
        is_sender: 按发送方向过滤，True=仅本人发送 / False=仅对方发送 /
            None=不过滤（默认）。SQL 层过滤，无法定位本人时返回空

    Returns:
        MessagesBySessionResponse：含 messages / next_cursor / has_more / talker

    Raises:
        BridgeError(INVALID_PARAMS): talker 为空、limit 越界、direction 非法（HTTP 400）
        BridgeError(DB_NOT_FOUND): 未找到微信消息 DB（HTTP 500）
        BridgeError(DB_ENCRYPTED): DB 已加密（HTTP 503）

    Notes:
        - before 模式返回升序消息（旧在前、新在后），便于客户端追加到列表头部
        - cursor=0 + before 取最新 limit 条；cursor=某消息时间 + before 取更旧的消息
        - has_more=true 时用 next_cursor 继续翻页
        - 表不存在时返回空列表而非报错（会话可能从未有过消息）
    """
    if not talker:
        raise BridgeError(code="INVALID_PARAMS", message="talker 不能为空")
    if limit < 1 or limit > 200:
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"limit 必须在 1~200 之间，收到 {limit}",
        )
    if direction not in ("before", "after"):
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"direction 必须为 before 或 after，收到 {direction}",
        )
    if cursor < 0:
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"cursor 必须 >= 0，收到 {cursor}",
        )

    db_reader = _require_db_reader()
    await _check_db_readable()

    result = await asyncio.to_thread(
        db_reader.get_messages_by_session, talker, cursor, limit, direction, is_sender
    )
    msg_count = len(result.get("messages", []))
    logger.info(
        "messages/by_session: talker=%s cursor=%d direction=%s limit=%d is_sender=%s → 返回 %d 条, next_cursor=%s has_more=%s",
        talker, cursor, direction, limit, is_sender, msg_count,
        result.get("next_cursor"), result.get("has_more"),
    )
    return MessagesBySessionResponse(**result)


# ---------------------------------------------------------------------------
# 路由：GET /api/messages/search
# ---------------------------------------------------------------------------
@router.get("/api/messages/search", response_model=MessageSearchResponse)
@with_db_retry
async def search_messages(
    keyword: str,
    talker: str = "",
    start_time: int = 0,
    end_time: int = 0,
    limit: int = 50,
    is_sender: Optional[bool] = Query(
        None,
        description="按发送方向过滤：True=仅本人发送 / False=仅对方发送 / None=不过滤（默认）",
    ),
) -> MessageSearchResponse:
    """按关键词搜索历史消息。

    遍历所有 Msg_* 分片表（或指定 talker 的单表），对 message_content
    做 LIKE 模糊匹配，合并后按 create_time 降序返回（最新匹配在前）。

    Args:
        keyword: 搜索关键词（非空）
        talker: 可选，限定在指定会话内搜索
        start_time: 可选，起始时间戳（含），0 表示不限制
        end_time: 可选，结束时间戳（含），0 表示不限制
        limit: 最多返回条数，1~200，默认 50
        is_sender: 按发送方向过滤，True=仅本人发送 / False=仅对方发送 /
            None=不过滤（默认）。SQL 层过滤，无法定位本人时返回空

    Returns:
        MessageSearchResponse：含 messages / total

    Raises:
        BridgeError(INVALID_PARAMS): keyword 为空或 limit 越界（HTTP 400）
        BridgeError(DB_NOT_FOUND): 未找到微信消息 DB（HTTP 500）
        BridgeError(DB_ENCRYPTED): DB 已加密（HTTP 503）

    Notes:
        - 搜索基于 SQL LIKE，性能取决于 DB 大小与索引；大库可能较慢
        - 限制：无法搜索 zstd 压缩消息（WCDB_CT_message_content==4 时存储
          压缩 BLOB，LIKE 不匹配二进制）；压缩消息会被静默漏掉
        - keyword 中的 % 和 _ 已转义为字面量，按子串匹配
        - total 为合并后返回的条数（受每表 LIMIT 截断，可能小于实际命中数）
        - 群消息的 sender 前缀（"wxid:\\n"）会先被剥离再匹配
    """
    if not keyword:
        raise BridgeError(code="INVALID_PARAMS", message="keyword 不能为空")
    if limit < 1 or limit > 200:
        raise BridgeError(
            code="INVALID_PARAMS",
            message=f"limit 必须在 1~200 之间，收到 {limit}",
        )

    db_reader = _require_db_reader()
    await _check_db_readable()

    result = await asyncio.to_thread(
        db_reader.search_messages,
        keyword,
        talker if talker else None,
        start_time if start_time > 0 else None,
        end_time if end_time > 0 else None,
        limit,
        is_sender,
    )
    msg_count = len(result.get("messages", []))
    logger.info(
        "messages/search: keyword=%r talker=%s is_sender=%s → 命中 %d 条（返回 %d 条）",
        keyword, talker or "(全部)", is_sender, result.get("total", 0), msg_count,
    )
    return MessageSearchResponse(**result)