import time

import kubernetes
from dagster_k8s.job import DagsterK8sJobConfig, construct_dagster_k8s_job, get_k8s_job_name

from dagster import DagsterInstance, Field, Noneable, StringSource, check
from dagster.core.host_representation import ExternalSchedule
from dagster.core.scheduler import DagsterSchedulerError, Scheduler
from dagster.serdes import ConfigurableClass, ConfigurableClassData
from dagster.utils import merge_dicts


class K8sScheduler(Scheduler, ConfigurableClass):
    def __init__(
        self,
        job_image,
        dagster_home,
        service_account_name,
        instance_config_map,
        postgres_password_secret,
        load_incluster_config=True,
        scheduler_namespace='default',
        image_pull_policy='Always',
        image_pull_secrets=None,
        kubeconfig_file=None,
        inst_data=None,
        env_config_maps=None,
        env_secrets=None,
    ):
        self._inst_data = check.opt_inst_param(inst_data, 'inst_data', ConfigurableClassData)

        if load_incluster_config:
            check.invariant(
                kubeconfig_file is None,
                '`kubeconfig_file` is set but `load_incluster_config` is True.',
            )
            kubernetes.config.load_incluster_config()
        else:
            check.opt_str_param(kubeconfig_file, 'kubeconfig_file')
            kubernetes.config.load_kube_config(kubeconfig_file)

        self._api = kubernetes.client.BatchV1beta1Api()
        self._namespace = check.str_param(scheduler_namespace, 'scheduler_namespace')
        self.grace_period_seconds = 5  # This should be passed in via config

        self.job_config = DagsterK8sJobConfig(
            job_image=check.str_param(job_image, 'job_image'),
            dagster_home=check.str_param(dagster_home, 'dagster_home'),
            image_pull_policy=check.str_param(image_pull_policy, 'image_pull_policy'),
            image_pull_secrets=check.opt_list_param(
                image_pull_secrets, 'image_pull_secrets', of_type=dict
            ),
            service_account_name=check.str_param(service_account_name, 'service_account_name'),
            instance_config_map=check.str_param(instance_config_map, 'instance_config_map'),
            postgres_password_secret=check.str_param(
                postgres_password_secret, 'postgres_password_secret'
            ),
            env_config_maps=check.opt_list_param(env_config_maps, 'env_config_maps', of_type=str),
            env_secrets=check.opt_list_param(env_secrets, 'env_secrets', of_type=str),
        )

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        job_cfg = DagsterK8sJobConfig.config_type()

        scheduler_extra_cfg = {
            'scheduler_namespace': Field(StringSource, is_required=True),
            'load_incluster_config': Field(bool, is_required=False, default_value=True),
            'kubeconfig_file': Field(Noneable(StringSource), is_required=False, default_value=None),
        }
        return merge_dicts(job_cfg, scheduler_extra_cfg)

    @classmethod
    def from_config_value(cls, inst_data, config_value):
        return cls(inst_data=inst_data, **config_value)

    def debug_info(self):
        return "Running K8s CronJob(s):\n{jobs}\n".format(
            jobs="\n".join([str(job) for job in self.get_all_cron_jobs()])
        )

    def wipe(self, instance):
        # Note: This method deletes schedules from ALL repositories
        check.inst_param(instance, 'instance', DagsterInstance)

        self._api.delete_collection_namespaced_cron_job(namespace=self._namespace)
        time.sleep(self.grace_period_seconds)

        # Verify that no cron jobs are running
        running_cron_job_count = len(self.get_all_cron_jobs())
        if running_cron_job_count != 0:
            raise DagsterSchedulerError(
                'Attempted to delete all K8s CronJobs but failed. There are '
                'still {} running schedules'.format(running_cron_job_count)
            )

    def _job_template(self, external_schedule):
        check.inst_param(external_schedule, 'external_schedule', ExternalSchedule)

        local_target = external_schedule.get_origin()

        job_config = self.job_config

        external_schedule_name = external_schedule.name
        job_name = get_k8s_job_name(external_schedule_name)
        pod_name = job_name

        job_template = construct_dagster_k8s_job(
            job_config=job_config,
            command=['dagster'],
            args=[
                'api',
                'launch_scheduled_execution',
                '/tmp/launch_scheduled_execution_output',  # https://bugs.python.org/issue20074 prevents using /dev/stdout
                '--schedule_name',
                external_schedule_name,
            ]
            + local_target.get_repo_cli_args().split(' '),
            job_name=job_name,
            pod_name=pod_name,
            component='scheduled_job',
        )
        return job_template

    def _start_cron_job(self, external_schedule, should_patch_schedule):
        check.inst_param(external_schedule, 'external_schedule', ExternalSchedule)
        check.bool_param(should_patch_schedule, 'should_patch_schedule')

        job_template = self._job_template(external_schedule)

        cron_job_spec = kubernetes.client.V1beta1CronJobSpec(
            schedule=external_schedule.cron_schedule, job_template=job_template
        )

        schedule_origin_id = external_schedule.get_origin_id()
        cron_job = kubernetes.client.V1beta1CronJob(
            spec=cron_job_spec, metadata={'name': schedule_origin_id}
        )

        if not should_patch_schedule:
            self._api.create_namespaced_cron_job(body=cron_job, namespace=self._namespace)
        else:
            self._api.patch_namespaced_cron_job(
                name=schedule_origin_id, body=cron_job, namespace=self._namespace
            )
        time.sleep(self.grace_period_seconds)

    def start_schedule(self, instance, external_schedule):
        check.inst_param(instance, 'instance', DagsterInstance)
        check.inst_param(external_schedule, 'external_schedule', ExternalSchedule)

        schedule_origin_id = external_schedule.get_origin_id()
        should_patch_schedule = (
            True if self.get_cron_job(schedule_origin_id=schedule_origin_id) else False
        )

        self._start_cron_job(external_schedule, should_patch_schedule)

        # Verify that the cron job is running
        cron_job = self.get_cron_job(schedule_origin_id)
        if not cron_job:
            raise DagsterSchedulerError(
                "Attempted to add K8s CronJob for schedule {schedule_name}, but failed. "
                "The schedule {schedule_name} is not running.".format(
                    schedule_name=external_schedule.name
                )
            )
        return

    def running_schedule_count(self, schedule_origin_id):
        check.str_param(schedule_origin_id, 'schedule_origin_id')

        if self.get_cron_job(schedule_origin_id):
            return 1
        else:
            return 0

    def get_cron_job(self, schedule_origin_id):
        check.str_param(schedule_origin_id, 'schedule_origin_id')

        cron_jobs = self._api.list_namespaced_cron_job(namespace=self._namespace)
        for item in cron_jobs.items:
            if schedule_origin_id == item.metadata.name:
                return item
        return None

    def get_all_cron_jobs(self):
        return self._api.list_namespaced_cron_job(namespace=self._namespace).items

    def _end_cron_job(self, schedule_origin_id):
        check.str_param(schedule_origin_id, 'schedule_origin_id')

        self._api.delete_namespaced_cron_job(name=schedule_origin_id, namespace=self._namespace)
        time.sleep(self.grace_period_seconds)

    def stop_schedule(self, instance, schedule_origin_id):
        check.inst_param(instance, 'instance', DagsterInstance)
        check.str_param(schedule_origin_id, 'schedule_origin_id')

        if self.get_cron_job(schedule_origin_id):
            self._end_cron_job(schedule_origin_id=schedule_origin_id)

        cron_job = self.get_cron_job(schedule_origin_id)
        if cron_job:
            schedule = self._get_schedule_state(instance, schedule_origin_id)

            raise DagsterSchedulerError(
                "Attempted to remove existing K8s CronJob for schedule "
                "{schedule_name}, but failed. Schedule is still running.".format(
                    schedule_name=schedule.name
                )
            )

    def get_logs_path(self, instance, schedule_origin_id):
        raise NotImplementedError('To get logs, inspect the corresponding K8s CronJob')
