import datetime
import hashlib
import itertools
import os
import shutil
import sys
import threading
import unittest
import uuid
from queue import Queue

try:
    from redis.connection import ConnectionPool
    from redis import Redis
    from redis.exceptions import ConnectionError as RedisConnectionError
except ImportError:
    ConnectionPool = Redis = RedisConnectionError = None

from huey.api import Huey
from huey.api import MemoryHuey
from huey.api import PriorityRedisHuey
from huey.api import RedisExpireHuey
from huey.api import RedisHuey
from huey.api import SqliteHuey
from huey.api import chord
from huey.constants import EmptyData
from huey.consumer import Consumer
from huey.exceptions import ConfigurationError
from huey.exceptions import ResultTimeout
from huey.storage import FileStorage
from huey.storage import MemoryStorage
from huey.storage import RedisExpireStorage
from huey.tests.base import BaseTestCase
from huey.tests.base import TRAVIS
from huey.tests.base import slow_test


def get_redis_version():
    # Major version of the locally-running redis server, or 0 if the redis
    # client library or server are not available.
    if Redis is None:
        return 0
    try:
        info = Redis().info()
    except Exception:
        return 0
    return int(info['redis_version'].split('.', 1)[0])

REDIS_VERSION = get_redis_version()
requires_redis = unittest.skipIf(REDIS_VERSION == 0, 'requires redis server')


class StorageTests(object):
    destructive_reads = True

    def setUp(self):
        super(StorageTests, self).setUp()
        self.s = self.huey.storage
        self.s.flush_all()

    def tearDown(self):
        super(StorageTests, self).tearDown()
        self.s.flush_all()

    def test_queue_methods(self):
        for i in range(3):
            self.s.enqueue(b'item-%d' % i)

        # Remove two items (this API is not used, but we'll test it anyways).
        self.assertEqual(self.s.dequeue(), b'item-0')
        self.assertEqual(self.s.queue_size(), 2)
        self.assertEqual(self.s.enqueued_items(), [b'item-1', b'item-2'])
        self.assertEqual(self.s.dequeue(), b'item-1')
        self.assertEqual(self.s.dequeue(), b'item-2')
        self.assertTrue(self.s.dequeue() is None)

        self.assertEqual(self.s.queue_size(), 0)

        # Test flushing the queue.
        self.s.enqueue(b'item-3')
        self.assertEqual(self.s.queue_size(), 1)
        self.s.flush_queue()
        self.assertEqual(self.s.queue_size(), 0)

    def test_schedule_methods(self):
        timestamp = datetime.datetime(2000, 1, 2, 3, 4, 5)
        second = datetime.timedelta(seconds=1)

        items = ((b'p1', timestamp + second),
                 (b'p0', timestamp),
                 (b'n1', timestamp - second),
                 (b'p2', timestamp + second + second))
        for data, ts in items:
            self.s.add_to_schedule(data, ts)

        self.assertEqual(self.s.schedule_size(), 4)

        # Read from the schedule up-to the "p0" timestamp.
        sched = self.s.read_schedule(timestamp)
        self.assertEqual(sched, [b'n1', b'p0'])

        self.assertEqual(self.s.scheduled_items(), [b'p1', b'p2'])
        self.assertEqual(self.s.schedule_size(), 2)
        sched = self.s.read_schedule(datetime.datetime.now())
        self.assertEqual(sched, [b'p1', b'p2'])
        self.assertEqual(self.s.schedule_size(), 0)
        self.assertEqual(self.s.read_schedule(datetime.datetime.now()), [])

    def test_result_store_methods(self):
        # Put and peek at data. Verify missing keys return EmptyData sentinel.
        self.s.put_data(b'k1', b'v1')
        self.s.put_data(b'k2', b'v2')
        self.assertEqual(self.s.peek_data(b'k2'), b'v2')
        self.assertEqual(self.s.peek_data(b'k1'), b'v1')
        self.assertTrue(self.s.peek_data(b'kx') is EmptyData)
        self.assertEqual(self.s.result_store_size(), 2)

        # Verify we can overwrite existing keys and that pop will remove the
        # key/value pair. Subsequent pop on missing key will return EmptyData.
        self.s.put_data(b'k1', b'v1-x')
        self.assertEqual(self.s.peek_data(b'k1'), b'v1-x')
        self.assertEqual(self.s.pop_data(b'k1'), b'v1-x')
        if self.destructive_reads:
            self.assertTrue(self.s.pop_data(b'k1') is EmptyData)
        else:
            self.assertEqual(self.s.pop_data(b'k1'), b'v1-x')
            self.assertTrue(self.s.delete_data(b'k1'))

        self.assertFalse(self.s.has_data_for_key(b'k1'))
        self.assertTrue(self.s.has_data_for_key(b'k2'))
        self.assertEqual(self.s.result_store_size(), 1)

        # Test put-if-empty.
        self.assertTrue(self.s.put_if_empty(b'k1', b'v1-y'))
        self.assertFalse(self.s.put_if_empty(b'k1', b'v1-z'))
        self.assertEqual(self.s.peek_data(b'k1'), b'v1-y')

        # Test deletion.
        self.assertTrue(self.s.put_if_empty(b'k3', b'v3'))
        self.assertTrue(self.s.delete_data(b'k3'))
        self.assertFalse(self.s.delete_data(b'k3'))

        # Test introspection.
        state = self.s.result_items()  # Normalize keys to unicode strings.
        clean = {k.decode('utf8') if isinstance(k, bytes) else k: v
                 for k, v in state.items()}
        self.assertEqual(clean, {'k1': b'v1-y', 'k2': b'v2'})
        self.s.flush_results()
        self.assertEqual(self.s.result_store_size(), 0)
        self.assertEqual(self.s.result_items(), {})

    def test_wait_result(self):
        key = str(uuid.uuid4())
        self.assertFalse(self.s.wait_result(key, timeout=0.1))
        self.s.put_data(key, b'v1', is_result=True)
        self.assertTrue(self.s.wait_result(key, timeout=1))

    def test_priority(self):
        if not self.s.priority:
            raise unittest.SkipTest('priority support required')

        priorities = (1, None, 5, None, 3, None, 9, None, 7, 0)
        for i, priority in enumerate(priorities):
            item = 'i%s-%s' % (i, priority)
            self.s.enqueue(item.encode('utf8'), priority)

        expected = [b'i6-9', b'i8-7', b'i2-5', b'i4-3', b'i0-1',
                    b'i1-None', b'i3-None', b'i5-None', b'i7-None', b'i9-0']
        self.assertEqual([self.s.dequeue() for _ in range(10)], expected)

    def test_counter(self):
        self.assertEqual(self.s.incr('k1'), 1)
        self.assertEqual(self.s.incr('k1'), 2)
        self.assertEqual(self.s.incr('k2', amount=10), 10)
        self.assertEqual(self.s.incr('k2', amount=-5), 5)
        self.assertEqual(self.s.incr('k2', amount=0), 5)
        self.assertEqual(self.s.incr('k3', 0), 0)

        self.s.delete_counter('k1')
        self.s.delete_counter('kx')
        self.assertEqual(self.s.incr('k1'), 1)
        self.assertEqual(self.s.incr('k2', amount=2), 7)

        self.s.flush_counters()
        self.assertEqual(self.s.incr('k1'), 1)
        self.assertEqual(self.s.incr('k2'), 1)

        # Ensure we don't collide w/result store.
        self.assertEqual(self.s.incr('k1'), 2)
        self.assertFalse(self.s.has_data_for_key('k1'))
        self.s.put_data('k1', b'test')
        self.assertEqual(self.s.incr('k1'), 3)
        self.assertEqual(self.s.pop_data('k1'), b'test')
        self.assertEqual(self.s.incr('k1'), 4)

    @slow_test()
    def test_consumer_integration(self):
        @self.huey.task()
        def task_a(n):
            return n + 1

        @self.huey.task()
        def total(ns):
            return sum(ns)

        with self.consumer_context():
            r1 = task_a(1)
            r2 = task_a(2)
            r3 = task_a(3)

            c = chord([task_a.s(1), task_a.s(2)], total)
            cr = self.huey.enqueue(c)

            self.assertEqual(r1.get(blocking=True, timeout=5), 2)
            self.assertEqual(r2.get(blocking=True, timeout=5), 3)
            self.assertEqual(r3.get(blocking=True, timeout=5), 4)
            self.assertEqual(cr.get(blocking=True, timeout=5), 5)
            self.assertEqual(cr.results(), [2, 3])

            task_a.revoke()
            self.assertTrue(task_a.is_revoked())
            self.assertTrue(task_a.restore())


class TestMemoryStorage(StorageTests, BaseTestCase):
    def get_huey(self):
        return MemoryHuey(utc=False)

    def test_put_if_empty_concurrent(self):
        nthreads, nkeys = 8, 20
        barrier = threading.Barrier(nthreads)
        winners = []

        def run(n):
            barrier.wait()
            for i in range(nkeys):
                if self.s.put_if_empty(b'k-%d' % i, b'%d' % n):
                    winners.append(i)

        threads = [threading.Thread(target=run, args=(n,))
                   for n in range(nthreads)]
        for t in threads: t.start()
        for t in threads: t.join()

        # Each key was acquired by exactly one thread.
        self.assertEqual(sorted(winners), list(range(nkeys)))


@requires_redis
class TestRedisStorage(StorageTests, BaseTestCase):
    def get_huey(self):
        return RedisHuey(utc=False)

    def test_conflicting_init_args(self):
        options = {'host': 'localhost', 'url': 'redis://localhost',
                   'connection_pool': ConnectionPool()}
        combinations = itertools.combinations(options.items(), 2)
        for kwargs in (dict(item) for item in combinations):
            self.assertRaises(ConfigurationError, lambda: RedisHuey(**kwargs))

        # None values are fine, however.
        RedisHuey(host=None, port=None, db=None, url='redis://localhost')


class TestRedisStorageWaitResult(TestRedisStorage):
    def get_huey(self):
        return RedisHuey(utc=False, notify_result=True, notify_result_ttl=30)


@requires_redis
class TestRedisExpireStorage(StorageTests, BaseTestCase):
    # Note that this does not subclass the StorageTests. This is partly because
    # the functionality should already be covered by the TestRedisStorage, as
    # the RedisExpireStorage is a subclass of RedisStorage -- but also because
    # the way the result store functions is fundamentally different, relying on
    # the database to handle result removal via expiration.
    destructive_reads = False

    def get_huey(self):
        return RedisExpireHuey(expire_time=3600, utc=False, blocking=False)

    def test_expire_results(self):
        self.s.put_data(b'k1', b'v1')
        self.s.put_data(b'k2', b'v2', is_result=True)

        conn = self.s.conn  # Underlying Redis client.

        # By default the put_data() API treats keys as being persistent. If we
        # specifically included the "is_result=True" flag, then the key will be
        # given a TTL.
        self.assertEqual(conn.ttl(self.s.result_key(b'k1')), -1)
        self.assertTrue(3580 <= conn.ttl(self.s.result_key(b'k2')) <= 3600)

        # Non-existent keys return -2. See redis docs for TTL command.
        self.assertEqual(conn.ttl(self.s.result_key(b'k3')), -2)

        # Non-expired keys return -1.
        conn.set(self.s.result_key(b'k3'), b'v3')
        self.assertEqual(conn.ttl(self.s.result_key(b'k3')), -1)

        # Verify behavior of put_if_empty and has_data_for_key.
        self.assertTrue(self.s.has_data_for_key(b'k2'))
        self.assertFalse(self.s.put_if_empty(b'k2', b'v2-x'))
        self.assertFalse(self.s.has_data_for_key(b'k4'))
        self.assertTrue(self.s.put_if_empty(b'k4', b'v4'))

        # Verify behavior of delete.
        self.assertTrue(self.s.delete_data(b'k2'))
        self.assertFalse(self.s.delete_data(b'k2'))

        # Check the result items.
        self.assertEqual(self.s.result_items(), {
            b'k1': b'v1',
            b'k3': b'v3',
            b'k4': b'v4'})
        self.assertEqual(self.s.result_store_size(), 3)

    def test_integration_2(self):
        @self.huey.task()
        def task_a(n):
            return n + 1

        r1, r2, r3 = [task_a(i) for i in (1, 2, 3)]
        r2.revoke()
        self.assertTrue(r2.is_revoked())
        self.assertEqual(self.huey.result_count(), 1)  # Revoke key.

        self.assertEqual(self.execute_next(), 2)
        self.assertEqual(self.huey.result_count(), 2)  # Revoke key and r1.

        self.assertTrue(self.execute_next() is None)
        self.assertEqual(self.huey.result_count(), 1)  # Just r1 now.
        self.assertFalse(r2.is_revoked())

        self.assertEqual(self.execute_next(), 4)
        self.assertEqual(self.huey.result_count(), 2)  # r1 and r3.

        for _ in range(3):
            self.assertEqual(r1(), 2)
            self.assertEqual(r3(), 4)
            r1.reset()
            r3.reset()
        self.assertEqual(self.huey.result_count(), 2)  # r1 and r3 still there.


@unittest.skipIf(REDIS_VERSION < 5, 'Requires Redis >= 5.0')
class TestPriorityRedisStorage(TestRedisStorage):
    def get_huey(self):
        return PriorityRedisHuey(utc=False)


@unittest.skipIf(REDIS_VERSION < 5, 'Requires Redis >= 5.0')
class TestPriorityRedisStorageNotBlocking(TestRedisStorage):
    def get_huey(self):
        return PriorityRedisHuey(utc=False, blocking=False)


@unittest.skipIf(Redis is None, 'requires redis python module')
class TestRedisStorageOffline(BaseTestCase):
    # Verify client-side behavior using a stub client -- these tests do not
    # require a live redis server.
    def get_huey(self):
        return RedisHuey(utc=False)

    def test_dequeue_error_handling(self):
        s = self.huey.storage

        class StubConnErr(object):
            def brpop(self, key, timeout=None):
                raise RedisConnectionError('cannot connect')
        s.conn = StubConnErr()

        # Connection errors propagate, so the worker logs the error and
        # applies backoff rather than treating it as an empty queue.
        self.assertRaises(RedisConnectionError, s.dequeue)

        class StubEmpty(object):
            def brpop(self, key, timeout=None):
                return None  # BRPOP timed out, queue is empty.
        s.conn = StubEmpty()
        self.assertTrue(s.dequeue() is None)

    def test_wait_result_timeout_clamp(self):
        huey = RedisHuey(utc=False, notify_result=True)
        s = huey.storage
        captured = []

        class StubConn(object):
            def hexists(self, key, k):
                return False
            def blpop(self, key, timeout=None):
                captured.append(timeout)
                return None
        s.conn = StubConn()

        s.redis_version = (5, 0, 0)
        self.assertFalse(s.wait_result('k1', timeout=30))
        self.assertFalse(s.wait_result('k1', timeout=0.5))

        s.redis_version = (7, 4, 0)
        self.assertFalse(s.wait_result('k1', timeout=0.5))

        # For redis < 6 the timeout is coerced to an int >= 1 (rather than
        # being clamped *down* to at-most 1 second, or, for sub-second floats,
        # truncated to zero -- which blocks indefinitely). Newer servers
        # receive the timeout unmodified.
        self.assertEqual(captured, [30, 1, 0.5])

    def test_enqueued_items_limit(self):
        s = self.huey.storage
        calls = []

        class StubConn(object):
            def lrange(self, key, start, stop):
                calls.append((start, stop))
                # Head of the redis list is the most-recently enqueued item.
                return [b'i2', b'i1', b'i0']
        s.conn = StubConn()

        self.assertEqual(s.enqueued_items(), [b'i0', b'i1', b'i2'])
        self.assertEqual(s.enqueued_items(3), [b'i0', b'i1', b'i2'])

        # With a limit, items are read from the consumption end of the list,
        # e.g. the next-N items to be dequeued.
        self.assertEqual(calls, [(0, -1), (-3, -1)])


class TestSqliteStorage(StorageTests, BaseTestCase):
    def tearDown(self):
        super(TestSqliteStorage, self).tearDown()
        if os.path.exists('huey_storage.db'):
            os.unlink('huey_storage.db')

    def get_huey(self):
        return SqliteHuey(filename='huey_storage.db', timeout=3)

    def test_timeout(self):
        self.assertEqual(self.s._timeout, 3)
        curs = self.s.conn.execute('pragma busy_timeout')
        self.assertEqual(curs.fetchone(), (3000,))


class TestFileStorageMethods(StorageTests, BaseTestCase):
    path = '/tmp/test-huey-storage'
    result_path = '/tmp/test-huey-storage/results'
    queue_path = '/tmp/test-huey-storage/queue'

    def tearDown(self):
        super(TestFileStorageMethods, self).tearDown()
        if os.path.exists(self.path):
            shutil.rmtree(self.path)

    def get_huey(self):
        return Huey('test-file-storage', storage_class=FileStorage,
                    path=self.path, levels=2, use_thread_lock=True)

    def test_filesystem_result_store(self):
        s = self.huey.storage
        self.assertEqual(s.result_items(), {})

        keys = (b'k1', b'k2', b'kx')
        for key in keys:
            checksum = hashlib.md5(key).hexdigest()
            b0, b1 = checksum[0], checksum[1]
            # Default is to use two levels.
            key_path = os.path.join(self.result_path, b0, b1)
            key_filename = os.path.join(key_path, checksum)

            self.assertFalse(os.path.exists(key_filename))
            self.assertFalse(os.path.exists(key_path))

            s.put_data(key, b'test-%s' % key)
            self.assertTrue(os.path.exists(key_path))
            self.assertTrue(os.path.exists(key_filename))

            self.assertEqual(s.pop_data(key), b'test-%s' % key)
            self.assertTrue(os.path.exists(key_path))
            self.assertFalse(os.path.exists(key_filename))

        # Flushing the results blows away everything.
        s.flush_results()
        self.assertTrue(os.path.exists(self.result_path))
        self.assertEqual(os.listdir(self.result_path), [])

    def test_fs_multithreaded(self):
        l = threading.Lock()

        def create_tasks(t, n, q):
            for i in range(n):
                with l:
                    message = str((t * n) + i)
                    self.huey.storage.enqueue(message.encode('utf8'))
                    q.put(message)

        def dequeue_tasks(q):
            while True:
                with l:
                    data = self.huey.storage.dequeue()
                    if data is None:
                        break
                    q.put(data.decode('utf8'))

        nthreads = 10
        ntasks = 50
        in_q = Queue()
        threads = []
        for i in range(nthreads):
            t = threading.Thread(target=create_tasks, args=(i, ntasks, in_q))
            t.daemon = True
            threads.append(t)

        for t in threads: t.start()
        for t in threads: t.join(timeout=10.)

        self.assertEqual(self.huey.pending_count(), nthreads * ntasks)

        out_q = Queue()
        threads = []
        for i in range(nthreads):
            t = threading.Thread(target=dequeue_tasks, args=(out_q,))
            t.daemon = True
            threads.append(t)

        for t in threads: t.start()
        for t in threads: t.join(timeout=10.)

        self.assertEqual(out_q.qsize(), nthreads * ntasks)
        self.assertEqual(self.huey.pending_count(), 0)

        # Ensure that the order in which tasks were enqueued is the order in
        # which they are dequeued.
        for i in range(nthreads * ntasks):
            self.assertEqual(in_q.get(), out_q.get())

    @unittest.skipIf(TRAVIS, 'skipping test that is flaky on travis-ci')
    def test_consumer_integration(self):
        return super(TestFileStorageMethods, self).test_consumer_integration()


try:
    from huey.contrib.valkey_glide import ValkeyGlideHuey
except ImportError:
    ValkeyGlideHuey = None


@unittest.skipIf(ValkeyGlideHuey is None, 'valkey-glide-sync not installed')
class TestValkeyGlideStorage(StorageTests, BaseTestCase):
    def get_huey(self):
        return ValkeyGlideHuey(utc=False)
