"""Unit tests for task graph runner module."""

import threading
import time
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import Mock, patch

import pytest

from abacus_common_logic.concurrent import (
    Task,
    TaskGraph,
    TaskGraphError,
    TaskGraphRunner,
    TaskOutput,
    TaskStatus,
)


class TestTaskGraphRunner:
    """Tests for TaskGraphRunner."""

    def test_initialization_with_single_pool(self):
        """Test runner initializes with single pool."""

        def handler():
            pass

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            assert runner.graph == graph
            assert runner._default_pool == 'default'

    def test_initialization_with_multiple_pools(self):
        """Test runner initializes with multiple pools."""

        def handler():
            pass

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with (
            ThreadPoolExecutor(max_workers=2) as exec1,
            ThreadPoolExecutor(max_workers=4) as exec2,
        ):
            runner = TaskGraphRunner(
                graph, pools={'cpu': exec1, 'io': exec2}, default_pool='cpu'
            )
            assert runner._default_pool == 'cpu'

    def test_initialization_without_pools_raises_error(self):
        """Test initialization without pools raises ValueError."""

        def handler():
            pass

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with pytest.raises(ValueError, match='At least one pool must be provided'):
            TaskGraphRunner(graph, pools={})

    def test_default_pool_selection(self):
        """Test default pool is selected when not specified."""

        def handler():
            pass

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'only_pool': executor})
            assert runner._default_pool == 'only_pool'

    def test_single_task_execution(self):
        """Test execution of single task."""
        result_value = [0]

        def handler():
            result_value[0] = 42
            return result_value[0]

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['task1'] == 42
        assert result_value[0] == 42

    def test_linear_task_chain(self):
        """Test execution of linear task chain."""
        counter = [0]
        lock = threading.Lock()

        def make_handler(task_id, expected_min):
            def handler():
                # Verify dependencies have completed by checking counter
                with lock:
                    assert counter[0] >= expected_min, (
                        f'{task_id} ran too early (counter={counter[0]}, '
                        f'expected >= {expected_min})'
                    )
                    counter[0] += 1
                return task_id

            return handler

        tasks = [
            Task(id='task1', handler=make_handler('task1', 0)),
            Task(id='task2', handler=make_handler('task2', 1), depends_on={'task1'}),
            Task(id='task3', handler=make_handler('task3', 2), depends_on={'task2'}),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        # All tasks should complete and results should be stored
        assert results == {'task1': 'task1', 'task2': 'task2', 'task3': 'task3'}
        # All tasks should have completed
        assert len(runner._tasks_complete) == 3

    def test_parallel_independent_tasks(self):
        """Test parallel execution of independent tasks."""
        start_times = {}
        lock = threading.Lock()

        def make_handler(task_id):
            def handler():
                with lock:
                    start_times[task_id] = time.time()
                time.sleep(0.05)
                return task_id

            return handler

        tasks = [
            Task(id='task1', handler=make_handler('task1')),
            Task(id='task2', handler=make_handler('task2')),
            Task(id='task3', handler=make_handler('task3')),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=3) as executor:
            start = time.time()
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()
            elapsed = time.time() - start

        # All tasks should complete
        assert len(results) == 3
        # Should execute in parallel (total time < sum of individual times)
        assert elapsed < 0.15, f'Took {elapsed}s, expected < 0.15s'

    def test_diamond_dependency_pattern(self):
        """Test diamond dependency execution."""
        execution_order = []
        lock = threading.Lock()

        def make_handler(task_id):
            def handler():
                with lock:
                    execution_order.append(task_id)
                time.sleep(0.01)
                return task_id

            return handler

        tasks = [
            Task(id='start', handler=make_handler('start')),
            Task(id='left', handler=make_handler('left'), depends_on={'start'}),
            Task(id='right', handler=make_handler('right'), depends_on={'start'}),
            Task(id='end', handler=make_handler('end'), depends_on={'left', 'right'}),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=3) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        # All tasks should execute
        assert len(execution_order) == 4
        # start must execute before left and right
        assert execution_order.index('start') < execution_order.index('left')
        assert execution_order.index('start') < execution_order.index('right')
        # end must execute after both left and right
        assert execution_order.index('left') < execution_order.index('end')
        assert execution_order.index('right') < execution_order.index('end')
        assert len(results) == 4

    def test_task_with_parameters(self):
        """Test task execution with parameters."""

        def add(x, y):
            return x + y

        tasks = [Task(id='add_task', handler=add, params={'x': 10, 'y': 20})]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['add_task'] == 30

    def test_task_output_resolution(self):
        """Test TaskOutput placeholder resolution."""

        def fetch():
            return {'data': [1, 2, 3]}

        def process(input_data):
            return sum(input_data['data'])

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'input_data': TaskOutput('fetch')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['fetch'] == {'data': [1, 2, 3]}
        assert results['process'] == 6

    def test_task_output_with_path_dict(self):
        """Test TaskOutput with path for dictionary access."""

        def fetch():
            return {'results': {'count': 42, 'items': ['a', 'b', 'c']}}

        def process(count):
            return count * 2

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'count': TaskOutput('fetch', path='results.count')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 84

    def test_task_output_with_path_list(self):
        """Test TaskOutput with path for list access."""

        def fetch():
            return {'items': [10, 20, 30]}

        def process(value):
            return value + 5

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'value': TaskOutput('fetch', path='items.1')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 25

    def test_task_output_with_path_nested(self):
        """Test TaskOutput with deeply nested path."""

        def fetch():
            return {'level1': {'level2': {'level3': {'value': 100}}}}

        def process(value):
            return value / 2

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={
                    'value': TaskOutput('fetch', path='level1.level2.level3.value')
                },
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 50

    def test_task_output_with_path_object_attribute(self):
        """Test TaskOutput with path for object attribute access."""

        class Result:
            def __init__(self):
                self.status = 'success'
                self.data = {'count': 5}

        def fetch():
            return Result()

        def process(status):
            return status.upper()

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'status': TaskOutput('fetch', path='status')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 'SUCCESS'

    def test_task_failure_stops_workflow(self):
        """Test task failure stops workflow execution."""
        executed = []
        lock = threading.Lock()

        def make_handler(task_id, should_fail=False):
            def handler():
                with lock:
                    executed.append(task_id)
                if should_fail:
                    raise ValueError(f'Task {task_id} failed')
                return task_id

            return handler

        tasks = [
            Task(id='task1', handler=make_handler('task1')),
            Task(
                id='task2',
                handler=make_handler('task2', should_fail=True),
                depends_on={'task1'},
            ),
            Task(id='task3', handler=make_handler('task3'), depends_on={'task2'}),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(RuntimeError, match='Workflow failed'):
                runner.run()

        # task1 and task2 should execute, task3 should be skipped
        assert 'task1' in executed
        assert 'task2' in executed
        assert 'task3' not in executed

    def test_first_error_is_preserved(self):
        """Test first error is preserved when multiple tasks fail."""

        def make_handler(task_id):
            def handler():
                time.sleep(0.01 * int(task_id[-1]))
                raise ValueError(f'Error from {task_id}')

            return handler

        tasks = [
            Task(id='task1', handler=make_handler('task1')),
            Task(id='task2', handler=make_handler('task2')),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(RuntimeError) as exc_info:
                runner.run()

        # Should be a TaskGraphError
        assert isinstance(exc_info.value.__cause__, TaskGraphError)

    def test_task_graph_error_wrapping(self):
        """Test exceptions are wrapped in TaskGraphError."""

        def failing_task():
            raise ValueError('Something went wrong')

        tasks = [Task(id='failing', handler=failing_task)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(RuntimeError) as exc_info:
                runner.run()

        error = exc_info.value.__cause__
        assert isinstance(error, TaskGraphError)
        assert error.task_id == 'failing'
        assert 'failing' in str(error)

    def test_downstream_tasks_skipped_after_failure(self):
        """Test downstream tasks are skipped after failure."""

        def failing_task():
            raise ValueError('Failure')

        def skipped_task():
            pytest.fail('This task should not execute')

        tasks = [
            Task(id='fail', handler=failing_task),
            Task(id='skip', handler=skipped_task, depends_on={'fail'}),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(RuntimeError):
                runner.run()

        # Check that fail task is in the failed set
        assert 'fail' in runner._tasks_failed
        # Skip task should still be in pending (never submitted after error)
        assert 'skip' in runner._tasks_pending
        # Skip task should not be in completed or running
        assert 'skip' not in runner._tasks_complete
        assert 'skip' not in runner._tasks_running

    def test_workflow_timeout(self):
        """Test workflow timeout."""

        def slow_task():
            time.sleep(0.25)
            return 'done'

        tasks = [Task(id='slow', handler=slow_task)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(TimeoutError, match='Workflow timed out'):
                runner.run(timeout=0.1)

    def test_cannot_reuse_runner(self):
        """Test runner cannot be reused."""

        def handler():
            return 42

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            runner.run()

            with pytest.raises(RuntimeError, match='Runner cannot be reused'):
                runner.run()

    def test_task_specific_pool(self):
        """Test task executes in specified pool."""
        pool_used = []
        lock = threading.Lock()

        def make_handler(pool_name):
            def handler(executor):
                # Check which pool we're running in
                with lock:
                    pool_used.append(pool_name)
                return pool_name

            return handler

        tasks = [
            Task(id='cpu_task', handler=make_handler('cpu'), pool='cpu'),
            Task(id='io_task', handler=make_handler('io'), pool='io'),
        ]
        graph = TaskGraph(tasks)

        with (
            ThreadPoolExecutor(max_workers=2) as cpu_exec,
            ThreadPoolExecutor(max_workers=4) as io_exec,
        ):
            runner = TaskGraphRunner(graph, pools={'cpu': cpu_exec, 'io': io_exec})
            results = runner.run()

        assert len(pool_used) == 2
        assert 'cpu' in pool_used
        assert 'io' in pool_used

    def test_invalid_pool_raises_error(self):
        """Test invalid pool reference raises error."""

        def handler():
            return 42

        tasks = [Task(id='task1', handler=handler, pool='nonexistent')]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(ValueError, match='nonexistent') as exc_info:
                runner.run()

    def test_executor_parameter_injection(self):
        """Test executor parameter is injected when handler expects it."""
        executor_captured = [None]

        def handler(executor):
            executor_captured[0] = executor
            return 'done'

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            runner.run()

        # Executor should be injected
        assert executor_captured[0] is executor

    def test_executor_parameter_not_overridden(self):
        """Test executor parameter is not injected if already provided."""
        custom_executor = Mock()

        def handler(executor):
            return executor

        tasks = [
            Task(id='task1', handler=handler, params={'executor': custom_executor})
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as pool_executor:
            runner = TaskGraphRunner(graph, pools={'default': pool_executor})
            results = runner.run()

        # Should use the custom executor from params
        assert results['task1'] is custom_executor

    def test_task_id_context_propagation(self):
        """Test task ID context is propagated to handlers."""
        from abacus_common_logic.utils.logging import task_id_var

        captured_task_ids = {}
        lock = threading.Lock()

        def make_handler(task_id):
            def handler():
                with lock:
                    captured_task_ids[task_id] = task_id_var.get()
                return task_id

            return handler

        tasks = [
            Task(id='task1', handler=make_handler('task1')),
            Task(id='task2', handler=make_handler('task2')),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            runner.run()

        # Each task should have its own ID in context
        assert captured_task_ids['task1'] == 'task1'
        assert captured_task_ids['task2'] == 'task2'

    def test_results_stored_correctly(self):
        """Test task results are stored correctly."""

        def task1():
            return {'result': 'from_task1'}

        def task2():
            return [1, 2, 3]

        def task3():
            return 42

        tasks = [
            Task(id='task1', handler=task1),
            Task(id='task2', handler=task2),
            Task(id='task3', handler=task3),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['task1'] == {'result': 'from_task1'}
        assert results['task2'] == [1, 2, 3]
        assert results['task3'] == 42

    def test_empty_graph(self):
        """Test runner with empty graph."""
        graph = TaskGraph([])

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results == {}

    def test_complex_graph_execution(self):
        """Test execution of complex graph with multiple paths."""
        execution_order = []
        lock = threading.Lock()

        def make_handler(task_id):
            def handler():
                with lock:
                    execution_order.append(task_id)
                time.sleep(0.01)
                return task_id

            return handler

        tasks = [
            Task(id='a', handler=make_handler('a')),
            Task(id='b', handler=make_handler('b'), depends_on={'a'}),
            Task(id='c', handler=make_handler('c'), depends_on={'a'}),
            Task(id='d', handler=make_handler('d'), depends_on={'b', 'c'}),
            Task(id='e', handler=make_handler('e'), depends_on={'b'}),
            Task(id='f', handler=make_handler('f'), depends_on={'d', 'e'}),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=3) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        # All tasks should execute
        assert len(execution_order) == 6
        # Verify dependency constraints
        assert execution_order.index('a') < execution_order.index('b')
        assert execution_order.index('a') < execution_order.index('c')
        assert execution_order.index('b') < execution_order.index('d')
        assert execution_order.index('c') < execution_order.index('d')
        assert execution_order.index('b') < execution_order.index('e')
        assert execution_order.index('d') < execution_order.index('f')
        assert execution_order.index('e') < execution_order.index('f')
        assert len(results) == 6

    def test_multiple_task_outputs_in_params(self):
        """Test multiple TaskOutput references in params."""

        def task1():
            return 10

        def task2():
            return 20

        def task3(x, y):
            return x + y

        tasks = [
            Task(id='task1', handler=task1),
            Task(id='task2', handler=task2),
            Task(
                id='task3',
                handler=task3,
                depends_on={'task1', 'task2'},
                params={'x': TaskOutput('task1'), 'y': TaskOutput('task2')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['task3'] == 30

    def test_task_output_mixed_with_regular_params(self):
        """Test TaskOutput mixed with regular parameters."""

        def fetch():
            return 5

        def process(multiplier, value):
            return multiplier * value

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'multiplier': 3, 'value': TaskOutput('fetch')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 15

    def test_cleanup_after_execution(self):
        """Test cleanup is performed after execution."""

        def handler():
            return 42

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            runner.run()

        # in_degrees should be cleared
        assert len(runner._in_degrees) == 0

    def test_cleanup_after_failure(self):
        """Test cleanup is performed even after failure."""

        def failing_task():
            raise ValueError('Failure')

        tasks = [Task(id='fail', handler=failing_task)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            with pytest.raises(RuntimeError):
                runner.run()

        # in_degrees should be cleared even after failure
        assert len(runner._in_degrees) == 0

    @patch('abacus_common_logic.concurrent.task_graph_runner.task_id_var')
    def test_task_id_context_cleanup(self, mock_task_id_var):
        """Test task ID context is properly cleaned up."""
        mock_token = Mock()
        mock_task_id_var.set.return_value = mock_token

        def handler():
            return 42

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            runner.run()

        # set should be called once
        mock_task_id_var.set.assert_called_once_with('task1')
        # reset should be called once
        mock_task_id_var.reset.assert_called_once_with(mock_token)

    def test_task_status_enum(self):
        """Test TaskStatus enum values."""
        assert TaskStatus.SUCCESS
        assert TaskStatus.FAILED
        assert TaskStatus.SKIPPED

    def test_parallel_execution_with_different_durations(self):
        """Test parallel execution with tasks of different durations."""
        execution_times = {}
        lock = threading.Lock()

        def make_handler(task_id, duration):
            def handler():
                start = time.time()
                time.sleep(duration)
                with lock:
                    execution_times[task_id] = time.time() - start
                return task_id

            return handler

        tasks = [
            Task(id='fast', handler=make_handler('fast', 0.01)),
            Task(id='medium', handler=make_handler('medium', 0.05)),
            Task(id='slow', handler=make_handler('slow', 0.1)),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=3) as executor:
            start = time.time()
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()
            elapsed = time.time() - start

        # All should execute in parallel
        assert len(results) == 3
        # Total time should be close to slowest task (0.1s), not sum (0.16s)
        assert elapsed < 0.15, f'Took {elapsed}s, expected < 0.15s'

    def test_none_return_value(self):
        """Test task with None return value."""

        def handler():
            return None

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['task1'] is None

    def test_task_with_no_parameters(self):
        """Test task handler with no parameters."""

        def handler():
            return 'no params'

        tasks = [Task(id='task1', handler=handler)]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['task1'] == 'no params'

    def test_resolve_path_with_tuple(self):
        """Test path resolution with tuple."""

        def fetch():
            return {'items': (1, 2, 3)}

        def process(value):
            return value * 2

        tasks = [
            Task(id='fetch', handler=fetch),
            Task(
                id='process',
                handler=process,
                depends_on={'fetch'},
                params={'value': TaskOutput('fetch', path='items.2')},
            ),
        ]
        graph = TaskGraph(tasks)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = TaskGraphRunner(graph, pools={'default': executor})
            results = runner.run()

        assert results['process'] == 6