"""
Tests for Group module.
"""

import pytest
from unittest.mock import Mock

from bgx.group import TaskGroup, GroupPattern, TaskDependency, GroupNode
from bgx.task import Task, TaskState, Priority, RetryPolicy


class TestTaskDependency:
    """Test TaskDependency functionality."""
    
    def test_task_dependency_creation(self):
        """Test creating task dependencies."""
        dep = TaskDependency(task_id="task1", depends_on=["task2"])
        
        assert dep.task_id == "task1"
        assert dep.depends_on == ["task2"]
        assert dep.condition is None
        
        # With condition
        def condition(result):
            return result > 0
        
        dep_with_condition = TaskDependency(
            task_id="task1", 
            depends_on=["task2"], 
            condition=condition
        )
        
        assert dep_with_condition.condition == condition


class TestGroupNode:
    """Test GroupNode functionality."""
    
    def test_group_node_creation(self):
        """Test creating group nodes."""
        def sample_func():
            return "result"
        
        task = Task(func=sample_func)
        
        node = GroupNode(
            task_id="task1",
            task=task,
            dependencies={"task2"},
            dependents={"task3"},
            in_degree=1,
            out_degree=2
        )
        
        assert node.task_id == "task1"
        assert node.task == task
        assert node.dependencies == {"task2"}
        assert node.dependents == {"task3"}
        assert node.in_degree == 1
        assert node.out_degree == 2


class TestTaskGroup:
    """Test TaskGroup functionality."""
    
    def test_task_group_creation(self):
        """Test creating task groups."""
        engine = Mock()
        
        group = TaskGroup(engine)
        
        assert group.engine == engine
        assert len(group._nodes) == 0
        assert len(group._edges) == 0
        assert len(group._completed_tasks) == 0
        assert len(group._failed_tasks) == 0
    
    def test_add_task(self):
        """Test adding tasks to group."""
        engine = Mock()
        group = TaskGroup(engine)
        
        def sample_func():
            return "result"
        
        task = Task(func=sample_func)
        
        # Add task without dependencies
        task_id = group.add_task(task)
        assert task_id == task.id
        assert len(group._nodes) == 1
        assert group._nodes[task_id].task == task
        assert group._nodes[task_id].dependencies == set()
        
        # Add task with dependencies
        task2 = Task(func=sample_func)
        task_id2 = group.add_task(task2, dependencies=[task_id])
        
        assert task_id2 == task2.id
        assert len(group._nodes) == 2
        assert group._nodes[task_id2].dependencies == {task_id}
        assert group._nodes[task_id].dependents == {task_id2}
        assert group._nodes[task_id].in_degree == 0  # First task has no dependencies
        assert group._nodes[task_id2].in_degree == 1  # Second task has one dependency
        
        # Test adding duplicate task
        with pytest.raises(ValueError):
            group.add_task(task)
    
    def test_add_dependency(self):
        """Test adding dependencies between tasks."""
        engine = Mock()
        group = TaskGroup(engine)
        
        def sample_func():
            return "result"
        
        task1 = Task(func=sample_func)
        task2 = Task(func=sample_func)
        
        # Add tasks
        task_id1 = group.add_task(task1)
        task_id2 = group.add_task(task2)
        
        # Add dependency
        group.add_dependency(task_id1, task_id2)
        
        assert task_id2 in group._nodes[task_id1].dependencies
        assert task_id1 in group._nodes[task_id2].dependents
        assert group._nodes[task_id1].in_degree == 1
        assert group._nodes[task_id2].in_degree == 0
        assert group._nodes[task_id2].out_degree == 1
    
    def test_build_execution_plan(self):
        """Test building execution plan."""
        engine = Mock()
        group = TaskGroup(engine)
        
        def sample_func():
            return "result"
        
        # Create linear chain: A -> B -> C
        task_a = Task(func=sample_func)
        task_b = Task(func=sample_func)
        task_c = Task(func=sample_func)
        
        group.add_task(task_a)
        group.add_task(task_b, dependencies=[task_a.id])
        group.add_task(task_c, dependencies=[task_b.id])
        
        plan = group._build_execution_plan()
        
        # Should execute in order: A, B, C
        assert plan == [task_a.id, task_b.id, task_c.id]
    
    def test_execute_sequential(self):
        """Test sequential execution."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=2) as engine:
            group = TaskGroup(engine)
            
            def sample_func(x):
                return f"result_{x}"
            
            # Create tasks
            task_a = Task(func=sample_func, args=["a"])
            task_b = Task(func=sample_func, args=["b"])
            task_c = Task(func=sample_func, args=["c"])
            
            group.add_task(task_a)
            group.add_task(task_b, dependencies=[task_a.id])
            group.add_task(task_c, dependencies=[task_b.id])
            
            # Execute sequential
            results = group.execute_sequential([task_a.id, task_b.id, task_c.id])
            
            assert results[task_a.id] == "result_a"
            assert results[task_b.id] == "result_b"
            assert results[task_c.id] == "result_c"
            
            # Verify all tasks completed
            assert group.get_task_status(task_a.id) == TaskState.SUCCESS.value
            assert group.get_task_status(task_b.id) == TaskState.SUCCESS.value
            assert group.get_task_status(task_c.id) == TaskState.SUCCESS.value
    
    def test_execute_fan_out(self):
        """Test fan-out execution."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=4) as engine:
            group = TaskGroup(engine)
            
            def sample_func(x):
                return f"result_{x}"
            
            # Create tasks
            main_task = Task(func=sample_func, args=["main"])
            dep_a = Task(func=sample_func, args=["a"])
            dep_b = Task(func=sample_func, args=["b"])
            dep_c = Task(func=sample_func, args=["c"])
            
            group.add_task(main_task)
            group.add_task(dep_a, dependencies=[main_task.id])
            group.add_task(dep_b, dependencies=[main_task.id])
            group.add_task(dep_c, dependencies=[main_task.id])
            
            # Execute fan-out
            results = group.execute_fan_out(main_task.id, [dep_a.id, dep_b.id, dep_c.id])
            
            assert results[main_task.id] == "result_main"
            assert results[dep_a.id] == "result_a"
            assert results[dep_b.id] == "result_b"
            assert results[dep_c.id] == "result_c"
            
            # Verify all tasks completed
            assert group.get_task_status(main_task.id) == TaskState.SUCCESS.value
            assert group.get_task_status(dep_a.id) == TaskState.SUCCESS.value
            assert group.get_task_status(dep_b.id) == TaskState.SUCCESS.value
            assert group.get_task_status(dep_c.id) == TaskState.SUCCESS.value
    
    def test_execute_fan_in(self):
        """Test fan-in execution."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=4) as engine:
            group = TaskGroup(engine)
            
            def sample_func(x):
                return f"result_{x}"
            
            # Create tasks
            task_a = Task(func=sample_func, args=["a"])
            task_b = Task(func=sample_func, args=["b"])
            
            group.add_task(task_a)
            group.add_task(task_b)
            
            # Execute fan-in
            results = group.execute_fan_in([task_a.id, task_b.id])
            
            assert results[task_a.id] == "result_a"
            assert results[task_b.id] == "result_b"
            
            # Verify all tasks completed
            assert group.get_task_status(task_a.id) == TaskState.SUCCESS.value
            assert group.get_task_status(task_b.id) == TaskState.SUCCESS.value
    
    def test_execute_diamond(self):
        """Test diamond execution."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=4) as engine:
            group = TaskGroup(engine)
            
            def sample_func(x):
                return f"result_{x}"
            
            # Create tasks
            start_task = Task(func=sample_func, args=["start"])
            parallel_a = Task(func=sample_func, args=["a"])
            parallel_b = Task(func=sample_func, args=["b"])
            end_task = Task(func=sample_func, args=["end"])
            
            group.add_task(start_task)
            group.add_task(parallel_a, dependencies=[start_task.id])
            group.add_task(parallel_b, dependencies=[start_task.id])
            group.add_task(end_task, dependencies=[parallel_a.id, parallel_b.id])
            
            # Execute diamond
            results = group.execute_diamond(
                start_task.id, 
                [parallel_a.id, parallel_b.id], 
                end_task.id
            )
            
            assert results[start_task.id] == "result_start"
            assert results[parallel_a.id] == "result_a"
            assert results[parallel_b.id] == "result_b"
            assert results[end_task.id] == "result_end"
            
            # Verify all tasks completed
            assert group.get_task_status(start_task.id) == TaskState.SUCCESS.value
            assert group.get_task_status(parallel_a.id) == TaskState.SUCCESS.value
            assert group.get_task_status(parallel_b.id) == TaskState.SUCCESS.value
            assert group.get_task_status(end_task.id) == TaskState.SUCCESS.value
    
    def test_execute_with_pattern(self):
        """Test execution with different patterns."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=4) as engine:
            group = TaskGroup(engine)
            
            def sample_func(x):
                return f"result_{x}"
            
            # Create tasks
            task_a = Task(func=sample_func, args=["a"])
            task_b = Task(func=sample_func, args=["b"])
            task_c = Task(func=sample_func, args=["c"])
            
            group.add_task(task_a)
            group.add_task(task_b, dependencies=[task_a.id])
            group.add_task(task_c, dependencies=[task_b.id])
            
            # Test sequential pattern
            results = group.execute(GroupPattern.SEQUENTIAL, [task_a.id, task_b.id, task_c.id])
            assert results[task_a.id] == "result_a"
            assert results[task_b.id] == "result_b"
            assert results[task_c.id] == "result_c"
    
    def test_task_status(self):
        """Test getting task status."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=2) as engine:
            group = TaskGroup(engine)
            
            def sample_func():
                return "result"
            
            task = Task(func=sample_func)
            task_id = group.add_task(task)
            
            # Task not executed yet
            assert group.get_task_status(task_id) == "pending"
            
            # Execute task
            group._execute_task(task_id)
            
            assert group.get_task_status(task_id) == TaskState.SUCCESS.value
            
            # Non-existent task
            assert group.get_task_status("non_existent") is None
    
    def test_group_status(self):
        """Test getting group status."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=2) as engine:
            group = TaskGroup(engine)
            
            def sample_func():
                return "result"
            
            task_a = Task(func=sample_func)
            task_b = Task(func=sample_func)
            
            group.add_task(task_a)
            group.add_task(task_b, dependencies=[task_a.id])
            
            status = group.get_group_status()
            
            assert len(status) == 2
            assert status[task_a.id] == "pending"
            assert status[task_b.id] == "pending"
            
            # Execute one task
            group._execute_task(task_a.id)
            
            status = group.get_group_status()
            assert status[task_a.id] == TaskState.SUCCESS.value
            assert status[task_b.id] == "pending"
    
    def test_clear_completed(self):
        """Test clearing completed tasks."""
        from bgx import BackgroundEngine
        with BackgroundEngine(workers=2) as engine:
            group = TaskGroup(engine)
            
            def sample_func():
                return "result"
            
            task_a = Task(func=sample_func)
            task_b = Task(func=sample_func)
            
            group.add_task(task_a)
            group.add_task(task_b, dependencies=[task_a.id])
            
            # Execute tasks
            group._execute_task(task_a.id)
            group._execute_task(task_b.id)
            
            assert len(group._completed_tasks) == 2
            assert group.get_task_status(task_a.id) == TaskState.SUCCESS.value
            assert group.get_task_status(task_b.id) == TaskState.SUCCESS.value
            
            # Clear completed
            group.clear_completed()
            
            assert len(group._completed_tasks) == 0
            assert len(group._failed_tasks) == 0
    
    def test_reset(self):
        """Test resetting group."""
        engine = Mock()
        group = TaskGroup(engine)
        
        def sample_func():
            return "result"
        
        task = Task(func=sample_func)
        group.add_task(task)
        
        assert len(group._nodes) == 1
        
        group.reset()
        
        assert len(group._nodes) == 0
        assert len(group._edges) == 0
        assert len(group._completed_tasks) == 0
        assert len(group._failed_tasks) == 0
