import datetime
import logging
import threading
import time
import unittest

from huey import crontab
from huey import Huey
from huey.backends.dummy import DummyDataStore
from huey.backends.dummy import DummyQueue
from huey.bin.huey_consumer import Consumer
from huey.bin.huey_consumer import IterableQueue
from huey.bin.huey_consumer import WorkerThread
from huey.exceptions import QueueException
from huey.registry import registry
from huey.utils import local_to_utc


# store some global state
state = {}

# create a queue, result store and invoker for testing
test_queue = DummyQueue('test-queue')
test_result_store = DummyDataStore('test-queue')
test_huey = Huey(test_queue, test_result_store)

@test_huey.task()
def modify_state(k, v):
    state[k] = v
    return v

@test_huey.task()
def blow_up():
    raise Exception('blowed up')

@test_huey.task(retries=3)
def retry_command(k, always_fail=True):
    if k not in state:
        if not always_fail:
            state[k] = 'fixed'
        raise Exception('fappsk')
    return state[k]

@test_huey.task(retries=3, retry_delay=10)
def retry_command_slow(k, always_fail=True):
    if k not in state:
        if not always_fail:
            state[k] = 'fixed'
        raise Exception('fappsk')
    return state[k]

@test_huey.periodic_task(crontab(minute='0'))
def every_hour():
    state['p'] = 'y'

# create a log handler that will track messages generated by the consumer
class TestLogHandler(logging.Handler):
    def __init__(self, *args, **kwargs):
        self.messages = []
        logging.Handler.__init__(self, *args, **kwargs)

    def emit(self, record):
        self.messages.append(record.getMessage())

logger = logging.getLogger('huey.consumer')

class ConsumerTestCase(unittest.TestCase):
    def setUp(self):
        global state
        state = {}

        self.orig_pc = registry._periodic_tasks
        registry._periodic_commands = [every_hour.task_class()]

        self.orig_sleep = time.sleep
        time.sleep = lambda x: None

        self.consumer = Consumer(test_huey, workers=2)
        self.consumer.create_threads()
        self.consumer.huey.queue._queue = []
        self.consumer.huey.result_store._results = {}
        self.consumer.huey._schedule = {}

        self.handler = TestLogHandler()
        logger.addHandler(self.handler)

    def tearDown(self):
        self.consumer.shutdown()
        logger.removeHandler(self.handler)
        registry._periodic_tasks = self.orig_pc
        time.sleep = self.orig_sleep

    def spawn(self, func, *args, **kwargs):
        t = threading.Thread(target=func, args=args, kwargs=kwargs)
        t.start()
        return t

    def run_worker(self, task, ts=None, pool=None, release=True):
        pool = pool or threading.Semaphore()
        ts = ts or datetime.datetime.utcnow()
        worker_t = WorkerThread(test_huey, task, ts, pool, release)
        worker_t.start()
        worker_t.join()

    def run_one_task(self, wait=True, ts=None):
        inbox = self.consumer._executor_inbox
        pool = self.consumer._pool
        pool.acquire()  # Simulate acquiring one worker, leaving one free.

        # Run through one iteration of the message loop.
        self.consumer.message_t.loop(now=ts)
        inbox.put(StopIteration)
        self.consumer.executor_t.run()

        # Block until the task is finished then release.
        if wait:
            pool.acquire()
            pool.release()

        pool.release()  # Release the fake worker acquired above.

    def test_iterable_queue(self):
        store = []
        q = IterableQueue()

        def do_queue(queue, result):
            for message in queue:
                result.append(message)

        t = self.spawn(do_queue, q, store)
        q.put(1)
        q.put(2)
        q.put(StopIteration)

        t.join()
        self.assertFalse(t.is_alive())
        self.assertEqual(store, [1, 2])

    def test_message_processing(self):
        self.consumer.message_t.start()
        self.consumer.executor_t.start()

        self.assertFalse('k' in state)

        res = modify_state('k', 'v')
        res.get(blocking=True)

        self.assertTrue('k' in state)
        self.assertEqual(res.get(), 'v')

    def test_worker(self):
        res = modify_state('k', 'w')
        task = test_huey.dequeue()
        self.run_worker(task)
        self.assertEqual(state, {'k': 'w'})

    def test_worker_exception(self):
        res = blow_up()
        task = test_huey.dequeue()

        self.run_worker(task)
        self.assertTrue(
            'Unhandled exception in worker thread' in self.handler.messages)

    def test_retries_and_logging(self):
        # this will continually fail
        res = retry_command('blampf')

        for i in reversed(range(4)):
            task = test_huey.dequeue()
            self.assertEqual(task.retries, i)
            self.run_worker(task)
            if i > 0:
                self.assertEqual(
                    self.handler.messages[-1],
                    'Re-enqueueing task %s, %s tries left' % (
                        task.task_id, i - 1))
                last_idx = -2
            else:
                last_idx = -1
            self.assertEqual(self.handler.messages[last_idx],
                             'Unhandled exception in worker thread')

        self.assertEqual(test_huey.dequeue(), None)

    def test_retries_with_success(self):
        # this will fail once, then succeed
        res = retry_command('blampf', False)
        self.assertFalse('blampf' in state)

        task = test_huey.dequeue()
        self.run_worker(task)
        self.assertEqual(self.handler.messages, [
            'Unhandled exception in worker thread',
            'Re-enqueueing task %s, 2 tries left' % task.task_id])

        task = test_huey.dequeue()
        self.assertEqual(task.retries, 2)
        self.run_worker(task)

        self.assertEqual(state['blampf'], 'fixed')
        self.assertEqual(test_huey.dequeue(), None)

    def test_pooling(self):
        # Simulate one worker being occupied.
        pool = self.consumer._pool
        pool.acquire()

        res = modify_state('x', 'y')

        # instantiate a message receiver
        self.assertEqual(len(test_queue._queue), 1)

        # process the pending message.
        self.consumer.message_t.loop(datetime.datetime.now())
        self.assertEqual(len(test_queue._queue), 0)

        # We cannot acquire the pool again until a worker releases it.
        self.assertFalse(pool.acquire(False))

        # Spawn executor and wait til it finishes the task.
        self.consumer.executor_t.start()
        res.get(blocking=True)
        self.assertEqual(state, {'x': 'y'})

        # The pool can now be acquired again.
        self.assertTrue(pool.acquire(False))

    def test_scheduling(self):
        dt = datetime.datetime(2011, 1, 1, 0, 0)
        dt2 = datetime.datetime(2037, 1, 1, 0, 0)
        r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
        r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)

        pool = self.consumer._pool
        inbox = self.consumer._executor_inbox

        # dequeue the past-timestamped task and run it.
        self.run_one_task()

        self.assertTrue('k' in state)
        self.assertEqual(test_huey._schedule, {})

        # dequeue the future-timestamped task.
        self.run_one_task(wait=False)

        # verify the task got stored in the schedule instead of executing
        self.assertFalse('k2' in state)
        self.assertEqual(len(test_huey._schedule), 1)

        # run through an iteration of the scheduler
        self.consumer.scheduler_t.loop(dt)

        # our command was not enqueued
        self.assertEqual(len(test_queue._queue), 0)
        self.assertEqual(len(test_huey._schedule), 1)

        # run through an iteration of the scheduler
        self.consumer.scheduler_t.loop(dt2)

        # our command was enqueued
        self.assertEqual(len(test_queue._queue), 1)
        self.assertEqual(len(test_huey._schedule), 0)

    def test_retry_scheduling(self):
        # this will continually fail
        res = retry_command_slow('blampf')
        self.assertEqual(test_huey._schedule, {})

        cur_time = datetime.datetime.utcnow()

        task = test_huey.dequeue()
        self.run_worker(task, ts=cur_time)
        self.assertEqual(self.handler.messages, [
            'Unhandled exception in worker thread',
            'Re-enqueueing task %s, 2 tries left' % task.task_id,
        ])

        self.assertEqual(test_huey._schedule, {task.task_id: task})
        task_from_sched = test_huey._schedule[task.task_id]
        self.assertEqual(task_from_sched.retries, 2)
        exec_time = task.execute_time

        self.assertEqual((exec_time - cur_time).seconds, 10)

    def test_schedule_persistence(self):
        # should not error out as it will be EmptyData
        test_huey.load_schedule()
        self.assertEqual(test_huey._schedule, {})

        dt = datetime.datetime(2037, 1, 1, 0, 0)
        dt2 = datetime.datetime(2037, 1, 1, 0, 1)
        r = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
        r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt2, convert_utc=False)

        # two messages in the queue
        self.assertEqual(len(test_huey.queue), 2)

        # pull 'em down
        self.run_one_task(wait=False)
        self.run_one_task(wait=False)

        test_huey.save_schedule()
        test_huey._schedule = {}
        test_huey.load_schedule()

        self.assertTrue(r.task.task_id in test_huey._schedule)
        self.assertTrue(r2.task.task_id in test_huey._schedule)

        task1 = test_huey._schedule[r.task.task_id]
        task2 = test_huey._schedule[r2.task.task_id]

        self.assertEqual(task1.execute_time, dt)
        self.assertEqual(task2.execute_time, dt2)

        # check w/conversion
        r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt)
        self.run_one_task(wait=False)

        test_huey.save_schedule()
        test_huey._schedule = {}

        test_huey.load_schedule()
        task3 = test_huey._schedule[r3.task.task_id]
        self.assertEqual(task3.execute_time, local_to_utc(dt))

    def test_revoking_normal(self):
        # enqueue 2 normal commands
        r1 = modify_state('k', 'v')
        r2 = modify_state('k2', 'v2')

        # revoke the first *before it has been checked*
        r1.revoke()
        self.assertTrue(test_huey.is_revoked(r1.task))
        self.assertFalse(test_huey.is_revoked(r2.task))

        # dequeue a *single* message (r1)
        self.run_one_task()

        # no changes and the task was not added to the schedule
        self.assertFalse('k' in state)
        self.assertEqual(test_huey._schedule, {})

        # dequeue a *single* message
        self.run_one_task()

        self.assertTrue('k2' in state)
        self.assertEqual(test_huey._schedule, {})

    def test_revoking_schedule(self):
        global state
        dt = datetime.datetime(2011, 1, 1)
        dt2 = datetime.datetime(2037, 1, 1)

        r1 = modify_state.schedule(args=('k', 'v'), eta=dt, convert_utc=False)
        r2 = modify_state.schedule(args=('k2', 'v2'), eta=dt, convert_utc=False)
        r3 = modify_state.schedule(args=('k3', 'v3'), eta=dt2, convert_utc=False)
        r4 = modify_state.schedule(args=('k4', 'v4'), eta=dt2, convert_utc=False)

        # revoke r1 and r3
        r1.revoke()
        r3.revoke()
        self.assertTrue(test_huey.is_revoked(r1.task))
        self.assertFalse(test_huey.is_revoked(r2.task))
        self.assertTrue(test_huey.is_revoked(r3.task))
        self.assertFalse(test_huey.is_revoked(r4.task))

        expected = [
            #state,        schedule
            ({},           {}),
            ({'k2': 'v2'}, {}),
            ({'k2': 'v2'}, {r3.task.task_id: r3.task}),
            ({'k2': 'v2'}, {r3.task.task_id: r3.task, r4.task.task_id: r4.task}),
        ]

        for i in range(4):
            estate, esc = expected[i]

            # dequeue a *single* message
            self.run_one_task()

            self.assertEqual(state, estate)
            self.assertEqual(test_huey._schedule, esc)

        # lets pretend its 2037
        future = dt2 + datetime.timedelta(seconds=1)
        self.consumer.scheduler_t.loop(future)
        self.assertEqual(test_huey._schedule, {})

        # There are two tasks in the queue now (r3 and r4) -- process both.
        for i in range(2):
            self.consumer.message_t.loop(now=future)
            self.consumer._executor_inbox.put(StopIteration)
            self.consumer.executor_t.run()

        self.assertEqual(state, {'k2': 'v2', 'k4': 'v4'})

    def test_revoking_periodic(self):
        global state
        def loop_periodic(ts):
            self.consumer.periodic_t.loop(ts)
            for i in range(len(test_queue._queue)):
                self.consumer.message_t.loop(ts)
                self.consumer._executor_inbox.put(StopIteration)
                self.consumer.executor_t.run()

        # revoke the command once
        every_hour.revoke(revoke_once=True)
        self.assertTrue(every_hour.is_revoked())

        # it will be skipped the first go-round
        dt = datetime.datetime(2011, 1, 1, 0, 0)
        loop_periodic(dt)

        # it has not been run
        self.assertEqual(state, {})

        # the next go-round it will be enqueued
        loop_periodic(dt)

        # our command was run
        self.assertEqual(state, {'p': 'y'})

        # reset state
        state = {}

        # revoke the command
        every_hour.revoke()
        self.assertTrue(every_hour.is_revoked())

        # it will no longer be enqueued
        loop_periodic(dt)
        loop_periodic(dt)
        self.assertEqual(state, {})

        # restore
        every_hour.restore()
        self.assertFalse(every_hour.is_revoked())

        # it will now be enqueued
        loop_periodic(dt)
        self.assertEqual(state, {'p': 'y'})

        # reset
        state = {}

        # revoke for an hour
        td = datetime.timedelta(seconds=3600)
        every_hour.revoke(revoke_until=dt + td)

        loop_periodic(dt)
        self.assertEqual(state, {})

        # after an hour it is back
        loop_periodic(dt + td)
        self.assertEqual(state, {'p': 'y'})

        # our data store should reflect the delay
        task_obj = every_hour.task_class()
        self.assertEqual(len(test_huey.result_store._results), 1)
        self.assertTrue(task_obj.revoke_id in test_huey.result_store._results)
