import os
from datetime import datetime
from pickle import dumps, loads
from unittest.mock import Mock, patch

import pytest

from celery import states, uuid
from celery.app.task import Context
from celery.exceptions import ImproperlyConfigured

pytest.importorskip('sqlalchemy')

from celery.backends.database import DatabaseBackend, retry, session, session_cleanup  # noqa
from celery.backends.database.models import Task, TaskSet  # noqa
from celery.backends.database.session import PREPARE_MODELS_MAX_RETRIES, ResultModelBase, SessionManager  # noqa
from t import skip  # noqa

DB_PATH = "test.db"


class SomeClass:

    def __init__(self, data):
        self.data = data

    def __eq__(self, cmp):
        return self.data == cmp.data


class test_session_cleanup:

    def test_context(self):
        session = Mock(name='session')
        with session_cleanup(session):
            pass
        session.close.assert_called_with()

    def test_context_raises(self):
        session = Mock(name='session')
        with pytest.raises(KeyError):
            with session_cleanup(session):
                raise KeyError()
        session.rollback.assert_called_with()
        session.close.assert_called_with()


@skip.if_pypy
class test_ModelsIdFieldTypeVariations:

    def test_for_mssql_dialect(self):
        """Test that ID columns use BigInteger for MSSQL and Integer for other dialects."""
        from sqlalchemy import BigInteger, Integer
        from sqlalchemy.dialects import mssql, mysql, oracle, postgresql, sqlite

        models = [Task, TaskSet]
        id_columns = [m.__table__.columns['id'] for m in models]

        for dialect in [mssql, postgresql, mysql, sqlite, oracle]:
            for id_column in id_columns:
                compiled_type = id_column.type.dialect_impl(dialect.dialect())
                if dialect == mssql:
                    assert isinstance(compiled_type, BigInteger)
                else:
                    assert isinstance(compiled_type, Integer)


@skip.if_pypy
class test_DateDoneIndex:
    """Test that date_done columns have index=True on Task and TaskSet models."""

    def test_task_date_done_has_index(self):
        col = Task.__table__.columns['date_done']
        assert col.index is True, "Task.date_done should have index=True"

    def test_taskset_date_done_has_index(self):
        col = TaskSet.__table__.columns['date_done']
        assert col.index is True, "TaskSet.date_done should have index=True"


class test_DateDoneColumnDefaults:
    """Test that date_done column defaults are callables, not fixed values.

    The default and onupdate values must be callables (lambdas) so that
    datetime.now() is evaluated per-row at INSERT/UPDATE time, not once
    at module import time.
    """

    def test_task_date_done_default_is_callable(self):
        """Task.date_done default should be a callable."""
        col = Task.__table__.columns['date_done']
        assert col.default is not None, \
            "Task.date_done should have a default"
        assert callable(col.default.arg), \
            "Task.date_done default should be a callable, not a fixed datetime"

    def test_task_date_done_onupdate_is_callable(self):
        """Task.date_done onupdate should be a callable (lambda)."""
        col = Task.__table__.columns['date_done']
        assert col.onupdate is not None, \
            "Task.date_done should have an onupdate"
        assert callable(col.onupdate.arg), \
            "Task.date_done onupdate should be a callable, not a fixed datetime"

    def test_taskset_date_done_default_is_callable(self):
        """TaskSet.date_done default should be a callable."""
        col = TaskSet.__table__.columns['date_done']
        assert col.default is not None, \
            "TaskSet.date_done should have a default"
        assert callable(col.default.arg), \
            "TaskSet.date_done default should be a callable, not a fixed datetime"


@skip.if_pypy
class test_DatabaseBackend:

    @pytest.fixture(autouse=True)
    def remmove_db(self):
        yield
        if os.path.exists(DB_PATH):
            os.remove(DB_PATH)

    def setup_method(self):
        self.uri = 'sqlite:///' + DB_PATH
        self.app.conf.result_serializer = 'pickle'

    def test_retry_helper(self):
        from celery.backends.database import DatabaseError

        calls = [0]

        @retry
        def raises():
            calls[0] += 1
            raise DatabaseError(1, 2, 3)

        with pytest.raises(DatabaseError):
            raises(max_retries=5)
        assert calls[0] == 5

    def test_retry_helper_calls_on_backend_retryable_error(self):
        from celery.backends.database import DatabaseError

        calls = [0]
        hook_calls = []

        mock_backend = Mock()
        mock_backend.on_backend_retryable_error = Mock(side_effect=lambda exc: hook_calls.append(exc))

        @retry
        def raises_with_backend(backend):
            calls[0] += 1
            raise DatabaseError(1, 2, 3)

        with pytest.raises(DatabaseError):
            raises_with_backend(mock_backend, max_retries=3)

        assert calls[0] == 3
        assert mock_backend.on_backend_retryable_error.call_count == 3
        for exc in hook_calls:
            assert isinstance(exc, DatabaseError)

    def test_retry_helper_without_hook(self):
        from celery.backends.database import DatabaseError

        calls = [0]

        mock_backend = Mock(spec=[])

        @retry
        def raises_with_backend(backend):
            calls[0] += 1
            raise DatabaseError(1, 2, 3)

        with pytest.raises(DatabaseError):
            raises_with_backend(mock_backend, max_retries=3)

        assert calls[0] == 3

    def test_retry_helper_hook_failure_continues(self):
        from celery.backends.database import DatabaseError

        calls = [0]

        mock_backend = Mock()
        mock_backend.on_backend_retryable_error = Mock(side_effect=RuntimeError("hook failed"))

        @retry
        def raises_with_backend(backend):
            calls[0] += 1
            raise DatabaseError(1, 2, 3)

        with pytest.raises(DatabaseError):
            raises_with_backend(mock_backend, max_retries=3)

        assert calls[0] == 3
        assert mock_backend.on_backend_retryable_error.call_count == 3

    def test_missing_dburi_raises_ImproperlyConfigured(self):
        self.app.conf.database_url = None
        with pytest.raises(ImproperlyConfigured):
            DatabaseBackend(app=self.app)

    def test_engine_options_include_pool_health_defaults(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.engine_options["pool_pre_ping"] is True
        assert tb.engine_options["pool_recycle"] == 3600

    def test_engine_options_explicit_values_override_defaults(self):
        self.app.conf.database_engine_options = {"pool_pre_ping": False}
        tb = DatabaseBackend(
            self.uri,
            app=self.app,
            engine_options={"pool_recycle": 15},
        )
        assert tb.engine_options["pool_pre_ping"] is False
        assert tb.engine_options["pool_recycle"] == 15

    def test_exception_safe_to_retry(self):
        from celery.backends.database import DatabaseError, InvalidRequestError, StaleDataError

        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.exception_safe_to_retry(DatabaseError("", "", Exception("db error")))
        assert tb.exception_safe_to_retry(InvalidRequestError())
        assert tb.exception_safe_to_retry(StaleDataError())
        assert not tb.exception_safe_to_retry(RuntimeError("not retryable"))

    def test_exception_safe_to_retry_with_interface_error(self):
        from celery.backends.database import InterfaceError

        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.exception_safe_to_retry(InterfaceError("", None, Exception("connection lost")))

    def test_on_backend_retryable_error_invalidates_session(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tb.session_manager.invalidate = Mock()

        tb.on_backend_retryable_error(RuntimeError("retryable"))
        tb.session_manager.invalidate.assert_called_once_with(tb.url)

    def test_on_backend_retryable_error_called_with_exception(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tb.session_manager.invalidate = Mock()
        mock_exc = RuntimeError("connection lost")

        tb.on_backend_retryable_error(mock_exc)

        tb.session_manager.invalidate.assert_called_once_with(tb.url)

    def test_table_schema_config(self):
        self.app.conf.database_table_schemas = {
            'task': 'foo',
            'group': 'bar',
        }
        # disable table creation because schema foo and bar do not exist
        # and aren't created if they don't exist.
        self.app.conf.database_create_tables_at_setup = False
        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.task_cls.__table__.schema == 'foo'
        assert tb.task_cls.__table__.c.id.default.schema == 'foo'
        assert tb.taskset_cls.__table__.schema == 'bar'
        assert tb.taskset_cls.__table__.c.id.default.schema == 'bar'

    def test_table_name_config(self):
        self.app.conf.database_table_names = {
            'task': 'foo',
            'group': 'bar',
        }
        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.task_cls.__table__.name == 'foo'
        assert tb.taskset_cls.__table__.name == 'bar'

    def test_table_creation_at_setup_config(self):
        from sqlalchemy import inspect
        self.app.conf.database_create_tables_at_setup = True
        tb = DatabaseBackend(self.uri, app=self.app)
        engine = tb.session_manager.get_engine(tb.url)
        inspect(engine).has_table("celery_taskmeta")
        inspect(engine).has_table("celery_tasksetmeta")

    def test_missing_task_id_is_PENDING(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        assert tb.get_state('xxx-does-not-exist') == states.PENDING

    def test_missing_task_meta_is_dict_with_pending(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        meta = tb.get_task_meta('xxx-does-not-exist-at-all')
        assert meta['status'] == states.PENDING
        assert meta['task_id'] == 'xxx-does-not-exist-at-all'
        assert meta['result'] is None
        assert meta['traceback'] is None

    def test_mark_as_done(self):
        tb = DatabaseBackend(self.uri, app=self.app)

        tid = uuid()

        assert tb.get_state(tid) == states.PENDING
        assert tb.get_result(tid) is None

        tb.mark_as_done(tid, 42)
        assert tb.get_state(tid) == states.SUCCESS
        assert tb.get_result(tid) == 42

    def test_is_pickled(self):
        tb = DatabaseBackend(self.uri, app=self.app)

        tid2 = uuid()
        result = {'foo': 'baz', 'bar': SomeClass(12345)}
        tb.mark_as_done(tid2, result)
        # is serialized properly.
        rindb = tb.get_result(tid2)
        assert rindb.get('foo') == 'baz'
        assert rindb.get('bar').data == 12345

    def test_mark_as_started(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tid = uuid()
        tb.mark_as_started(tid)
        assert tb.get_state(tid) == states.STARTED

    def test_mark_as_revoked(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tid = uuid()
        tb.mark_as_revoked(tid)
        assert tb.get_state(tid) == states.REVOKED

    def test_mark_as_retry(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tid = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            import traceback
            trace = '\n'.join(traceback.format_stack())
            tb.mark_as_retry(tid, exception, traceback=trace)
            assert tb.get_state(tid) == states.RETRY
            assert isinstance(tb.get_result(tid), KeyError)
            assert tb.get_traceback(tid) == trace

    def test_mark_as_failure(self):
        tb = DatabaseBackend(self.uri, app=self.app)

        tid3 = uuid()
        try:
            raise KeyError('foo')
        except KeyError as exception:
            import traceback
            trace = '\n'.join(traceback.format_stack())
            tb.mark_as_failure(tid3, exception, traceback=trace)
            assert tb.get_state(tid3) == states.FAILURE
            assert isinstance(tb.get_result(tid3), KeyError)
            assert tb.get_traceback(tid3) == trace

    def test_forget(self):
        tb = DatabaseBackend(self.uri, backend='memory://', app=self.app)
        tid = uuid()
        tb.mark_as_done(tid, {'foo': 'bar'})
        tb.mark_as_done(tid, {'foo': 'bar'})
        x = self.app.AsyncResult(tid, backend=tb)
        x.forget()
        assert x.result is None

    def test_process_cleanup(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        tb.process_cleanup()

    @pytest.mark.usefixtures('depends_on_current_app')
    def test_reduce(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        assert loads(dumps(tb))

    def test_save__restore__delete_group(self):
        tb = DatabaseBackend(self.uri, app=self.app)

        tid = uuid()
        res = {'something': 'special'}
        assert tb.save_group(tid, res) == res

        res2 = tb.restore_group(tid)
        assert res2 == res

        tb.delete_group(tid)
        assert tb.restore_group(tid) is None

        assert tb.restore_group('xxx-nonexisting-id') is None

    def test_cleanup(self):
        tb = DatabaseBackend(self.uri, app=self.app)
        for i in range(10):
            tb.mark_as_done(uuid(), 42)
            tb.save_group(uuid(), {'foo': 'bar'})
        s = tb.ResultSession()
        for t in s.query(Task).all():
            t.date_done = datetime.now() - tb.expires * 2
        for t in s.query(TaskSet).all():
            t.date_done = datetime.now() - tb.expires * 2
        s.commit()
        s.close()

        tb.cleanup()

    def test_Task__repr__(self):
        assert 'foo' in repr(Task('foo'))

    def test_TaskSet__repr__(self):
        assert 'foo', repr(TaskSet('foo' in None))


@skip.if_pypy
class test_DatabaseBackend_result_extended():
    def setup_method(self):
        self.uri = 'sqlite:///' + DB_PATH
        self.app.conf.result_serializer = 'pickle'
        self.app.conf.result_extended = True

    @pytest.mark.parametrize(
        'result_serializer, args, kwargs',
        [
            ('pickle', (SomeClass(1), SomeClass(2)), {'foo': SomeClass(123)}),
            ('json', ['a', 'b'], {'foo': 'bar'}),
        ],
        ids=['using pickle', 'using json']
    )
    def test_store_result(self, result_serializer, args, kwargs):
        self.app.conf.result_serializer = result_serializer
        tb = DatabaseBackend(self.uri, app=self.app)
        tid = uuid()

        request = Context(args=args, kwargs=kwargs,
                          task='mytask', retries=2,
                          hostname='celery@worker_1',
                          delivery_info={'routing_key': 'celery'})

        tb.store_result(tid, {'fizz': 'buzz'}, states.SUCCESS, request=request)
        meta = tb.get_task_meta(tid)

        assert meta['result'] == {'fizz': 'buzz'}
        assert meta['args'] == args
        assert meta['kwargs'] == kwargs
        assert meta['queue'] == 'celery'
        assert meta['name'] == 'mytask'
        assert meta['retries'] == 2
        assert meta['worker'] == "celery@worker_1"

    @pytest.mark.parametrize(
        'result_serializer, args, kwargs',
        [
            ('pickle', (SomeClass(1), SomeClass(2)), {'foo': SomeClass(123)}),
            ('json', ['a', 'b'], {'foo': 'bar'}),
        ],
        ids=['using pickle', 'using json']
    )
    def test_store_none_result(self, result_serializer, args, kwargs):
        self.app.conf.result_serializer = result_serializer
        tb = DatabaseBackend(self.uri, app=self.app)
        tid = uuid()

        request = Context(args=args, kwargs=kwargs,
                          task='mytask', retries=2,
                          hostname='celery@worker_1',
                          delivery_info={'routing_key': 'celery'})

        tb.store_result(tid, None, states.SUCCESS, request=request)
        meta = tb.get_task_meta(tid)

        assert meta['result'] is None
        assert meta['args'] == args
        assert meta['kwargs'] == kwargs
        assert meta['queue'] == 'celery'
        assert meta['name'] == 'mytask'
        assert meta['retries'] == 2
        assert meta['worker'] == "celery@worker_1"

    @pytest.mark.parametrize(
        'result_serializer, args, kwargs',
        [
            ('pickle', (SomeClass(1), SomeClass(2)),
             {'foo': SomeClass(123)}),
            ('json', ['a', 'b'], {'foo': 'bar'}),
        ],
        ids=['using pickle', 'using json']
    )
    def test_get_result_meta(self, result_serializer, args, kwargs):
        self.app.conf.result_serializer = result_serializer
        tb = DatabaseBackend(self.uri, app=self.app)

        request = Context(args=args, kwargs=kwargs,
                          task='mytask', retries=2,
                          hostname='celery@worker_1',
                          delivery_info={'routing_key': 'celery'})

        meta = tb._get_result_meta(result={'fizz': 'buzz'},
                                   state=states.SUCCESS, traceback=None,
                                   request=request, format_date=False,
                                   encode=True)

        assert meta['result'] == {'fizz': 'buzz'}
        assert tb.decode(meta['args']) == args
        assert tb.decode(meta['kwargs']) == kwargs
        assert meta['queue'] == 'celery'
        assert meta['name'] == 'mytask'
        assert meta['retries'] == 2
        assert meta['worker'] == "celery@worker_1"

    @pytest.mark.parametrize(
        'result_serializer, args, kwargs',
        [
            ('pickle', (SomeClass(1), SomeClass(2)),
             {'foo': SomeClass(123)}),
            ('json', ['a', 'b'], {'foo': 'bar'}),
        ],
        ids=['using pickle', 'using json']
    )
    def test_get_result_meta_with_none(self, result_serializer, args, kwargs):
        self.app.conf.result_serializer = result_serializer
        tb = DatabaseBackend(self.uri, app=self.app)

        request = Context(args=args, kwargs=kwargs,
                          task='mytask', retries=2,
                          hostname='celery@worker_1',
                          delivery_info={'routing_key': 'celery'})

        meta = tb._get_result_meta(result=None,
                                   state=states.SUCCESS, traceback=None,
                                   request=request, format_date=False,
                                   encode=True)

        assert meta['result'] is None
        assert tb.decode(meta['args']) == args
        assert tb.decode(meta['kwargs']) == kwargs
        assert meta['queue'] == 'celery'
        assert meta['name'] == 'mytask'
        assert meta['retries'] == 2
        assert meta['worker'] == "celery@worker_1"


class test_SessionManager:

    def test_after_fork(self):
        s = SessionManager()
        assert not s.forked
        s._after_fork()
        assert s.forked

    @patch('celery.backends.database.session.create_engine')
    def test_get_engine_forked(self, create_engine):
        s = SessionManager()
        s._after_fork()
        engine = s.get_engine('dburi', foo=1)
        create_engine.assert_called_with('dburi', foo=1)
        assert engine is create_engine()
        engine2 = s.get_engine('dburi', foo=1)
        assert engine2 is engine

    @patch('celery.backends.database.session.create_engine')
    def test_get_engine_kwargs(self, create_engine):
        s = SessionManager()
        engine = s.get_engine('dbur', foo=1, pool_size=5)
        assert engine is create_engine()
        engine2 = s.get_engine('dburi', foo=1)
        assert engine2 is engine

    def test_invalidate_disposes_cached_engine(self):
        s = SessionManager()
        engine = Mock()
        s._engines['dburi'] = engine
        s._sessions['dburi'] = Mock()

        s.invalidate('dburi')

        assert 'dburi' not in s._engines
        assert 'dburi' not in s._sessions
        engine.dispose.assert_called_once_with()

    def test_invalidate_nonexistent_dburi_is_noop(self):
        s = SessionManager()
        s.invalidate('nonexistent-dburi')
        assert 'nonexistent-dburi' not in s._engines
        assert 'nonexistent-dburi' not in s._sessions

    def test_invalidate_only_engine_cached(self):
        s = SessionManager()
        engine = Mock()
        s._engines['dburi'] = engine

        s.invalidate('dburi')

        assert 'dburi' not in s._engines
        engine.dispose.assert_called_once_with()

    @patch('celery.backends.database.session.sessionmaker')
    def test_create_session_forked(self, sessionmaker):
        s = SessionManager()
        s.get_engine = Mock(name='get_engine')
        s._after_fork()
        engine, session = s.create_session('dburi', short_lived_sessions=True)
        sessionmaker.assert_called_with(bind=s.get_engine())
        assert session is sessionmaker()
        sessionmaker.return_value = Mock(name='new')
        engine, session2 = s.create_session('dburi', short_lived_sessions=True)
        sessionmaker.assert_called_with(bind=s.get_engine())
        assert session2 is not session
        sessionmaker.return_value = Mock(name='new2')
        engine, session3 = s.create_session(
            'dburi', short_lived_sessions=False)
        sessionmaker.assert_called_with(bind=s.get_engine())
        assert session3 is session2

    def test_coverage_madness(self):
        prev, session.register_after_fork = (
            session.register_after_fork, None,
        )
        try:
            SessionManager()
        finally:
            session.register_after_fork = prev

    @patch('celery.backends.database.session.create_engine')
    def test_prepare_models_terminates(self, create_engine):
        """SessionManager.prepare_models has retry logic because the creation
        of database tables by multiple workers is racy. This test patches
        the used method to always raise, so we can verify that it does
        eventually terminate.
        """
        from sqlalchemy.dialects.sqlite import dialect
        from sqlalchemy.exc import DatabaseError

        if hasattr(dialect, 'dbapi'):
            # Method name in SQLAlchemy < 2.0
            sqlite = dialect.dbapi()
        else:
            # Newer method name in SQLAlchemy 2.0
            sqlite = dialect.import_dbapi()
        manager = SessionManager()
        engine = manager.get_engine('dburi')

        def raise_err(bind):
            raise DatabaseError("", "", [], sqlite.DatabaseError)

        patch_create_all = patch.object(
            ResultModelBase.metadata, 'create_all', side_effect=raise_err)

        with pytest.raises(DatabaseError), patch_create_all as mock_create_all:
            manager.prepare_models(engine)

        assert mock_create_all.call_count == PREPARE_MODELS_MAX_RETRIES + 1

    @patch('celery.backends.database.session.create_engine')
    def test_get_engine_filters_nullpool_unsupported_kwargs(self, mock_create_engine):
        """
        Test that QueuePool-specific kwargs (like pool_size and max_overflow)
        are filtered out when creating an engine with NullPool.
        """
        from celery.backends.database.session import NullPool

        s = SessionManager()
        s.forked = False  # Ensure we're in the non-forked code path

        s.get_engine('dburi', echo_pool=True, pool_size=10, max_overflow=5)

        mock_create_engine.assert_called_once_with(
            'dburi',
            poolclass=NullPool,
        )
