#    Copyright (C) 2014 Yahoo! Inc. All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import collections
import contextlib
import threading

import futurist
from oslo_utils import uuidutils
import testscenarios
import testtools

from taskflow.conductors import backends
from taskflow import engines
from taskflow.jobs.backends import impl_zookeeper
from taskflow.jobs import base
from taskflow.patterns import linear_flow as lf
from taskflow.persistence.backends import impl_memory
from taskflow import states as st
from taskflow import test
from taskflow.tests import utils as test_utils
from taskflow.utils import kazoo_utils
from taskflow.utils import persistence_utils as pu
from taskflow.utils import threading_utils


TEST_PATH_TPL = '/taskflow/conductor-test/%s'
ZOOKEEPER_AVAILABLE = test_utils.zookeeper_available(
    impl_zookeeper.ZookeeperJobBoard.MIN_ZK_VERSION
)


def test_factory(blowup):
    f = lf.Flow("test")
    if not blowup:
        f.add(test_utils.ProgressingTask('test1'))
    else:
        f.add(test_utils.FailingTask("test1"))
    return f


def sleep_factory():
    f = lf.Flow("test")
    f.add(test_utils.SleepTask('test1'))
    f.add(test_utils.ProgressingTask('test2'))
    return f


def test_store_factory():
    f = lf.Flow("test")
    f.add(test_utils.TaskMultiArg('task1'))
    return f


def single_factory():
    return futurist.ThreadPoolExecutor(max_workers=1)


ComponentBundle = collections.namedtuple(
    'ComponentBundle', ['board', 'persistence', 'conductor']
)


@testtools.skipIf(not ZOOKEEPER_AVAILABLE, 'zookeeper is not available')
class ManyConductorTest(
    testscenarios.TestWithScenarios, test_utils.EngineTestBase, test.TestCase
):
    scenarios = [
        (
            'blocking',
            {'kind': 'blocking', 'conductor_kwargs': {'wait_timeout': 0.1}},
        ),
        (
            'nonblocking_many_thread',
            {'kind': 'nonblocking', 'conductor_kwargs': {'wait_timeout': 0.1}},
        ),
        (
            'nonblocking_one_thread',
            {
                'kind': 'nonblocking',
                'conductor_kwargs': {
                    'executor_factory': single_factory,
                    'wait_timeout': 0.1,
                },
            },
        ),
    ]

    def make_components(self):
        def cleanup_path(client, path):
            if not client.connected:
                return
            client.delete(path, recursive=True)

        client = kazoo_utils.make_client(test_utils.ZK_TEST_CONFIG.copy())
        path = TEST_PATH_TPL % uuidutils.generate_uuid()
        persistence = impl_memory.MemoryBackend()
        board = impl_zookeeper.ZookeeperJobBoard(
            'testing', {'path': path}, client=client, persistence=persistence
        )

        self.addCleanup(cleanup_path, client, path)
        self.addCleanup(kazoo_utils.finalize_client, client)

        conductor_kwargs = self.conductor_kwargs.copy()
        conductor_kwargs['persistence'] = persistence
        conductor = backends.fetch(
            self.kind, 'testing', board, **conductor_kwargs
        )
        return ComponentBundle(board, persistence, conductor)

    def test_connection(self):
        components = self.make_components()
        components.conductor.connect()
        with contextlib.closing(components.conductor):
            self.assertTrue(components.board.connected)
        self.assertFalse(components.board.connected)

    def test_run_empty(self):
        components = self.make_components()
        components.conductor.connect()
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)
            t.join()

    def test_run(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()
        job_consumed_event = threading.Event()
        job_abandoned_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        def on_job_consumed(event, details):
            if event == 'job_consumed':
                job_consumed_event.set()

        def on_job_abandoned(event, details):
            if event == 'job_abandoned':
                job_abandoned_event.set()

        components.board.notifier.register(base.REMOVAL, on_consume)
        components.conductor.notifier.register("job_consumed", on_job_consumed)
        components.conductor.notifier.register(
            "job_abandoned", on_job_abandoned
        )
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, test_factory, [False], {}, backend=components.persistence
            )
            components.board.post('poke', lb, details={'flow_uuid': fd.uuid})
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            self.assertTrue(job_consumed_event.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(job_abandoned_event.wait(1))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertEqual(st.SUCCESS, fd.state)

    def test_run_max_dispatches(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        components.board.notifier.register(base.REMOVAL, on_consume)
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(
                lambda: components.conductor.run(max_dispatches=5)
            )
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, test_factory, [False], {}, backend=components.persistence
            )
            for _ in range(5):
                components.board.post(
                    'poke', lb, details={'flow_uuid': fd.uuid}
                )
                self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            components.board.post('poke', lb, details={'flow_uuid': fd.uuid})
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

    def test_fail_run(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()
        job_consumed_event = threading.Event()
        job_abandoned_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        def on_job_consumed(event, details):
            if event == 'job_consumed':
                job_consumed_event.set()

        def on_job_abandoned(event, details):
            if event == 'job_abandoned':
                job_abandoned_event.set()

        components.board.notifier.register(base.REMOVAL, on_consume)
        components.conductor.notifier.register("job_consumed", on_job_consumed)
        components.conductor.notifier.register(
            "job_abandoned", on_job_abandoned
        )
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, test_factory, [True], {}, backend=components.persistence
            )
            components.board.post('poke', lb, details={'flow_uuid': fd.uuid})
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            self.assertTrue(job_consumed_event.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(job_abandoned_event.wait(1))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertEqual(st.REVERTED, fd.state)

    def test_missing_store(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        components.board.notifier.register(base.REMOVAL, on_consume)
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, test_store_factory, [], {}, backend=components.persistence
            )
            components.board.post('poke', lb, details={'flow_uuid': fd.uuid})
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertIsNone(fd.state)

    def test_job_store(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        store = {'x': True, 'y': False, 'z': None}

        components.board.notifier.register(base.REMOVAL, on_consume)
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, test_store_factory, [], {}, backend=components.persistence
            )
            components.board.post(
                'poke', lb, details={'flow_uuid': fd.uuid, 'store': store}
            )
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertEqual(st.SUCCESS, fd.state)

    def test_flowdetails_store(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        store = {'x': True, 'y': False, 'z': None}

        components.board.notifier.register(base.REMOVAL, on_consume)
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(
                components.persistence, meta={'store': store}
            )
            engines.save_factory_details(
                fd, test_store_factory, [], {}, backend=components.persistence
            )
            components.board.post('poke', lb, details={'flow_uuid': fd.uuid})
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertEqual(st.SUCCESS, fd.state)

    def test_combined_store(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()

        def on_consume(state, details):
            consumed_event.set()

        flow_store = {'x': True, 'y': False}
        job_store = {'z': None}

        components.board.notifier.register(base.REMOVAL, on_consume)
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(
                components.persistence, meta={'store': flow_store}
            )
            engines.save_factory_details(
                fd, test_store_factory, [], {}, backend=components.persistence
            )
            components.board.post(
                'poke', lb, details={'flow_uuid': fd.uuid, 'store': job_store}
            )
            self.assertTrue(consumed_event.wait(test_utils.WAIT_TIMEOUT))
            components.conductor.stop()
            self.assertTrue(components.conductor.wait(test_utils.WAIT_TIMEOUT))
            self.assertFalse(components.conductor.dispatching)

        persistence = components.persistence
        with contextlib.closing(persistence.get_connection()) as conn:
            lb = conn.get_logbook(lb.uuid)
            fd = lb.find(fd.uuid)
        self.assertIsNotNone(fd)
        self.assertEqual(st.SUCCESS, fd.state)

    def test_stop_aborts_engine(self):
        components = self.make_components()
        components.conductor.connect()
        consumed_event = threading.Event()
        job_consumed_event = threading.Event()
        job_abandoned_event = threading.Event()
        running_start_event = threading.Event()

        def on_running_start(event, details):
            running_start_event.set()

        def on_consume(state, details):
            consumed_event.set()

        def on_job_consumed(event, details):
            if event == 'job_consumed':
                job_consumed_event.set()

        def on_job_abandoned(event, details):
            if event == 'job_abandoned':
                job_abandoned_event.set()

        components.board.notifier.register(base.REMOVAL, on_consume)
        components.conductor.notifier.register("job_consumed", on_job_consumed)
        components.conductor.notifier.register(
            "job_abandoned", on_job_abandoned
        )
        components.conductor.notifier.register(
            "running_start", on_running_start
        )
        with contextlib.closing(components.conductor):
            t = threading_utils.daemon_thread(components.conductor.run)
            t.start()
            lb, fd = pu.temporary_flow_detail(components.persistence)
            engines.save_factory_details(
                fd, sleep_factory, [], {}, backend=components.persistence
            )
            components.board.post(
                'poke',
                lb,
                details={'flow_uuid': fd.uuid, 'store': {'duration': 2}},
            )
            running_start_event.wait(test_utils.WAIT_TIMEOUT)
            components.conductor.stop()
            job_abandoned_event.wait(test_utils.WAIT_TIMEOUT)
            self.assertTrue(job_abandoned_event.is_set())
            self.assertFalse(job_consumed_event.is_set())
            self.assertFalse(consumed_event.is_set())


@testtools.skipIf(not ZOOKEEPER_AVAILABLE, 'zookeeper is not available')
class NonBlockingExecutorTest(test.TestCase):
    def test_bad_wait_timeout(self):
        persistence = impl_memory.MemoryBackend()
        board = impl_zookeeper.ZookeeperJobBoard(
            'testing',
            test_utils.ZK_TEST_CONFIG.copy(),
            persistence=persistence,
        )
        self.assertRaises(
            ValueError,
            backends.fetch,
            'nonblocking',
            'testing',
            board,
            persistence=persistence,
            wait_timeout='testing',
        )

    def test_bad_factory(self):
        persistence = impl_memory.MemoryBackend()
        board = impl_zookeeper.ZookeeperJobBoard(
            'testing',
            test_utils.ZK_TEST_CONFIG.copy(),
            persistence=persistence,
        )
        self.assertRaises(
            ValueError,
            backends.fetch,
            'nonblocking',
            'testing',
            board,
            persistence=persistence,
            executor_factory='testing',
        )
