# tests/test_subagent_wait.py
"""Tests for wait_for_subagents with top-level and child subagents."""
import pytest
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock
from nanobot.agent.subagent import SubagentManager
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMProvider, LLMResponse


@pytest.mark.asyncio
async def test_wait_for_top_level_subagent():
    """Test that wait_for works for top-level subagents spawned from telegram."""
    bus = MessageBus()
    provider = MagicMock(spec=LLMProvider)
    provider.chat = AsyncMock(return_value=LLMResponse(
        content="Task completed",
        tool_calls=[]
    ))
    provider.get_default_model = MagicMock(return_value="test-model")
    provider.thinking_budget = 0

    workspace = Path("/tmp/test-subagent")
    workspace.mkdir(exist_ok=True)

    manager = SubagentManager(
        bus=bus,
        provider=provider,
        workspace=workspace
    )

    # Spawn a top-level subagent (origin channel = "telegram")
    task_id = await manager.spawn(
        task="Test task",
        label="test",
        model=None,
        origin_channel="telegram",
        origin_chat_id="12345"
    )

    # Wait for it to complete
    result = await manager.wait_for([task_id])

    # Should find the result (not "No result found")
    assert "No result found" not in result
    assert task_id in result
    assert "Task completed" in result or "completed" in result.lower()


@pytest.mark.asyncio
async def test_wait_for_child_subagent():
    """Test that wait_for works for child subagents (orchestrator pattern)."""
    bus = MessageBus()
    provider = MagicMock(spec=LLMProvider)
    provider.chat = AsyncMock(return_value=LLMResponse(
        content="Child task completed",
        tool_calls=[]
    ))
    provider.get_default_model = MagicMock(return_value="test-model")
    provider.thinking_budget = 0

    workspace = Path("/tmp/test-subagent")
    workspace.mkdir(exist_ok=True)

    manager = SubagentManager(
        bus=bus,
        provider=provider,
        workspace=workspace
    )

    # Spawn a child subagent (origin channel = "subagent")
    task_id = await manager.spawn(
        task="Child test task",
        label="test-child",
        model=None,
        origin_channel="subagent",
        origin_chat_id="parent-id"
    )

    # Wait for it to complete
    result = await manager.wait_for([task_id])

    # Should find the result (not "No result found")
    assert "No result found" not in result
    assert task_id in result
    assert "Child task completed" in result or "completed" in result.lower()


@pytest.mark.asyncio
async def test_wait_for_multiple_subagents():
    """Test waiting for multiple subagents of different types."""
    bus = MessageBus()
    call_count = 0

    async def chat_response(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        return LLMResponse(content=f"Task {call_count} completed", tool_calls=[])

    provider = MagicMock(spec=LLMProvider)
    provider.chat = AsyncMock(side_effect=chat_response)
    provider.get_default_model = MagicMock(return_value="test-model")
    provider.thinking_budget = 0

    workspace = Path("/tmp/test-subagent")
    workspace.mkdir(exist_ok=True)

    manager = SubagentManager(
        bus=bus,
        provider=provider,
        workspace=workspace
    )

    # Spawn one top-level and one child subagent
    task_id_1 = await manager.spawn(
        task="Top-level task",
        label="test-top",
        model=None,
        origin_channel="telegram",
        origin_chat_id="12345"
    )

    task_id_2 = await manager.spawn(
        task="Child task",
        label="test-child",
        model=None,
        origin_channel="subagent",
        origin_chat_id="parent"
    )

    # Wait for both
    result = await manager.wait_for([task_id_1, task_id_2])

    # Should find both results
    assert "No result found" not in result
    assert task_id_1 in result
    assert task_id_2 in result
    assert "Task 1 completed" in result or "completed" in result.lower()
    assert "Task 2 completed" in result or "completed" in result.lower()
