import sys
from time import monotonic
from unittest.mock import Mock, patch

import pytest

from celery.concurrency.base import BasePool
from celery.utils.objects import Bunch
from celery.worker import autoscale, state


class MockPool(BasePool):

    shrink_raises_exception = False
    shrink_raises_ValueError = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._pool = Bunch(_processes=self.limit)

    def grow(self, n=1):
        self._pool._processes += n

    def shrink(self, n=1):
        if self.shrink_raises_exception:
            raise KeyError('foo')
        if self.shrink_raises_ValueError:
            raise ValueError('foo')
        self._pool._processes -= n

    @property
    def num_processes(self):
        return self._pool._processes


class test_WorkerComponent:

    def test_register_with_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.consumer.on_task_message = set()
        w = autoscale.WorkerComponent(parent)
        assert parent.autoscaler is None
        assert w.enabled

        hub = Mock(name='hub')
        w.create(parent)
        w.register_with_event_loop(parent, hub)
        assert (parent.autoscaler.maybe_scale in
                parent.consumer.on_task_message)
        hub.call_repeatedly.assert_called_with(
            parent.autoscaler.keepalive, parent.autoscaler.maybe_scale,
        )

        parent.hub = hub
        hub.on_init = []
        w.instantiate = Mock()
        w.register_with_event_loop(parent, Mock(name='loop'))
        assert parent.consumer.on_task_message

    def test_info_without_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.max_concurrency = '10'
        parent.min_concurrency = '2'
        parent.use_eventloop = False
        w = autoscale.WorkerComponent(parent)
        w.create(parent)
        info = w.info(parent)
        assert 'autoscaler' in info
        assert parent.autoscaler_cls().info.called


class test_Autoscaler:

    def setup_method(self):
        self.pool = MockPool(3)

    def test_stop(self):

        class Scaler(autoscale.Autoscaler):
            alive = True
            joined = False

            def is_alive(self):
                return self.alive

            def join(self, timeout=None):
                self.joined = True

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        # Don't allow thread joining or event waiting to block the test
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert x.joined
        x.joined = False
        x.alive = False
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert not x.joined

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_body(self, sleepdeprived):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()
        assert x.pool.num_processes == 3
        _keep = [Mock(name=f'req{i}') for i in range(20)]
        [state.task_reserved(m) for m in _keep]
        x.body()
        x.body()
        assert x.pool.num_processes == 10
        state.reserved_requests.clear()
        x.body()
        assert x.pool.num_processes == 10
        x._last_scale_up = monotonic() - 10000
        x.body()
        assert x.pool.num_processes == 3

    def test_run(self):

        class Scaler(autoscale.Autoscaler):
            scale_called = False

            def body(self):
                self.scale_called = True
                getattr(self, "_bgThread__is_shutdown").set()

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        x.run()
        assert getattr(x, "_bgThread__is_shutdown").is_set()
        assert getattr(x, "_bgThread__is_stopped").is_set()
        assert x.scale_called

    def test_shrink_raises_exception(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x.pool.shrink_raises_exception = True
        x._shrink(1)

    @patch('celery.worker.autoscale.debug')
    def test_shrink_raises_ValueError(self, debug):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x._last_scale_up = monotonic() - 10000
        x.pool.shrink_raises_ValueError = True
        x.scale_down(1)
        assert debug.call_count

    def test_update(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.keepalive = -1
        assert x.processes == 3
        x.scale_up(5)
        x.update(7, None)
        assert x.processes == 7
        assert x.max_concurrency == 7
        x.scale_down(4)
        x.update(None, 6)
        assert x.processes == 6
        assert x.min_concurrency == 6

        x.update(max=300, min=10)
        x.update(max=300, min=2)
        x.update(max=None, min=None)

    def test_prefetch_count_on_updates(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_prefetch_count_on_updates_prefetch_multiplier_gt_one(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 4
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_info(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        info = x.info()
        assert info['max'] == 10
        assert info['min'] == 3
        assert info['current'] == 3

    @patch('os._exit')
    def test_thread_crash(self, _exit):

        class _Autoscaler(autoscale.Autoscaler):

            def body(self):
                getattr(self, "_bgThread__is_shutdown").set()
                raise OSError('foo')
        worker = Mock(name='worker')
        x = _Autoscaler(self.pool, 10, 3, worker=worker)

        stderr = Mock()
        p, sys.stderr = sys.stderr, stderr
        try:
            x.run()
        finally:
            sys.stderr = p
        _exit.assert_called_with(1)
        stderr.write.assert_called()

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_no_negative_scale(self, sleepdeprived):
        total_num_processes = []
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()  # the body func scales up or down

        _keep = [Mock(name=f'req{i}') for i in range(35)]
        for req in _keep:
            state.task_reserved(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        for req in _keep:
            state.task_ready(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        assert all(x.min_concurrency <= i <= x.max_concurrency
                   for i in total_num_processes)

    def test_disable_prefetch_respects_max_concurrency(self):
        """Test that disable_prefetch respects autoscale max_concurrency setting"""
        from celery.worker.consumer.tasks import Tasks

        # Create a mock consumer with autoscale and disable_prefetch enabled
        consumer = Mock()
        consumer.app = Mock()
        consumer.app.conf.worker_disable_prefetch = True
        consumer.pool = Mock()
        consumer.pool.num_processes = 10
        consumer.controller = Mock()
        consumer.controller.max_concurrency = 5  # Lower than pool processes

        # Mock task consumer setup
        consumer.task_consumer = Mock()
        consumer.task_consumer.channel = Mock()
        consumer.task_consumer.channel.qos = Mock()
        consumer.task_consumer.channel.qos.can_consume = Mock(return_value=True)

        # Mock the connection and other required attributes
        consumer.connection = Mock()
        consumer.connection.default_channel = Mock()
        consumer.connection.transport = Mock()
        consumer.connection.transport.driver_type = 'redis'
        consumer.initial_prefetch_count = 20
        consumer.update_strategies = Mock()
        consumer.on_decode_error = Mock()

        # Mock the amqp TaskConsumer
        consumer.app.amqp = Mock()
        consumer.app.amqp.TaskConsumer = Mock(return_value=consumer.task_consumer)

        tasks_instance = Tasks(consumer)

        # Mock 5 reserved requests (at autoscale limit of 5)
        mock_requests = [Mock() for _ in range(5)]
        with patch('celery.worker.state.reserved_requests', mock_requests):
            tasks_instance.start(consumer)

            # Should not be able to consume when at autoscale limit
            assert consumer.task_consumer.channel.qos.can_consume() is False

        # Test with 4 reserved requests (under autoscale limit of 5)
        mock_requests = [Mock() for _ in range(4)]
        with patch('celery.worker.state.reserved_requests', mock_requests):
            tasks_instance.start(consumer)

            # Should be able to consume when under autoscale limit
            assert consumer.task_consumer.channel.qos.can_consume() is True
