ForcePilot/backend/package/yuxi/channel/extensions/qqbot/pipeline.py

from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from yuxi.channel.extensions.qqbot.types import QQBotChatType, QueuedMessage

if TYPE_CHECKING:
    from yuxi.channel.common.attachment_cache import AttachmentCache


@dataclass
class PipelineContext:
    msg: QueuedMessage
    account_id: str
    config: dict = field(default_factory=dict)
    blocked: bool = False
    skip_reason: str | None = None
    metadata: dict = field(default_factory=dict)


class QQBotPipeline:
    def __init__(self):
        self._stages: list[Any] = []

    def add_stage(self, stage: Any) -> None:
        self._stages.append(stage)

    async def process(self, msg: QueuedMessage, account_id: str, config: dict) -> PipelineContext:
        ctx = PipelineContext(msg=msg, account_id=account_id, config=config)

        for stage in self._stages:
            if ctx.blocked:
                break
            try:
                ctx = await stage(ctx)
            except Exception:
                pass

        return ctx


class DedupeStage:
    def __init__(self, ttl_seconds: int = 300):
        self._seen: dict[str, float] = {}
        self._ttl = ttl_seconds

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        self._cleanup()
        msg_id = ctx.msg.msg_id
        if msg_id in self._seen:
            ctx.blocked = True
            ctx.skip_reason = "duplicate"
            return ctx
        import time

        self._seen[msg_id] = time.time()
        return ctx

    def _cleanup(self) -> None:
        import time

        now = time.time()
        expired = [k for k, v in self._seen.items() if now - v > self._ttl]
        for k in expired:
            del self._seen[k]


class AccessStage:
    def __init__(self, check_fn: Any | None = None):
        self._check_fn = check_fn

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        if self._check_fn:
            allowed = await self._check_fn(ctx.msg.sender_id, ctx.msg.chat_type)
            if not allowed:
                ctx.blocked = True
                ctx.skip_reason = "access_denied"
        return ctx


class ContentStage:
    def __init__(self, attachment_cache: AttachmentCache | None = None):
        self._attachment_cache = attachment_cache

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        session_id = ctx.msg.sender_id

        if self._attachment_cache is not None and ctx.msg.attachments:
            if not ctx.msg.content.strip():
                for att in ctx.msg.attachments:
                    await self._attachment_cache.add(
                        session_id,
                        att.url,
                        att.content_type or "",
                        filename=att.filename,
                    )
                ctx.blocked = True
                ctx.skip_reason = "attachment_cached"
                return ctx

            cached = await self._attachment_cache.consume(session_id)
            if cached:
                cached_tags = _build_cached_attachment_tags(cached)
                ctx.msg.content = f"{ctx.msg.content}\n{cached_tags}" if ctx.msg.content else cached_tags

        if ctx.msg.attachments:
            ctx.metadata["has_attachments"] = True
            ctx.metadata["attachment_count"] = len(ctx.msg.attachments)
            ctx.msg.content = _inject_attachment_tags(ctx.msg.content, ctx.msg.attachments)
        return ctx


def _build_cached_attachment_tags(cached: list) -> str:
    tags: list[str] = []
    for att in cached:
        ct = (att.content_type or "").lower()
        if ct.startswith("image/"):
            tag = f"[图片: {att.url}]"
        elif ct.startswith("video/"):
            tag = f"[视频: {att.url}]"
        elif ct.startswith("audio/") or ct.startswith("voice/"):
            tag = f"[语音: {att.url}]"
        else:
            label = att.filename or att.url
            tag = f"[文件: {label}]"
        tags.append(tag)
    return "\n".join(tags)


def _inject_attachment_tags(content: str, attachments: list) -> str:
    tags: list[str] = []
    for att in attachments:
        ct = (att.content_type or "").lower()
        if ct.startswith("image/"):
            tag = f"[图片: {att.url}]"
        elif ct.startswith("video/"):
            tag = f"[视频: {att.url}]"
        elif ct.startswith("audio/") or ct.startswith("voice/"):
            tag = f"[语音: {att.url}]"
        else:
            label = att.filename or att.url
            tag = f"[文件: {label}]"
        tags.append(tag)
    if tags:
        tag_section = "\n".join(tags)
        content = f"{content}\n{tag_section}" if content else tag_section
    return content


class GroupGateStage:
    def __init__(self, gate_fn: Any | None = None, history_fn: Any | None = None):
        self._gate_fn = gate_fn
        self._history_fn = history_fn

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        if ctx.msg.chat_type not in (QQBotChatType.GROUP, QQBotChatType.GUILD):
            return ctx

        if self._gate_fn:
            decision = await self._gate_fn(ctx)
            if decision and decision != "process":
                ctx.blocked = True
                ctx.skip_reason = decision
                if self._history_fn and decision in ("skip_no_mention", "drop_other_mention"):
                    await self._history_fn(ctx)
                return ctx

        if self._history_fn and not ctx.blocked:
            await self._history_fn(ctx)

        return ctx


class AssemblyStage:
    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        if ctx.msg.sender_name and ctx.msg.chat_type == QQBotChatType.GROUP:
            ctx.metadata["agent_content"] = f"[{ctx.msg.sender_name}]: {ctx.msg.content}"
        else:
            ctx.metadata["agent_content"] = ctx.msg.content
        return ctx


class EnvelopeStage:
    def __init__(self, dispatch_fn: Any | None = None):
        self._dispatch_fn = dispatch_fn

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        if ctx.blocked or not self._dispatch_fn:
            return ctx
        await self._dispatch_fn(ctx)
        return ctx


class RefIndexStage:
    def __init__(self, index_fn: Any | None = None):
        self._index_fn = index_fn

    async def __call__(self, ctx: PipelineContext) -> PipelineContext:
        if self._index_fn:
            await self._index_fn(ctx)
        return ctx


def create_default_pipeline(
    dispatch_fn=None,
    access_check_fn=None,
    gate_fn=None,
    history_fn=None,
    ref_index_fn=None,
    attachment_cache: AttachmentCache | None = None,
) -> QQBotPipeline:
    pipeline = QQBotPipeline()
    pipeline.add_stage(DedupeStage())
    pipeline.add_stage(AccessStage(check_fn=access_check_fn))
    pipeline.add_stage(ContentStage(attachment_cache=attachment_cache))
    pipeline.add_stage(GroupGateStage(gate_fn=gate_fn, history_fn=history_fn))
    pipeline.add_stage(AssemblyStage())
    pipeline.add_stage(RefIndexStage(index_fn=ref_index_fn))
    pipeline.add_stage(EnvelopeStage(dispatch_fn=dispatch_fn))
    return pipeline