LangBot/tests/factories/message.py

"""
Message and query factories for tests.

Provides reusable factories for creating message chains, events, and query objects.
"""

from __future__ import annotations

from unittest.mock import AsyncMock, Mock
import typing

import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.platform.events as platform_events
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext


# Counter for generating unique IDs
_query_counter = 0


def _next_query_id() -> int:
    """Generate a unique query ID."""
    global _query_counter
    _query_counter += 1
    return _query_counter


# ============== Message Chain Factories ==============


def text_chain(text: str = 'hello') -> platform_message.MessageChain:
    """Create a simple text message chain."""
    return platform_message.MessageChain(
        [
            platform_message.Plain(text=text),
        ]
    )


def group_text_chain(text: str = 'hello') -> platform_message.MessageChain:
    """Create a group text message chain (same as text_chain, context provided by event)."""
    return text_chain(text)


def mention_chain(
    text: str = 'hello',
    target: typing.Union[int, str] = 12345,
) -> platform_message.MessageChain:
    """Create a message chain with @mention."""
    return platform_message.MessageChain(
        [
            platform_message.At(target=target),
            platform_message.Plain(text=f' {text}'),
        ]
    )


def image_chain(
    text: str = '',
    url: str = 'https://example.com/image.png',
) -> platform_message.MessageChain:
    """Create a message chain with an image."""
    components = []
    if text:
        components.append(platform_message.Plain(text=text))
    components.append(platform_message.Image(url=url))
    return platform_message.MessageChain(components)


def command_chain(
    command: str = 'help',
    prefix: str = '/',
) -> platform_message.MessageChain:
    """Create a command message chain."""
    return platform_message.MessageChain(
        [
            platform_message.Plain(text=f'{prefix}{command}'),
        ]
    )


# ============== Message Event Factories ==============


def friend_message_event(
    message_chain: platform_message.MessageChain,
    sender_id: typing.Union[int, str] = 12345,
    nickname: str = 'TestUser',
) -> platform_events.FriendMessage:
    """Create a friend (private) message event."""
    sender = platform_entities.Friend(
        id=sender_id,
        nickname=nickname,
        remark=None,
    )
    return platform_events.FriendMessage(
        type='FriendMessage',
        sender=sender,
        message_chain=message_chain,
        time=1609459200,
    )


def group_message_event(
    message_chain: platform_message.MessageChain,
    sender_id: typing.Union[int, str] = 12345,
    sender_name: str = 'TestUser',
    group_id: typing.Union[int, str] = 99999,
    group_name: str = 'TestGroup',
) -> platform_events.GroupMessage:
    """Create a group message event."""
    group = platform_entities.Group(
        id=group_id,
        name=group_name,
        permission=platform_entities.Permission.Member,
    )
    sender = platform_entities.GroupMember(
        id=sender_id,
        member_name=sender_name,
        permission=platform_entities.Permission.Member,
        group=group,
    )
    return platform_events.GroupMessage(
        type='GroupMessage',
        sender=sender,
        message_chain=message_chain,
        time=1609459200,
    )


# ============== Mock Adapter Factory ==============


def mock_adapter() -> Mock:
    """Create a mock platform adapter."""
    adapter = AsyncMock()
    adapter.is_stream_output_supported = AsyncMock(return_value=False)
    adapter.reply_message = AsyncMock()
    adapter.reply_message_chunk = AsyncMock()
    return adapter


# ============== Query Factories ==============


def _base_query(
    message_chain: platform_message.MessageChain,
    message_event: platform_events.MessageEvent,
    launcher_type: provider_session.LauncherTypes,
    launcher_id: typing.Union[int, str],
    sender_id: typing.Union[int, str],
    adapter: Mock,
    **overrides,
) -> pipeline_query.Query:
    """Create a base query with model_construct to bypass validation."""
    query_id = _next_query_id()

    base_data = {
        'query_id': query_id,
        'launcher_type': launcher_type,
        'launcher_id': launcher_id,
        'sender_id': sender_id,
        'message_chain': message_chain,
        'message_event': message_event,
        'adapter': adapter,
        'pipeline_uuid': 'test-pipeline-uuid',
        'bot_uuid': 'test-bot-uuid',
        'pipeline_config': {
            'ai': {
                'runner': {'runner': 'local-agent'},
                'local-agent': {
                    'model': {'primary': 'test-model-uuid', 'fallbacks': []},
                    'prompt': 'test-prompt',
                },
            },
            'output': {'misc': {'at-sender': False, 'quote-origin': False}},
            'trigger': {'misc': {'combine-quote-message': False}},
        },
        'session': None,
        'prompt': None,
        'messages': [],
        'user_message': None,
        'use_funcs': [],
        'use_llm_model_uuid': None,
        'variables': {},
        'resp_messages': [],
        'resp_message_chain': None,
        'current_stage_name': None,
    }

    # Apply overrides
    for key, value in overrides.items():
        base_data[key] = value

    query = pipeline_query.Query.model_construct(**base_data)
    object.__setattr__(
        query,
        '_execution_context',
        ExecutionContext(
            instance_uuid='test-instance',
            workspace_uuid='test-workspace',
            placement_generation=1,
            bot_uuid=query.bot_uuid,
            pipeline_uuid=query.pipeline_uuid,
            query_uuid=query.query_uuid,
        ),
    )
    return query


def text_query(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a basic text query (private chat)."""
    chain = text_chain(text)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def private_text_query(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a private text query (alias for text_query)."""
    return text_query(text, sender_id, **overrides)


def group_text_query(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    group_id: typing.Union[int, str] = 99999,
    **overrides,
) -> pipeline_query.Query:
    """Create a group text query."""
    chain = text_chain(text)
    event = group_message_event(chain, sender_id, group_id=group_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.GROUP,
        launcher_id=group_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def command_query(
    command: str = 'help',
    prefix: str = '/',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a command-like query."""
    chain = command_chain(command, prefix)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def mention_query(
    text: str = 'hello',
    target: typing.Union[int, str] = 12345,
    sender_id: typing.Union[int, str] = 12345,
    group_id: typing.Union[int, str] = 99999,
    **overrides,
) -> pipeline_query.Query:
    """Create a mention-bot query (group chat with @mention)."""
    chain = mention_chain(text, target)
    event = group_message_event(chain, sender_id, group_id=group_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.GROUP,
        launcher_id=group_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def empty_query(**overrides) -> pipeline_query.Query:
    """Create an empty message query."""
    chain = platform_message.MessageChain([])
    event = friend_message_event(chain)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=12345,
        sender_id=12345,
        adapter=adapter,
        **overrides,
    )


def image_query(
    text: str = '',
    url: str = 'https://example.com/image.png',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create an image query."""
    chain = image_chain(text, url)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def file_query(
    url: str = 'https://example.com/document.pdf',
    name: str = 'document.pdf',
    text: str = '',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a file attachment query."""
    components = []
    if text:
        components.append(platform_message.Plain(text=text))
    components.append(platform_message.File(url=url, name=name))
    chain = platform_message.MessageChain(components)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def unsupported_query(
    unsupported_type: str = 'CustomComponent',
    text: str = '',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a query with unsupported/unknown message segment."""
    components = []
    if text:
        components.append(platform_message.Plain(text=text))
    # Use Unknown component for unsupported types
    components.append(platform_message.Unknown(text=f'Unsupported: {unsupported_type}'))
    chain = platform_message.MessageChain(components)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def query_with_session(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    session: provider_session.Session = None,
    **overrides,
) -> pipeline_query.Query:
    """Create a query with a session object.

    If session is None, creates a default session with empty conversation.
    """
    if session is None:
        # Create a default session
        session = provider_session.Session(
            launcher_type=provider_session.LauncherTypes.PERSON,
            launcher_id=sender_id,
            sender_id=sender_id,
            use_prompt_name='default',
            using_conversation=None,
            conversations=[],
        )

    return text_query(text, sender_id, session=session, **overrides)


def query_with_config(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    pipeline_config: dict = None,
    **overrides,
) -> pipeline_query.Query:
    """Create a query with custom pipeline configuration.

    If pipeline_config is None, uses default config.
    Useful for testing specific stage behaviors.
    """
    if pipeline_config is None:
        pipeline_config = {
            'ai': {
                'runner': {'runner': 'local-agent'},
                'local-agent': {
                    'model': {'primary': 'test-model-uuid', 'fallbacks': []},
                    'prompt': 'test-prompt',
                },
            },
            'output': {'misc': {'at-sender': False, 'quote-origin': False}},
            'trigger': {'misc': {'combine-quote-message': False}},
        }

    return text_query(text, sender_id, pipeline_config=pipeline_config, **overrides)


def voice_query(
    url: str = 'https://example.com/audio.mp3',
    sender_id: typing.Union[int, str] = 12345,
    **overrides,
) -> pipeline_query.Query:
    """Create a voice/audio query."""
    components = [
        platform_message.Voice(url=url),
    ]
    chain = platform_message.MessageChain(components)
    event = friend_message_event(chain, sender_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.PERSON,
        launcher_id=sender_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )


def at_all_query(
    text: str = 'hello',
    sender_id: typing.Union[int, str] = 12345,
    group_id: typing.Union[int, str] = 99999,
    **overrides,
) -> pipeline_query.Query:
    """Create a group query with @All mention."""
    components = [
        platform_message.AtAll(),
        platform_message.Plain(text=f' {text}'),
    ]
    chain = platform_message.MessageChain(components)
    event = group_message_event(chain, sender_id, group_id=group_id)
    adapter = mock_adapter()
    return _base_query(
        message_chain=chain,
        message_event=event,
        launcher_type=provider_session.LauncherTypes.GROUP,
        launcher_id=group_id,
        sender_id=sender_id,
        adapter=adapter,
        **overrides,
    )