# -*- coding: utf-8 -*-
#
# Copyright 2012-2015 Spotify AB
#
# 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 os
import pickle
import shutil
import tempfile
import time
from multiprocessing import Process

from helpers import unittest, with_config

import luigi.configuration
import luigi.scheduler
import luigi.server
from luigi.target import FileAlreadyExists


class SchedulerIoTest(unittest.TestCase):
    def test_pretty_id_unicode(self):
        scheduler = luigi.scheduler.Scheduler()
        scheduler.add_task(worker="A", task_id="1", params={"foo": "\u2192bar"})
        [task] = list(scheduler._state.get_active_tasks())
        task.pretty_id

    def test_load_old_state(self):
        tasks = {}
        active_workers = {"Worker1": 1e9, "Worker2": time.time()}

        with tempfile.NamedTemporaryFile(delete=True) as fn:
            with open(fn.name, "wb") as fobj:
                state = (tasks, active_workers)
                pickle.dump(state, fobj)

            state = luigi.scheduler.SimpleTaskState(state_path=fn.name)
            state.load()

            self.assertEqual(set(state.get_worker_ids()), {"Worker1", "Worker2"})

    def test_load_broken_state(self):
        with tempfile.NamedTemporaryFile(delete=True) as fn:
            with open(fn.name, "w") as fobj:
                print("b0rk", file=fobj)

            state = luigi.scheduler.SimpleTaskState(state_path=fn.name)
            state.load()  # bad if this crashes

            self.assertEqual(list(state.get_worker_ids()), [])

    @with_config({"scheduler": {"retry_count": "44", "worker_disconnect_delay": "55"}})
    def test_scheduler_with_config(self):
        scheduler = luigi.scheduler.Scheduler()
        self.assertEqual(44, scheduler._config.retry_count)
        self.assertEqual(55, scheduler._config.worker_disconnect_delay)

        # Override
        scheduler = luigi.scheduler.Scheduler(retry_count=66, worker_disconnect_delay=77)
        self.assertEqual(66, scheduler._config.retry_count)
        self.assertEqual(77, scheduler._config.worker_disconnect_delay)

    @with_config({"resources": {"a": "100", "b": "200"}})
    def test_scheduler_with_resources(self):
        scheduler = luigi.scheduler.Scheduler()
        self.assertEqual({"a": 100, "b": 200}, scheduler._resources)

    @with_config({"scheduler": {"record_task_history": "True"}, "task_history": {"db_connection": "sqlite:////none/existing/path/hist.db"}})
    def test_local_scheduler_task_history_status(self):
        ls = luigi.interface._WorkerSchedulerFactory().create_local_scheduler()
        self.assertEqual(False, ls._config.record_task_history)

    def test_load_recovers_tasks_index(self):
        scheduler = luigi.scheduler.Scheduler()
        scheduler.add_task(worker="A", task_id="1")
        scheduler.add_task(worker="B", task_id="2")
        scheduler.add_task(worker="C", task_id="3")
        scheduler.add_task(worker="D", task_id="4")
        self.assertEqual(scheduler.get_work(worker="A")["task_id"], "1")

        with tempfile.NamedTemporaryFile(delete=True) as fn:

            def reload_from_disk(scheduler):
                scheduler._state._state_path = fn.name
                scheduler.dump()
                scheduler = luigi.scheduler.Scheduler()
                scheduler._state._state_path = fn.name
                scheduler.load()
                return scheduler

            scheduler = reload_from_disk(scheduler=scheduler)
            self.assertEqual(scheduler.get_work(worker="B")["task_id"], "2")
            self.assertEqual(scheduler.get_work(worker="C")["task_id"], "3")
            scheduler = reload_from_disk(scheduler=scheduler)
            self.assertEqual(scheduler.get_work(worker="D")["task_id"], "4")

    def test_worker_prune_after_init(self):
        """
        See https://github.com/spotify/luigi/pull/1019
        """
        worker = luigi.scheduler.Worker(123)

        class TmpCfg:
            def __init__(self):
                self.worker_disconnect_delay = 10

        worker.prune(TmpCfg())

    def test_get_empty_retry_policy(self):
        retry_policy = luigi.scheduler._get_empty_retry_policy()
        self.assertEqual(3, len(retry_policy))
        self.assertEqual(["retry_count", "disable_hard_timeout", "disable_window"], list(retry_policy._asdict().keys()))
        self.assertEqual([None, None, None], list(retry_policy._asdict().values()))

    @with_config({"scheduler": {"retry_count": "9", "disable_hard_timeout": "99", "disable_window": "999"}})
    def test_scheduler_get_retry_policy(self):
        s = luigi.scheduler.Scheduler()
        self.assertEqual(luigi.scheduler.RetryPolicy(9, 99, 999), s._config._get_retry_policy())

    @with_config({"scheduler": {"retry_count": "9", "disable_hard_timeout": "99", "disable_window": "999"}})
    def test_generate_retry_policy(self):
        s = luigi.scheduler.Scheduler()

        try:
            s._generate_retry_policy({"inexist_attr": True})
            self.assertFalse(True, "'unexpected keyword argument' error must have been thrown")
        except TypeError:
            self.assertTrue(True)

        retry_policy = s._generate_retry_policy({})
        self.assertEqual(luigi.scheduler.RetryPolicy(9, 99, 999), retry_policy)

        retry_policy = s._generate_retry_policy({"retry_count": 1})
        self.assertEqual(luigi.scheduler.RetryPolicy(1, 99, 999), retry_policy)

        retry_policy = s._generate_retry_policy({"retry_count": 1, "disable_hard_timeout": 11, "disable_window": 111})
        self.assertEqual(luigi.scheduler.RetryPolicy(1, 11, 111), retry_policy)

    @with_config({"scheduler": {"retry_count": "44"}})
    def test_per_task_retry_policy(self):
        cps = luigi.scheduler.Scheduler()

        cps.add_task(worker="test_worker1", task_id="test_task_1", deps=["test_task_2", "test_task_3"])
        tasks = list(cps._state.get_active_tasks())
        self.assertEqual(3, len(tasks))

        tasks = sorted(tasks, key=lambda x: x.id)
        task_1 = tasks[0]
        task_2 = tasks[1]
        task_3 = tasks[2]

        self.assertEqual("test_task_1", task_1.id)
        self.assertEqual("test_task_2", task_2.id)
        self.assertEqual("test_task_3", task_3.id)

        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_1.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_2.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_3.retry_policy)

        cps._state._tasks = {}
        cps.add_task(
            worker="test_worker2",
            task_id="test_task_4",
            deps=["test_task_5", "test_task_6"],
            retry_policy_dict=luigi.scheduler.RetryPolicy(99, 999, 9999)._asdict(),
        )

        tasks = list(cps._state.get_active_tasks())
        self.assertEqual(3, len(tasks))

        tasks = sorted(tasks, key=lambda x: x.id)
        task_4 = tasks[0]
        task_5 = tasks[1]
        task_6 = tasks[2]

        self.assertEqual("test_task_4", task_4.id)
        self.assertEqual("test_task_5", task_5.id)
        self.assertEqual("test_task_6", task_6.id)

        self.assertEqual(luigi.scheduler.RetryPolicy(99, 999, 9999), task_4.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_5.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_6.retry_policy)

        cps._state._tasks = {}
        cps.add_task(worker="test_worker3", task_id="test_task_7", deps=["test_task_8", "test_task_9"])
        cps.add_task(worker="test_worker3", task_id="test_task_8", retry_policy_dict=luigi.scheduler.RetryPolicy(99, 999, 9999)._asdict())
        cps.add_task(worker="test_worker3", task_id="test_task_9", retry_policy_dict=luigi.scheduler.RetryPolicy(11, 111, 1111)._asdict())

        tasks = list(cps._state.get_active_tasks())
        self.assertEqual(3, len(tasks))

        tasks = sorted(tasks, key=lambda x: x.id)
        task_7 = tasks[0]
        task_8 = tasks[1]
        task_9 = tasks[2]

        self.assertEqual("test_task_7", task_7.id)
        self.assertEqual("test_task_8", task_8.id)
        self.assertEqual("test_task_9", task_9.id)

        self.assertEqual(luigi.scheduler.RetryPolicy(44, 999999999, 3600), task_7.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(99, 999, 9999), task_8.retry_policy)
        self.assertEqual(luigi.scheduler.RetryPolicy(11, 111, 1111), task_9.retry_policy)

        # Task 7 which is disable-failures 44 and its has_excessive_failures method returns False under 44
        for i in range(43):
            task_7.add_failure()
        self.assertFalse(task_7.has_excessive_failures())
        task_7.add_failure()
        self.assertTrue(task_7.has_excessive_failures())

        # Task 8 which is disable-failures 99 and its has_excessive_failures method returns False under 44
        for i in range(98):
            task_8.add_failure()
        self.assertFalse(task_8.has_excessive_failures())
        task_8.add_failure()
        self.assertTrue(task_8.has_excessive_failures())

        # Task 9 which is disable-failures 1 and its has_excessive_failures method returns False under 44
        for i in range(10):
            task_9.add_failure()
        self.assertFalse(task_9.has_excessive_failures())
        task_9.add_failure()
        self.assertTrue(task_9.has_excessive_failures())

    @with_config({"scheduler": {"record_task_history": "true"}})
    def test_has_task_history(self):
        cfg = luigi.configuration.get_config()
        with tempfile.NamedTemporaryFile(suffix=".db", delete=True) as fn:
            cfg.set("task_history", "db_connection", "sqlite:///" + fn.name)
            s = luigi.scheduler.Scheduler()
            self.assertTrue(s.has_task_history())

    @with_config({"scheduler": {"record_task_history": "false"}})
    def test_has_no_task_history(self):
        s = luigi.scheduler.Scheduler()
        self.assertFalse(s.has_task_history())

    @with_config({"scheduler": {"pause_enabled": "false"}})
    def test_pause_disabled(self):
        s = luigi.scheduler.Scheduler()
        self.assertFalse(s.is_pause_enabled()["enabled"])
        self.assertFalse(s.is_paused()["paused"])
        s.pause()
        self.assertFalse(s.is_paused()["paused"])

    def test_default_metrics_collector(self):
        from luigi.metrics import MetricsCollector

        s = luigi.scheduler.Scheduler()
        scheduler_state = s._state
        collector = scheduler_state._metrics_collector
        self.assertTrue(isinstance(collector, MetricsCollector))

    @with_config({"scheduler": {"metrics_collector": "datadog"}})
    def test_datadog_metrics_collector(self):
        from luigi.contrib.datadog_metric import DatadogMetricsCollector

        s = luigi.scheduler.Scheduler()
        scheduler_state = s._state
        collector = scheduler_state._metrics_collector
        self.assertTrue(isinstance(collector, DatadogMetricsCollector))

    @with_config({"scheduler": {"metrics_collector": "prometheus"}})
    def test_prometheus_metrics_collector(self):
        from luigi.contrib.prometheus_metric import PrometheusMetricsCollector

        s = luigi.scheduler.Scheduler()
        scheduler_state = s._state
        collector = scheduler_state._metrics_collector
        self.assertTrue(isinstance(collector, PrometheusMetricsCollector))

    @with_config({"scheduler": {"metrics_collector": "custom", "metrics_custom_import": "luigi.contrib.prometheus_metric.PrometheusMetricsCollector"}})
    def test_custom_metrics_collector(self):
        from luigi.contrib.prometheus_metric import PrometheusMetricsCollector

        s = luigi.scheduler.Scheduler()
        scheduler_state = s._state
        collector = scheduler_state._metrics_collector
        self.assertTrue(isinstance(collector, PrometheusMetricsCollector))


class SchedulerWorkerTest(unittest.TestCase):
    def get_pending_ids(self, worker, state):
        return {task.id for task in worker.get_tasks(state, "PENDING")}

    def test_get_pending_tasks_with_many_done_tasks(self):
        sch = luigi.scheduler.Scheduler()
        sch.add_task(worker="NON_TRIVIAL", task_id="A", resources={"a": 1})
        sch.add_task(worker="TRIVIAL", task_id="B", status="PENDING")
        sch.add_task(worker="TRIVIAL", task_id="C", status="DONE")
        sch.add_task(worker="TRIVIAL", task_id="D", status="DONE")

        scheduler_state = sch._state
        trivial_worker = scheduler_state.get_worker("TRIVIAL")
        self.assertEqual({"B"}, self.get_pending_ids(trivial_worker, scheduler_state))

        non_trivial_worker = scheduler_state.get_worker("NON_TRIVIAL")
        self.assertEqual({"A"}, self.get_pending_ids(non_trivial_worker, scheduler_state))


class FailingOnDoubleRunTask(luigi.Task):
    time_to_check_secs = 1
    time_to_run_secs = 2
    output_dir = luigi.Parameter(default="")

    def __init__(self, *args, **kwargs):
        super(FailingOnDoubleRunTask, self).__init__(*args, **kwargs)
        self.file_name = os.path.join(self.output_dir, "AnyTask")

    def complete(self):
        time.sleep(self.time_to_check_secs)  # e.g., establish connection
        exists = os.path.exists(self.file_name)
        time.sleep(self.time_to_check_secs)  # e.g., close connection
        return exists

    def run(self):
        time.sleep(self.time_to_run_secs)
        if os.path.exists(self.file_name):
            raise FileAlreadyExists(self.file_name)
        open(self.file_name, "w").close()


class StableDoneCooldownSecsTest(unittest.TestCase):
    def setUp(self):
        self.p = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.p)

    def run_task(self):
        return luigi.build([FailingOnDoubleRunTask(output_dir=self.p)], detailed_summary=True, parallel_scheduling=True, parallel_scheduling_processes=2)

    @with_config({"worker": {"keep_alive": "false"}})
    def get_second_run_result_on_double_run(self):
        server_process = Process(target=luigi.server.run)
        process = Process(target=self.run_task)
        try:
            # scheduler is started
            server_process.start()
            # first run is started
            process.start()
            time.sleep(FailingOnDoubleRunTask.time_to_run_secs + FailingOnDoubleRunTask.time_to_check_secs)
            # second run of the same task is started
            second_run_result = self.run_task()
            return second_run_result
        finally:
            process.join(1)
            server_process.terminate()
            server_process.join(1)

    @with_config({"scheduler": {"stable_done_cooldown_secs": "5"}})
    def test_sending_same_task_twice_with_cooldown_does_not_lead_to_double_run(self):
        second_run_result = self.get_second_run_result_on_double_run()
        self.assertEqual(second_run_result.scheduling_succeeded, True)

    @with_config({"scheduler": {"stable_done_cooldown_secs": "0"}})
    def test_sending_same_task_twice_without_cooldown_leads_to_double_run(self):
        second_run_result = self.get_second_run_result_on_double_run()
        self.assertEqual(second_run_result.scheduling_succeeded, False)
