from unittest.mock import ANY, Mock, PropertyMock, patch
from uuid import uuid4

import pytest
from billiard.einfo import ExceptionInfo
from kombu.exceptions import EncodeError

from celery import group, signals, states, uuid
from celery.app.task import Context
from celery.app.trace import (TraceInfo, build_tracer, fast_trace_task, get_log_policy, get_task_name,
                              log_policy_expected, log_policy_ignore, log_policy_internal, log_policy_reject,
                              log_policy_unexpected, reset_worker_optimizations, setup_worker_optimizations,
                              trace_task, trace_task_ret, traceback_clear)
from celery.backends.base import BaseDictBackend
from celery.backends.cache import CacheBackend
from celery.exceptions import BackendGetMetaError, Ignore, Reject, Retry
from celery.states import PENDING
from celery.worker.state import successful_requests


def trace(
    app, task, args=(), kwargs={}, propagate=False,
    eager=True, request=None, task_id='id-1', **opts
):
    t = build_tracer(task.name, task, eager=eager, propagate=propagate, app=app, **opts)
    ret = t(task_id, args, kwargs, request)
    return ret.retval, ret.info


class TraceCase:
    def setup_method(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        self.add = add

        @self.app.task(shared=False, ignore_result=True)
        def add_cast(x, y):
            return x + y

        self.add_cast = add_cast

        @self.app.task(shared=False)
        def raises(exc):
            raise exc

        self.raises = raises

    def trace(self, *args, **kwargs):
        return trace(self.app, *args, **kwargs)


class test_trace(TraceCase):
    def test_trace_successful(self):
        retval, info = self.trace(self.add, (2, 2), {})
        assert info is None
        assert retval == 4

    def test_trace_before_start(self):
        @self.app.task(shared=False, before_start=Mock())
        def add_with_before_start(x, y):
            return x + y

        self.trace(add_with_before_start, (2, 2), {})
        add_with_before_start.before_start.assert_called()

    def test_trace_on_success(self):
        @self.app.task(shared=False, on_success=Mock())
        def add_with_success(x, y):
            return x + y

        self.trace(add_with_success, (2, 2), {})
        add_with_success.on_success.assert_called()

    def test_get_log_policy(self):
        einfo = Mock(name='einfo')
        einfo.internal = False
        assert get_log_policy(self.add, einfo, Reject()) is log_policy_reject
        assert get_log_policy(self.add, einfo, Ignore()) is log_policy_ignore

        self.add.throws = (TypeError,)
        assert get_log_policy(self.add, einfo, KeyError()) is log_policy_unexpected
        assert get_log_policy(self.add, einfo, TypeError()) is log_policy_expected

        einfo2 = Mock(name='einfo2')
        einfo2.internal = True
        assert get_log_policy(self.add, einfo2, KeyError()) is log_policy_internal

    def test_get_task_name(self):
        assert get_task_name(Context({}), 'default') == 'default'
        assert get_task_name(Context({'shadow': None}), 'default') == 'default'
        assert get_task_name(Context({'shadow': ''}), 'default') == 'default'
        assert get_task_name(Context({'shadow': 'test'}), 'default') == 'test'

    def test_trace_after_return(self):
        @self.app.task(shared=False, after_return=Mock())
        def add_with_after_return(x, y):
            return x + y

        self.trace(add_with_after_return, (2, 2), {})
        add_with_after_return.after_return.assert_called()

    def test_with_prerun_receivers(self):
        on_prerun = Mock()
        signals.task_prerun.connect(on_prerun)
        try:
            self.trace(self.add, (2, 2), {})
            on_prerun.assert_called()
        finally:
            signals.task_prerun.receivers[:] = []

    def test_with_postrun_receivers(self):
        on_postrun = Mock()
        signals.task_postrun.connect(on_postrun)
        try:
            self.trace(self.add, (2, 2), {})
            on_postrun.assert_called()
        finally:
            signals.task_postrun.receivers[:] = []

    def test_with_success_receivers(self):
        on_success = Mock()
        signals.task_success.connect(on_success)
        try:
            self.trace(self.add, (2, 2), {})
            on_success.assert_called()
        finally:
            signals.task_success.receivers[:] = []

    def test_when_chord_part(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        add.backend = Mock()

        request = {'chord': uuid()}
        self.trace(add, (2, 2), {}, request=request)
        add.backend.mark_as_done.assert_called()
        args, kwargs = add.backend.mark_as_done.call_args
        assert args[0] == 'id-1'
        assert args[1] == 4
        assert args[2].chord == request['chord']
        assert not args[3]

    def test_when_backend_cleanup_raises(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        add.backend = Mock(name='backend')
        add.backend.process_cleanup.side_effect = KeyError()
        self.trace(add, (2, 2), {}, eager=False)
        add.backend.process_cleanup.assert_called_with()
        add.backend.process_cleanup.side_effect = MemoryError()
        with pytest.raises(MemoryError):
            self.trace(add, (2, 2), {}, eager=False)

    def test_eager_task_does_not_store_result_even_if_not_ignore_result(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        add.backend = Mock(name='backend')
        add.ignore_result = False

        self.trace(add, (2, 2), {}, eager=True)

        add.backend.mark_as_done.assert_called_once_with(
            'id-1',     # task_id
            4,          # result
            ANY,        # request
            False       # store_result
        )

    def test_eager_task_does_not_call_store_result(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = BaseDictBackend(app=self.app)
        backend.store_result = Mock()
        add.backend = backend
        add.ignore_result = False

        self.trace(add, (2, 2), {}, eager=True)

        add.backend.store_result.assert_not_called()

    def test_eager_task_will_store_result_if_proper_setting_is_set(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        add.backend = Mock(name='backend')
        add.store_eager_result = True
        add.ignore_result = False

        self.trace(add, (2, 2), {}, eager=True)

        add.backend.mark_as_done.assert_called_once_with(
            'id-1',     # task_id
            4,          # result
            ANY,        # request
            True        # store_result
        )

    def test_eager_task_with_setting_will_call_store_result(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = BaseDictBackend(app=self.app)
        backend.store_result = Mock()
        add.backend = backend
        add.store_eager_result = True
        add.ignore_result = False

        self.trace(add, (2, 2), {}, eager=True)

        add.backend.store_result.assert_called_once_with(
            'id-1',
            4,
            states.SUCCESS,
            request=ANY
        )

    def test_when_backend_raises_exception(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        add.backend = Mock(name='backend')
        add.backend.mark_as_done.side_effect = Exception()
        add.backend.mark_as_failure.side_effect = Exception("failed mark_as_failure")

        with pytest.raises(Exception):
            self.trace(add, (2, 2), {}, eager=False)

    def test_traceback_clear(self):
        import inspect
        import sys
        sys.exc_clear = Mock()
        frame_list = []

        def raise_dummy():
            frame_str_temp = str(inspect.currentframe().__repr__)
            frame_list.append(frame_str_temp)
            raise KeyError('foo')

        try:
            raise_dummy()
        except KeyError as exc:
            traceback_clear(exc)

            tb_ = exc.__traceback__
            while tb_ is not None:
                if str(tb_.tb_frame.__repr__) == frame_list[0]:
                    assert len(tb_.tb_frame.f_locals) == 0
                tb_ = tb_.tb_next

        try:
            raise_dummy()
        except KeyError as exc:
            traceback_clear()

            tb_ = exc.__traceback__
            while tb_ is not None:
                if str(tb_.tb_frame.__repr__) == frame_list[0]:
                    assert len(tb_.tb_frame.f_locals) == 0
                tb_ = tb_.tb_next

        try:
            raise_dummy()
        except KeyError as exc:
            traceback_clear(str(exc))

            tb_ = exc.__traceback__
            while tb_ is not None:
                if str(tb_.tb_frame.__repr__) == frame_list[0]:
                    assert len(tb_.tb_frame.f_locals) == 0
                tb_ = tb_.tb_next

    @patch('celery.app.trace.traceback_clear')
    def test_when_Ignore(self, mock_traceback_clear):
        @self.app.task(shared=False)
        def ignored():
            raise Ignore()

        retval, info = self.trace(ignored, (), {})
        assert info.state == states.IGNORED
        mock_traceback_clear.assert_called()

    @patch('celery.app.trace.traceback_clear')
    def test_when_Reject(self, mock_traceback_clear):
        @self.app.task(shared=False)
        def rejecting():
            raise Reject()

        retval, info = self.trace(rejecting, (), {})
        assert info.state == states.REJECTED
        mock_traceback_clear.assert_called()

    def test_backend_cleanup_raises(self):
        self.add.backend.process_cleanup = Mock()
        self.add.backend.process_cleanup.side_effect = RuntimeError()
        self.trace(self.add, (2, 2), {})

    @patch('celery.canvas.maybe_signature')
    def test_callbacks__scalar(self, maybe_signature):
        sig = Mock(name='sig')
        request = {'callbacks': [sig], 'root_id': 'root'}
        maybe_signature.return_value = sig
        retval, _ = self.trace(self.add, (2, 2), {}, request=request)
        sig.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', priority=None
        )

    @patch('celery.canvas.maybe_signature')
    def test_chain_proto2(self, maybe_signature):
        sig = Mock(name='sig')
        sig2 = Mock(name='sig2')
        request = {'chain': [sig2, sig], 'root_id': 'root'}
        maybe_signature.return_value = sig
        retval, _ = self.trace(self.add, (2, 2), {}, request=request)
        sig.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', chain=[sig2], priority=None
        )

    @patch('celery.canvas.maybe_signature')
    def test_chain_inherit_parent_priority(self, maybe_signature):
        self.app.conf.task_inherit_parent_priority = True
        sig = Mock(name='sig')
        sig2 = Mock(name='sig2')
        request = {
            'chain': [sig2, sig],
            'root_id': 'root',
            'delivery_info': {'priority': 42},
        }
        maybe_signature.return_value = sig
        retval, _ = self.trace(self.add, (2, 2), {}, request=request)
        sig.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', chain=[sig2], priority=42
        )

    @patch('celery.canvas.maybe_signature')
    def test_callbacks__EncodeError(self, maybe_signature):
        sig = Mock(name='sig')
        request = {'callbacks': [sig], 'root_id': 'root'}
        maybe_signature.return_value = sig
        sig.apply_async.side_effect = EncodeError()
        retval, einfo = self.trace(self.add, (2, 2), {}, request=request)
        assert einfo.state == states.FAILURE

    @patch('celery.canvas.maybe_signature')
    @patch('celery.app.trace.group.apply_async')
    def test_callbacks__sigs(self, group_, maybe_signature):
        sig1 = Mock(name='sig')
        sig2 = Mock(name='sig2')
        sig3 = group([Mock(name='g1'), Mock(name='g2')], app=self.app)
        sig3.apply_async = Mock(name='gapply')
        request = {'callbacks': [sig1, sig3, sig2], 'root_id': 'root'}

        def pass_value(s, *args, **kwargs):
            return s

        maybe_signature.side_effect = pass_value
        retval, _ = self.trace(self.add, (2, 2), {}, request=request)
        group_.assert_called_with((4,), parent_id='id-1', root_id='root', priority=None)
        sig3.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', priority=None
        )

    @patch('celery.canvas.maybe_signature')
    @patch('celery.app.trace.group.apply_async')
    def test_callbacks__only_groups(self, group_, maybe_signature):
        sig1 = group([Mock(name='g1'), Mock(name='g2')], app=self.app)
        sig2 = group([Mock(name='g3'), Mock(name='g4')], app=self.app)
        sig1.apply_async = Mock(name='gapply')
        sig2.apply_async = Mock(name='gapply')
        request = {'callbacks': [sig1, sig2], 'root_id': 'root'}

        def pass_value(s, *args, **kwargs):
            return s

        maybe_signature.side_effect = pass_value
        retval, _ = self.trace(self.add, (2, 2), {}, request=request)
        sig1.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', priority=None
        )
        sig2.apply_async.assert_called_with(
            (4,), parent_id='id-1', root_id='root', priority=None
        )

    def test_trace_SystemExit(self):
        with pytest.raises(SystemExit):
            self.trace(self.raises, (SystemExit(),), {})

    @patch('celery.app.trace.traceback_clear')
    def test_trace_Retry(self, mock_traceback_clear):
        exc = Retry('foo', 'bar')
        _, info = self.trace(self.raises, (exc,), {})
        assert info.state == states.RETRY
        assert info.retval is exc
        mock_traceback_clear.assert_called()

    @patch('celery.app.trace.traceback_clear')
    def test_trace_exception(self, mock_traceback_clear):
        exc = KeyError('foo')
        _, info = self.trace(self.raises, (exc,), {})
        assert info.state == states.FAILURE
        assert info.retval is exc
        mock_traceback_clear.assert_called()

    def test_trace_task_ret__no_content_type(self):
        trace_task_ret(
            self.add.name, 'id1', {}, ((2, 2), {}, {}), None, None, app=self.app,
        )

    def test_fast_trace_task__no_content_type(self):
        self.app.tasks[self.add.name].__trace__ = build_tracer(
            self.add.name, self.add, app=self.app,
        )
        fast_trace_task(
            self.add.name,
            'id1',
            {},
            ((2, 2), {}, {}),
            None,
            None,
            app=self.app,
            _loc=[self.app.tasks, {}, 'hostname'],
        )

    def test_trace_exception_propagate(self):
        with pytest.raises(KeyError):
            self.trace(self.raises, (KeyError('foo'),), {}, propagate=True)

    @patch('celery.app.trace.signals.task_internal_error.send')
    @patch('celery.app.trace.build_tracer')
    @patch('celery.app.trace.report_internal_error')
    def test_outside_body_error(self, report_internal_error, build_tracer, send):
        tracer = Mock()
        tracer.side_effect = KeyError('foo')
        build_tracer.return_value = tracer

        @self.app.task(shared=False)
        def xtask():
            pass

        trace_task(xtask, 'uuid', (), {})
        assert report_internal_error.call_count
        assert send.call_count
        assert xtask.__trace__ is tracer

    def test_backend_error_should_report_failure(self):
        """check internal error is reported as failure.

        In case of backend error, an exception may bubble up from trace and be
        caught by trace_task.
        """

        @self.app.task(shared=False)
        def xtask():
            pass

        xtask.backend = BaseDictBackend(app=self.app)
        xtask.backend.mark_as_done = Mock()
        xtask.backend.mark_as_done.side_effect = Exception()
        xtask.backend.mark_as_failure = Mock()
        xtask.backend.mark_as_failure.side_effect = Exception()

        ret, info, _, _ = trace_task(xtask, 'uuid', (), {}, app=self.app)
        assert info is not None
        assert isinstance(ret, ExceptionInfo)

    def test_deduplicate_successful_tasks__deduplication(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = CacheBackend(app=self.app, backend='memory')
        add.backend = backend
        add.store_eager_result = True
        add.ignore_result = False
        add.acks_late = True

        self.app.conf.worker_deduplicate_successful_tasks = True
        task_id = str(uuid4())
        request = {'id': task_id, 'delivery_info': {'redelivered': True}}

        assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (2, None)
        assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (None, None)

        self.app.conf.worker_deduplicate_successful_tasks = False

    def test_deduplicate_successful_tasks__no_deduplication(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = CacheBackend(app=self.app, backend='memory')
        add.backend = backend
        add.store_eager_result = True
        add.ignore_result = False
        add.acks_late = True

        self.app.conf.worker_deduplicate_successful_tasks = True
        task_id = str(uuid4())
        request = {'id': task_id, 'delivery_info': {'redelivered': True}}

        with patch('celery.app.trace.AsyncResult') as async_result_mock:
            async_result_mock().state.return_value = PENDING
            assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (2, None)
            assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (2, None)

        self.app.conf.worker_deduplicate_successful_tasks = False

    def test_deduplicate_successful_tasks__result_not_found(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = CacheBackend(app=self.app, backend='memory')
        add.backend = backend
        add.store_eager_result = True
        add.ignore_result = False
        add.acks_late = True

        self.app.conf.worker_deduplicate_successful_tasks = True
        task_id = str(uuid4())
        request = {'id': task_id, 'delivery_info': {'redelivered': True}}

        with patch('celery.app.trace.AsyncResult') as async_result_mock:
            assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (2, None)
            state_property = PropertyMock(side_effect=BackendGetMetaError)
            type(async_result_mock()).state = state_property
            assert trace(self.app, add, (1, 1), task_id=task_id, request=request) == (2, None)

        self.app.conf.worker_deduplicate_successful_tasks = False

    def test_deduplicate_successful_tasks__cached_request(self):
        @self.app.task(shared=False)
        def add(x, y):
            return x + y

        backend = CacheBackend(app=self.app, backend='memory')
        add.backend = backend
        add.store_eager_result = True
        add.ignore_result = False
        add.acks_late = True

        self.app.conf.worker_deduplicate_successful_tasks = True

        task_id = str(uuid4())
        request = {'id': task_id, 'delivery_info': {'redelivered': True}}

        successful_requests.add(task_id)

        assert trace(self.app, add, (1, 1), task_id=task_id,
                     request=request) == (None, None)

        successful_requests.clear()
        self.app.conf.worker_deduplicate_successful_tasks = False


class test_TraceInfo(TraceCase):
    class TI(TraceInfo):
        __slots__ = TraceInfo.__slots__ + ('__dict__',)

    def test_handle_error_state(self):
        x = self.TI(states.FAILURE)
        x.handle_failure = Mock()
        x.handle_error_state(self.add_cast, self.add_cast.request)
        x.handle_failure.assert_called_with(
            self.add_cast,
            self.add_cast.request,
            store_errors=self.add_cast.store_errors_even_if_ignored,
            call_errbacks=True,
        )

    def test_handle_error_state_for_eager_task(self):
        x = self.TI(states.FAILURE)
        x.handle_failure = Mock()

        x.handle_error_state(self.add, self.add.request, eager=True)
        x.handle_failure.assert_called_once_with(
            self.add,
            self.add.request,
            store_errors=False,
            call_errbacks=True,
        )

    def test_handle_error_for_eager_saved_to_backend(self):
        x = self.TI(states.FAILURE)
        x.handle_failure = Mock()

        self.add.store_eager_result = True

        x.handle_error_state(self.add, self.add.request, eager=True)
        x.handle_failure.assert_called_with(
            self.add,
            self.add.request,
            store_errors=True,
            call_errbacks=True,
        )

    @patch('celery.app.trace.ExceptionInfo')
    def test_handle_reject(self, ExceptionInfo):
        x = self.TI(states.FAILURE)
        x._log_error = Mock(name='log_error')
        req = Mock(name='req')
        x.handle_reject(self.add, req)
        x._log_error.assert_called_with(self.add, req, ExceptionInfo())


class test_stackprotection:
    def test_stackprotection(self):
        setup_worker_optimizations(self.app)
        try:

            @self.app.task(shared=False, bind=True)
            def foo(self, i):
                if i:
                    return foo(0)
                return self.request

            assert foo(1).called_directly
        finally:
            reset_worker_optimizations(self.app)

    def test_stackprotection_headers_passed_on_new_request_stack(self):
        setup_worker_optimizations(self.app)
        try:

            @self.app.task(shared=False, bind=True)
            def foo(self, i):
                if i:
                    return foo.apply(args=(i-1,), headers=456)
                return self.request

            task = foo.apply(args=(2,), headers=123, loglevel=5)
            assert task.result.result.result.args == (0,)
            assert task.result.result.result.headers == 456
            assert task.result.result.result.loglevel == 0
        finally:
            reset_worker_optimizations(self.app)

    def test_stackprotection_headers_persisted_calling_task_directly(self):
        setup_worker_optimizations(self.app)
        try:

            @self.app.task(shared=False, bind=True)
            def foo(self, i):
                if i:
                    return foo(i-1)
                return self.request

            task = foo.apply(args=(2,), headers=123, loglevel=5)
            assert task.result.args == (0,)
            assert task.result.headers == 123
            assert task.result.loglevel == 5
        finally:
            reset_worker_optimizations(self.app)