"""Tests for sub-agent spawn functionality."""

import asyncio
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock

import pytest

from oclawma.providers.base import (
    CompletionResponse,
    UsageStats,
)
from oclawma.subagent import (
    AggregationConfig,
    SubAgent,
    SubAgentConfig,
    SubAgentManager,
    SubAgentPool,
    SubAgentResult,
    SubAgentStatus,
)


class TestSubAgentConfig:
    """Test SubAgentConfig dataclass."""

    def test_default_config(self) -> None:
        """Test default configuration values."""
        config = SubAgentConfig()

        assert config.model is None
        assert config.temperature == 0.7
        assert config.max_tokens == 2000
        assert config.max_iterations == 5
        assert config.timeout_seconds == 300.0
        assert config.system_prompt is None
        assert config.tools_enabled is True
        assert config.context_budget == 4000

    def test_custom_config(self) -> None:
        """Test custom configuration values."""
        config = SubAgentConfig(
            model="gpt-4",
            temperature=0.5,
            max_tokens=1000,
            timeout_seconds=60.0,
            tools_enabled=False,
            context_budget=2000,
        )

        assert config.model == "gpt-4"
        assert config.temperature == 0.5
        assert config.max_tokens == 1000
        assert config.timeout_seconds == 60.0
        assert config.tools_enabled is False
        assert config.context_budget == 2000


class TestSubAgentResult:
    """Test SubAgentResult dataclass."""

    def test_success_property(self) -> None:
        """Test success property."""
        result_success = SubAgentResult(
            agent_id="test1",
            task="test task",
            status=SubAgentStatus.COMPLETED,
            output="result",
        )
        assert result_success.success is True

        result_failed = SubAgentResult(
            agent_id="test2",
            task="test task",
            status=SubAgentStatus.FAILED,
            error="something went wrong",
        )
        assert result_failed.success is False

        result_error = SubAgentResult(
            agent_id="test3",
            task="test task",
            status=SubAgentStatus.COMPLETED,
            error="has error but completed",
        )
        assert result_error.success is False

    def test_duration_seconds(self) -> None:
        """Test duration calculation."""
        start = datetime(2024, 1, 1, 12, 0, 0)
        end = datetime(2024, 1, 1, 12, 0, 5)

        result = SubAgentResult(
            agent_id="test",
            task="test",
            status=SubAgentStatus.COMPLETED,
            start_time=start,
            end_time=end,
        )

        assert result.duration_seconds == 5.0

    def test_duration_no_times(self) -> None:
        """Test duration with no timestamps."""
        result = SubAgentResult(
            agent_id="test",
            task="test",
            status=SubAgentStatus.PENDING,
        )

        assert result.duration_seconds == 0.0

    def test_to_dict(self) -> None:
        """Test conversion to dictionary."""
        start = datetime(2024, 1, 1, 12, 0, 0)
        end = datetime(2024, 1, 1, 12, 0, 5)

        result = SubAgentResult(
            agent_id="abc123",
            task="test task",
            status=SubAgentStatus.COMPLETED,
            output="test output",
            tokens_used=100,
            start_time=start,
            end_time=end,
            metadata={"key": "value"},
        )

        data = result.to_dict()

        assert data["agent_id"] == "abc123"
        assert data["task"] == "test task"
        assert data["status"] == "completed"
        assert data["output"] == "test output"
        assert data["tokens_used"] == 100
        assert data["duration_seconds"] == 5.0
        assert data["metadata"] == {"key": "value"}


class TestSubAgent:
    """Test SubAgent class."""

    @pytest.fixture
    def mock_provider(self) -> MagicMock:
        """Create a mock provider."""
        provider = MagicMock()
        provider.complete = AsyncMock()
        return provider

    def test_init(self, mock_provider: MagicMock) -> None:
        """Test sub-agent initialization."""
        agent = SubAgent(
            provider=mock_provider,
            task="test task",
        )

        assert agent.provider == mock_provider
        assert agent.task == "test task"
        assert agent.status == SubAgentStatus.PENDING
        assert len(agent.agent_id) == 8  # UUID truncated

    def test_init_with_custom_id(self, mock_provider: MagicMock) -> None:
        """Test sub-agent with custom ID."""
        agent = SubAgent(
            provider=mock_provider,
            task="test task",
            agent_id="custom123",
        )

        assert agent.agent_id == "custom123"

    def test_create_system_prompt_default(self, mock_provider: MagicMock) -> None:
        """Test default system prompt creation."""
        agent = SubAgent(
            provider=mock_provider,
            task="analyze this data",
        )

        prompt = agent._create_system_prompt()

        assert "analyze this data" in prompt
        assert "specialized sub-agent" in prompt
        assert "Guidelines:" in prompt

    def test_create_system_prompt_with_parent_context(self, mock_provider: MagicMock) -> None:
        """Test system prompt with parent context."""
        agent = SubAgent(
            provider=mock_provider,
            task="analyze this data",
            parent_context={"context_hint": "This is a database query"},
        )

        prompt = agent._create_system_prompt()

        assert "This is a database query" in prompt

    def test_create_system_prompt_custom(self, mock_provider: MagicMock) -> None:
        """Test custom system prompt."""
        config = SubAgentConfig(system_prompt="Custom prompt here")
        agent = SubAgent(
            provider=mock_provider,
            task="test",
            config=config,
        )

        prompt = agent._create_system_prompt()

        assert prompt == "Custom prompt here"

    @pytest.mark.asyncio
    async def test_run_success(self, mock_provider: MagicMock) -> None:
        """Test successful execution."""
        mock_response = CompletionResponse(
            content="Task completed successfully",
            model="test-model",
            usage=UsageStats(prompt_tokens=50, completion_tokens=50, total_tokens=100),
        )
        mock_provider.complete.return_value = mock_response

        agent = SubAgent(
            provider=mock_provider,
            task="do something",
            agent_id="test123",
        )

        result = await agent.run()

        assert result.status == SubAgentStatus.COMPLETED
        assert result.output == "Task completed successfully"
        assert result.tokens_used == 100
        assert result.agent_id == "test123"
        assert result.success is True

        # Verify provider was called correctly
        mock_provider.complete.assert_called_once()
        call_args = mock_provider.complete.call_args[0][0]
        assert call_args.model == "qwen2.5:3b"  # Default model

    @pytest.mark.asyncio
    async def test_run_failure(self, mock_provider: MagicMock) -> None:
        """Test failed execution."""
        mock_provider.complete.side_effect = Exception("Provider error")

        agent = SubAgent(
            provider=mock_provider,
            task="do something",
        )

        result = await agent.run()

        assert result.status == SubAgentStatus.FAILED
        assert result.error == "Provider error"
        assert result.success is False

    @pytest.mark.asyncio
    async def test_run_timeout(self, mock_provider: MagicMock) -> None:
        """Test timeout handling."""
        config = SubAgentConfig(timeout_seconds=0.001)

        async def slow_complete(*args, **kwargs):  # type: ignore
            await asyncio.sleep(1)  # Will timeout
            return CompletionResponse(
                content="slow",
                model="test",
                usage=UsageStats(0, 0, 0),
            )

        mock_provider.complete = slow_complete

        agent = SubAgent(
            provider=mock_provider,
            task="slow task",
            config=config,
        )

        result = await agent.run()

        assert result.status == SubAgentStatus.FAILED
        assert "timed out" in result.error
        assert result.success is False

    @pytest.mark.asyncio
    async def test_start(self, mock_provider: MagicMock) -> None:
        """Test async start."""
        mock_response = CompletionResponse(
            content="done",
            model="test",
            usage=UsageStats(0, 0, 0),
        )
        mock_provider.complete.return_value = mock_response

        agent = SubAgent(
            provider=mock_provider,
            task="test",
        )

        task = await agent.start()

        assert isinstance(task, asyncio.Task)

        result = await task

        assert result.success is True

    @pytest.mark.asyncio
    async def test_cancel(self, mock_provider: MagicMock) -> None:
        """Test cancellation."""

        async def slow_complete(*args, **kwargs):  # type: ignore
            await asyncio.sleep(10)
            return CompletionResponse(
                content="slow",
                model="test",
                usage=UsageStats(0, 0, 0),
            )

        mock_provider.complete = slow_complete

        agent = SubAgent(
            provider=mock_provider,
            task="slow task",
        )

        # Start the agent
        await agent.start()

        # Cancel immediately
        await agent.cancel()

        assert agent.status == SubAgentStatus.CANCELLED


class TestSubAgentPool:
    """Test SubAgentPool class."""

    @pytest.fixture
    def mock_provider(self) -> MagicMock:
        """Create a mock provider."""
        provider = MagicMock()
        provider.complete = AsyncMock()
        return provider

    def test_init(self, mock_provider: MagicMock) -> None:
        """Test pool initialization."""
        pool = SubAgentPool(
            provider=mock_provider,
            max_concurrent=8,
        )

        assert pool.provider == mock_provider
        assert pool.max_concurrent == 8
        assert pool._semaphore._value == 8

    @pytest.mark.asyncio
    async def test_spawn(self, mock_provider: MagicMock) -> None:
        """Test spawning an agent."""
        pool = SubAgentPool(provider=mock_provider)

        agent = await pool.spawn("test task")

        assert agent.task == "test task"
        assert len(pool._agents) == 1
        assert pool._agents[0] == agent

    @pytest.mark.asyncio
    async def test_spawn_many(self, mock_provider: MagicMock) -> None:
        """Test spawning multiple agents."""
        pool = SubAgentPool(provider=mock_provider)

        tasks = ["task1", "task2", "task3"]
        agents = await pool.spawn_many(tasks)

        assert len(agents) == 3
        assert len(pool._agents) == 3
        assert agents[0].task == "task1"
        assert agents[1].task == "task2"
        assert agents[2].task == "task3"

    @pytest.mark.asyncio
    async def test_execute_concurrent(self, mock_provider: MagicMock) -> None:
        """Test concurrent execution."""
        mock_response = CompletionResponse(
            content="done",
            model="test",
            usage=UsageStats(10, 10, 20),
        )
        mock_provider.complete.return_value = mock_response

        pool = SubAgentPool(provider=mock_provider, max_concurrent=2)

        # Spawn agents
        await pool.spawn_many(["task1", "task2"])

        # Execute
        results = await pool.execute_concurrent()

        assert len(results) == 2
        assert all(r.success for r in results)
        assert results[0].output == "done"
        assert results[1].output == "done"

    @pytest.mark.asyncio
    async def test_execute_concurrent_with_failures(self, mock_provider: MagicMock) -> None:
        """Test concurrent execution with failures."""
        # First call succeeds, second fails
        mock_provider.complete.side_effect = [
            CompletionResponse(
                content="success",
                model="test",
                usage=UsageStats(10, 10, 20),
            ),
            Exception("Provider error"),
        ]

        pool = SubAgentPool(provider=mock_provider)

        await pool.spawn_many(["task1", "task2"])
        results = await pool.execute_concurrent()

        assert len(results) == 2
        assert results[0].success is True
        assert results[1].success is False
        assert "Provider error" in results[1].error

    @pytest.mark.asyncio
    async def test_execute_concurrent_progress_callback(self, mock_provider: MagicMock) -> None:
        """Test progress callback."""
        mock_response = CompletionResponse(
            content="done",
            model="test",
            usage=UsageStats(0, 0, 0),
        )
        mock_provider.complete.return_value = mock_response

        pool = SubAgentPool(provider=mock_provider)
        await pool.spawn_many(["task1", "task2", "task3"])

        progress_calls = []

        def progress(completed: int, total: int) -> None:
            progress_calls.append((completed, total))

        await pool.execute_concurrent(progress_callback=progress)

        assert len(progress_calls) == 3
        assert progress_calls[-1] == (3, 3)

    def test_aggregate_concatenate(self, mock_provider: MagicMock) -> None:
        """Test concatenate aggregation."""
        pool = SubAgentPool(
            provider=mock_provider,
            aggregation_config=AggregationConfig(strategy="concatenate"),
        )

        results = [
            SubAgentResult(
                agent_id="a1",
                task="t1",
                status=SubAgentStatus.COMPLETED,
                output="Result 1",
            ),
            SubAgentResult(
                agent_id="a2",
                task="t2",
                status=SubAgentStatus.COMPLETED,
                output="Result 2",
            ),
        ]

        aggregated = pool.aggregate_results(results)

        assert "Result 1" in aggregated
        assert "Result 2" in aggregated
        assert "Agent: a1" in aggregated
        assert "Agent: a2" in aggregated

    def test_aggregate_summarize(self, mock_provider: MagicMock) -> None:
        """Test summarize aggregation."""
        pool = SubAgentPool(
            provider=mock_provider,
            aggregation_config=AggregationConfig(strategy="summarize"),
        )

        results = [
            SubAgentResult(
                agent_id="a1",
                task="t1",
                status=SubAgentStatus.COMPLETED,
                output="Result 1",
                tokens_used=100,
            ),
            SubAgentResult(
                agent_id="a2",
                task="t2",
                status=SubAgentStatus.FAILED,
                error="Failed",
            ),
        ]

        aggregated = pool.aggregate_results(results)

        assert "Total agents: 2" in aggregated
        assert "Successful: 1" in aggregated
        assert "Failed: 1" in aggregated
        assert "Total tokens used: 100" in aggregated

    def test_aggregate_vote(self, mock_provider: MagicMock) -> None:
        """Test vote aggregation."""
        pool = SubAgentPool(
            provider=mock_provider,
            aggregation_config=AggregationConfig(strategy="vote"),
        )

        results = [
            SubAgentResult(
                agent_id="a1",
                task="t1",
                status=SubAgentStatus.COMPLETED,
                output="Option A",
            ),
            SubAgentResult(
                agent_id="a2",
                task="t2",
                status=SubAgentStatus.COMPLETED,
                output="Option A",  # Same answer
            ),
            SubAgentResult(
                agent_id="a3",
                task="t3",
                status=SubAgentStatus.COMPLETED,
                output="Option B",
            ),
        ]

        aggregated = pool.aggregate_results(results)

        assert "Voting Results:" in aggregated
        assert "Total votes: 3" in aggregated
        assert "Winning result:" in aggregated
        assert "Option A" in aggregated

    def test_aggregate_merge(self, mock_provider: MagicMock) -> None:
        """Test merge aggregation."""
        pool = SubAgentPool(
            provider=mock_provider,
            aggregation_config=AggregationConfig(strategy="merge"),
        )

        results = [
            SubAgentResult(
                agent_id="a1",
                task="t1",
                status=SubAgentStatus.COMPLETED,
                output="Unique result 1",
            ),
            SubAgentResult(
                agent_id="a2",
                task="t2",
                status=SubAgentStatus.COMPLETED,
                output="Unique result 1",  # Duplicate
            ),
            SubAgentResult(
                agent_id="a3",
                task="t3",
                status=SubAgentStatus.COMPLETED,
                output="Unique result 2",
            ),
        ]

        aggregated = pool.aggregate_results(results)

        assert "2 unique from 3 agents" in aggregated

    def test_get_stats(self, mock_provider: MagicMock) -> None:
        """Test getting pool statistics."""
        pool = SubAgentPool(provider=mock_provider)
        pool._results = [
            SubAgentResult(
                agent_id="a1",
                task="t1",
                status=SubAgentStatus.COMPLETED,
                tokens_used=100,
                start_time=datetime(2024, 1, 1, 12, 0, 0),
                end_time=datetime(2024, 1, 1, 12, 0, 5),
            ),
            SubAgentResult(
                agent_id="a2",
                task="t2",
                status=SubAgentStatus.FAILED,
                tokens_used=50,
                start_time=datetime(2024, 1, 1, 12, 0, 0),
                end_time=datetime(2024, 1, 1, 12, 0, 3),
            ),
        ]

        stats = pool.get_stats()

        assert stats["total_agents"] == 2
        assert stats["successful"] == 1
        assert stats["failed"] == 1
        assert stats["success_rate"] == 0.5
        assert stats["total_tokens"] == 150
        assert stats["total_duration"] == 8.0
        assert stats["avg_duration"] == 4.0


class TestSubAgentManager:
    """Test SubAgentManager class."""

    @pytest.fixture
    def mock_provider(self) -> MagicMock:
        """Create a mock provider."""
        provider = MagicMock()
        provider.complete = AsyncMock()
        return provider

    def test_init(self, mock_provider: MagicMock) -> None:
        """Test manager initialization."""
        manager = SubAgentManager(mock_provider, default_max_concurrent=8)

        assert manager.provider == mock_provider
        assert manager.default_max_concurrent == 8

    @pytest.mark.asyncio
    async def test_parallel_map(self, mock_provider: MagicMock) -> None:
        """Test parallel map execution."""
        mock_response = CompletionResponse(
            content="done",
            model="test",
            usage=UsageStats(10, 10, 20),
        )
        mock_provider.complete.return_value = mock_response

        manager = SubAgentManager(mock_provider)

        tasks = ["task1", "task2"]
        results, aggregated = await manager.parallel_map(
            tasks=tasks,
            aggregation_strategy="concatenate",
        )

        assert len(results) == 2
        assert "done" in aggregated

    @pytest.mark.asyncio
    async def test_parallel_vote(self, mock_provider: MagicMock) -> None:
        """Test parallel vote execution."""
        mock_response = CompletionResponse(
            content="Answer: 42",
            model="test",
            usage=UsageStats(10, 10, 20),
        )
        mock_provider.complete.return_value = mock_response

        manager = SubAgentManager(mock_provider)

        results, voted = await manager.parallel_vote(
            task="What is the answer?",
            num_agents=3,
        )

        assert len(results) == 3
        assert "Voting Results:" in voted
        assert "Answer: 42" in voted

    @pytest.mark.asyncio
    async def test_divide_and_conquer(self, mock_provider: MagicMock) -> None:
        """Test divide and conquer execution."""
        mock_response = CompletionResponse(
            content="Partial result",
            model="test",
            usage=UsageStats(10, 10, 20),
        )
        mock_provider.complete.return_value = mock_response

        manager = SubAgentManager(mock_provider)

        subtasks = ["subtask1", "subtask2"]
        results, merged = await manager.divide_and_conquer(
            subtasks=subtasks,
        )

        assert len(results) == 2
        assert "unique from 2 agents" in merged


class TestAggregationConfig:
    """Test AggregationConfig dataclass."""

    def test_defaults(self) -> None:
        """Test default values."""
        config = AggregationConfig()

        assert config.strategy == "concatenate"
        assert config.separator == "\n\n---\n\n"
        assert config.include_metadata is False
        assert config.max_output_length == 10000

    def test_custom_values(self) -> None:
        """Test custom values."""
        config = AggregationConfig(
            strategy="vote",
            separator=" | ",
            include_metadata=True,
            max_output_length=5000,
        )

        assert config.strategy == "vote"
        assert config.separator == " | "
        assert config.include_metadata is True
        assert config.max_output_length == 5000
