"""Unit tests for relay 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 RelayRunner


class TestRelayRunner:
    """Tests for RelayRunner."""

    def test_basic_execution(self):
        """Test basic execution of relay runner."""
        items = [1, 2, 3, 4, 5]
        results = []
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                results.append((value, index))

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, items, handler)
            runner.run()

        # All items should be processed
        assert len(results) == 5
        # Indices should be correct
        values_by_index = sorted(results, key=lambda x: x[1])
        assert [v for v, _ in values_by_index] == items

    def test_empty_iterable(self):
        """Test relay runner with empty iterable."""
        results = []

        def handler(value, index):
            results.append((value, index))

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [], handler)
            runner.run()

        assert len(results) == 0

    def test_single_item(self):
        """Test relay runner with single item."""
        results = []

        def handler(value, index):
            results.append((value, index))

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [42], handler)
            runner.run()

        assert results == [(42, 0)]

    def test_handler_receives_correct_indices(self):
        """Test handler receives correct indices."""
        items = ['a', 'b', 'c', 'd', 'e']
        indices = []
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                indices.append(index)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, items, handler)
            runner.run()

        # Indices should be 0, 1, 2, 3, 4
        assert sorted(indices) == [0, 1, 2, 3, 4]

    def test_error_in_handler_stops_processing(self):
        """Test error in handler stops processing."""

        def handler(value, index):
            if value == 3:
                raise ValueError('Error at 3')
            # Simulate some work
            time.sleep(0.01)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [1, 2, 3, 4, 5], handler)
            with pytest.raises(ValueError, match='Error at 3'):
                runner.run()

    def test_first_error_is_raised(self):
        """Test first error is raised when multiple errors occur."""
        errors_raised = []
        lock = threading.Lock()

        def handler(value, index):
            # Introduce delay to control timing
            time.sleep(0.001 * index)
            with lock:
                errors_raised.append(value)
            raise ValueError(f'Error at {value}')

        with ThreadPoolExecutor(max_workers=5) as executor:
            runner = RelayRunner(executor, [1, 2, 3, 4, 5], handler)
            with pytest.raises(ValueError):
                runner.run()

        # At least one error should have been raised
        assert len(errors_raised) > 0

    def test_runner_cannot_be_reused(self):
        """Test runner raises error if reused."""

        def handler(value, index):
            pass

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [1, 2, 3], handler)
            runner.run()

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

    def test_concurrent_execution(self):
        """Test tasks execute concurrently."""
        start_times = {}
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                start_times[value] = time.time()
            # Simulate work
            time.sleep(0.05)

        with ThreadPoolExecutor(max_workers=3) as executor:
            start = time.time()
            runner = RelayRunner(executor, [1, 2, 3, 4, 5], handler)
            runner.run()
            elapsed = time.time() - start

        # With 3 workers and 5 tasks of 0.05s each, should complete in ~0.1s
        # If sequential, would take 0.25s
        # Allow some margin for thread overhead
        assert elapsed < 0.20, f'Took {elapsed}s, expected < 0.20s'

    def test_task_id_context(self):
        """Test task ID context is preserved."""
        from abacus_common_logic.utils.logging import task_id_var

        task_id_var.set('test-task-id')

        captured_task_ids = []
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                captured_task_ids.append(task_id_var.get())

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [1, 2, 3], handler)
            runner.run()

        # All handlers should have the same task ID
        assert all(tid == 'test-task-id' for tid in captured_task_ids)
        assert len(captured_task_ids) == 3

    def test_with_generator(self):
        """Test relay runner works with generators."""

        def generator():
            for i in range(5):
                yield i * 2

        results = []
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                results.append(value)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, generator(), handler)
            runner.run()

        assert sorted(results) == [0, 2, 4, 6, 8]

    def test_handler_modifications_dont_affect_iteration(self):
        """Test handler can't affect iteration by modifying values."""
        items = [1, 2, 3, 4, 5]
        processed = []
        lock = threading.Lock()

        def handler(value, index):
            # Try to modify value (shouldn't affect other iterations)
            value = value * 100
            with lock:
                processed.append(value)

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, items, handler)
            runner.run()

        # Original iteration values should be processed
        assert sorted(processed) == [100, 200, 300, 400, 500]

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

        def handler(value, index):
            pass

        with ThreadPoolExecutor(max_workers=2) as executor:
            runner = RelayRunner(executor, [1, 2, 3], handler)
            runner.run()

        # Reset should be called for each task (3 items + 1 sentinel task for StopIteration)
        assert mock_task_id_var.reset.call_count == 4

    def test_large_batch(self):
        """Test relay runner with large batch."""
        items = list(range(1000))
        count = [0]
        lock = threading.Lock()

        def handler(value, index):
            with lock:
                count[0] += 1

        with ThreadPoolExecutor(max_workers=10) as executor:
            runner = RelayRunner(executor, items, handler)
            runner.run()

        assert count[0] == 1000

    def test_counter_accuracy(self):
        """Test internal counter tracks progress accurately."""

        def handler(value, index):
            time.sleep(0.001)

        with ThreadPoolExecutor(max_workers=5) as executor:
            runner = RelayRunner(executor, list(range(10)), handler)
            runner.run()

        # After completion, counter should equal length
        # Length is items + 1 (includes sentinel task that gets StopIteration)
        assert runner._counter == runner._length
        assert runner._length == 11  # 10 items + 1 sentinel task