# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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 datetime
import logging
import os
import shutil
from abc import ABC
from tempfile import NamedTemporaryFile
from typing import Dict, Text, List, Optional

import cloudpickle
from airflow.exceptions import DagNotFound

from ai_flow_plugins.scheduler_plugins.airflow.dag_generator import DAGGenerator
from ai_flow.context.project_context import ProjectContext
from ai_flow.plugin_interface.scheduler_interface import Scheduler, \
    WorkflowInfo, JobExecutionInfo, WorkflowExecutionInfo
from ai_flow_plugins.scheduler_plugins.airflow.airflow_restful_util import AirFlowRestfulUtil
from ai_flow.workflow.workflow import Workflow
from ai_flow.workflow import status
from ai_flow.util.time_utils import datetime_to_int64
from airflow.executors.scheduling_action import SchedulingAction
from airflow.models.taskexecution import TaskExecution
from airflow.models.dag import DagModel
from airflow.models.dagrun import DagRun
from airflow.utils.db import create_session
from airflow.utils.state import State
from airflow.contrib.jobs.scheduler_client import EventSchedulerClient, SCHEDULER_NAMESPACE, ExecutionContext
from airflow.utils.dates import parse_execution_date
from airflow.api.common.experimental import delete_dag

logger = logging.getLogger(__name__)

TIMEOUT = 5 * 60


class AirFlowSchedulerBase(Scheduler, ABC):
    """
    AirFlowScheduler is an implementation of a Scheduler interface based on AirFlow.
    AirFlowScheduler contains two configuration items:
    1. notification_server_uri: The address of NotificationService.
    2. airflow_deploy_path: AirFlow dag file deployment directory.
    """

    def __init__(self, config: Dict):
        if 'notification_server_uri' not in config:
            raise Exception('`notification_server_uri` option of scheduler config is not configured. '
                            'Please add the `notification_server_uri` option under `scheduler_config` option!')
        if 'airflow_deploy_path' not in config:
            from airflow import settings
            dags_folder = settings.DAGS_FOLDER
            config['airflow_deploy_path'] = dags_folder
            logger.info("airflow_deploy_path is set to {} from airflow config".format(dags_folder))
        super().__init__(config)
        self.dag_generator = DAGGenerator()
        self._airflow_client = None

    @classmethod
    def airflow_dag_id(cls, namespace, workflow_name):
        return '{}.{}'.format(namespace, workflow_name)

    @classmethod
    def dag_id_to_namespace_workflow(cls, dag_id: Text):
        tmp = dag_id.split('.')
        return tmp[0], tmp[1]

    @classmethod
    def generate_workflow_execution_id(cls, dag_id, run_id) -> Text:
        return '{}|{}'.format(dag_id, run_id)

    @classmethod
    def generate_job_execution_id(cls, dag_id, run_id, task_id, seq_num) -> Text:
        return '{}|{}|{}|{}'.format(dag_id, run_id, task_id, seq_num)

    @classmethod
    def parse_dag_id_and_run_id(cls, workflow_execution_id: Text) -> (Text, Text):
        return workflow_execution_id.split('|')

    @classmethod
    def airflow_state_to_status(cls, state) -> status.Status:
        if State.SUCCESS == state:
            return status.Status.FINISHED
        elif State.FAILED == state:
            return status.Status.FAILED
        elif State.RUNNING == state:
            return status.Status.RUNNING
        elif State.KILLED == state or State.SHUTDOWN == state \
                or State.KILLING == state:
            # We map airflow state KILLING to KILLED in the assumption that KILLING is a transient state,
            # and it is the best we can do.
            return status.Status.KILLED
        else:
            return status.Status.INIT

    @classmethod
    def status_to_airflow_state(cls, status_: status.Status) -> Text:
        if status.Status.FINISHED == status_:
            return State.SUCCESS
        elif status.Status.FAILED == status_:
            return State.FAILED
        elif status.Status.RUNNING == status_:
            return State.RUNNING
        elif status.Status.KILLED == status_:
            return State.KILLED
        else:
            return State.NONE

    @property
    def airflow_client(self):
        if self._airflow_client is None:
            self._airflow_client = EventSchedulerClient(
                notification_server_uri=self.config.get('notification_server_uri'),
                namespace=SCHEDULER_NAMESPACE)
        return self._airflow_client

    def submit_workflow(self, workflow: Workflow, context_extractor, project_context: ProjectContext) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_context.project_name, workflow.workflow_name)
        self._check_configurable_path('resource_dir')
        code_text = self.dag_generator.generate(workflow=workflow,
                                                project_name=project_context.project_name,
                                                resource_dir=self.config.get('resource_dir'),
                                                base_log_folder=self.config.get('base_log_folder'),
                                                context_extractor=context_extractor)
        self._check_configurable_path('airflow_deploy_path')
        airflow_file_path = self._write_to_deploy_path(code_text, dag_id + ".py",
                                                       self.config.get('airflow_deploy_path'))

        self.airflow_client.trigger_parse_dag(airflow_file_path, timeout=TIMEOUT)
        return WorkflowInfo(namespace=project_context.project_name,
                            workflow_name=workflow.workflow_name,
                            properties={'dag_file': airflow_file_path})

    @staticmethod
    def _write_to_deploy_path(content, filename, deploy_path, mode='w+t'):
        airflow_file_path = os.path.join(deploy_path, filename)
        if os.path.exists(airflow_file_path):
            os.remove(airflow_file_path)
        with NamedTemporaryFile(mode=mode, prefix=filename, dir='/tmp', delete=False) as f:
            f.write(content)
        shutil.move(f.name, airflow_file_path)
        return airflow_file_path

    def delete_workflow(self, project_name: Text, workflow_name: Text) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        deploy_path = self.config.get('airflow_deploy_path')
        if deploy_path is None:
            raise Exception("airflow_deploy_path config not set!")

        self.pause_workflow_scheduling(project_name=project_name,
                                       workflow_name=workflow_name)
        self.stop_all_workflow_execution(project_name=project_name,
                                         workflow_name=workflow_name)
        try:
            delete_dag.delete_dag(dag_id)
        except DagNotFound:
            raise Exception("DAG with id {} not found. Cannot delete".format(dag_id))

        airflow_file_path = os.path.join(deploy_path,
                                         dag_id + '.py')
        if os.path.exists(airflow_file_path):
            os.remove(airflow_file_path)
            return WorkflowInfo(namespace=project_name,
                                workflow_name=workflow_name,
                                properties={'dag_file': airflow_file_path})
        else:
            return None

    def _check_configurable_path(self, config_str):
        config_path = self.config.get(config_str)
        if config_path is None:
            raise Exception("{} config not set!".format(config_str))
        if not os.path.exists(config_path):
            os.makedirs(config_path)

    def stop_scheduling_job(self, workflow_execution_id: Text, job_name: Text):
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        self.airflow_client.stop_scheduling_task(dag_id=dag_id,
                                                 dagrun_id=run_id,
                                                 task_id=job_name,
                                                 timeout=TIMEOUT)

    def resume_scheduling_job(self, workflow_execution_id: Text, job_name: Text):
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        self.airflow_client.resume_scheduling_task(dag_id=dag_id,
                                                   dagrun_id=run_id,
                                                   task_id=job_name,
                                                   timeout=TIMEOUT)


class AirFlowScheduler(AirFlowSchedulerBase):
    """
    AirFlowScheduler is an implementation of a Scheduler interface based on AirFlow.
    AirFlowScheduler contains two configuration items:
    1. notification_server_uri: The address of server of NotificationService.
    2. airflow_deploy_path: AirFlow dag file deployment directory.
    """

    def __init__(self, config: Dict):
        super().__init__(config)

    @classmethod
    def dag_exist(cls, dag_id):
        with create_session() as session:
            dag = session.query(DagModel).filter(DagModel.dag_id == dag_id).first()
            if dag is None:
                return False
            else:
                return True

    def pause_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        DagModel.get_dagmodel(dag_id=dag_id).set_is_paused(is_paused=True)
        return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)

    def resume_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        DagModel.get_dagmodel(dag_id=dag_id).set_is_paused(is_paused=False)
        return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)

    def start_new_workflow_execution(self, project_name: Text, workflow_name: Text,
                                     workflow_execution_context: Text = None) \
            -> Optional[WorkflowExecutionInfo]:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        deploy_path = self.config.get('airflow_deploy_path')
        if deploy_path is None:
            raise Exception("airflow_deploy_path config not set!")
        if not self.dag_exist(dag_id):
            return None
        context: ExecutionContext = self.airflow_client.schedule_dag(dag_id, workflow_execution_context, timeout=TIMEOUT)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(session=session, dag_id=dag_id, run_id=context.dagrun_id)
            if dagrun is None:
                return None
            else:
                return WorkflowExecutionInfo(
                    workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name),
                    workflow_execution_id=self.generate_workflow_execution_id(dag_id=dag_id, run_id=dagrun.run_id),
                    status=status.Status.INIT)

    def stop_all_workflow_execution(self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]:
        workflow_execution_list = self.list_workflow_executions(project_name, workflow_name)
        for we in workflow_execution_list:
            if we.status == status.Status.RUNNING:
                self.stop_workflow_execution(we.workflow_execution_id)
        return workflow_execution_list

    def stop_workflow_execution_by_context(self, workflow_name: Text, context: Text) -> Optional[WorkflowExecutionInfo]:
        # TODO: impl
        pass

    def stop_workflow_execution(self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
            if dagrun.state in State.finished:
                return WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                        workflow_name=workflow_name),
                                             workflow_execution_id=workflow_execution_id,
                                             status=self.airflow_state_to_status(dagrun.state))
            context: ExecutionContext = ExecutionContext(dagrun_id=dagrun.run_id)
            current_context = self.airflow_client.stop_dag_run(dagrun.dag_id, context, timeout=TIMEOUT)
            return WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                    workflow_name=workflow_name),
                                         workflow_execution_id=workflow_execution_id,
                                         status=status.Status.KILLED)

    def get_workflow_execution(self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            else:
                status_ = self.airflow_state_to_status(dagrun.state)
                project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
                return WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                        workflow_name=workflow_name),
                                             workflow_execution_id=workflow_execution_id,
                                             status=status_,
                                             start_date=str(datetime_to_int64(dagrun.start_date)),
                                             end_date=str(datetime_to_int64(dagrun.end_date))
                                             )

    def list_workflow_executions(self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        with create_session() as session:
            dagrun_list = session.query(DagRun).filter(DagRun.dag_id == dag_id).all()
            if dagrun_list is None:
                return []
            else:
                result = []
                for dagrun in dagrun_list:
                    status_ = self.airflow_state_to_status(dagrun.state)
                    result.append(WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                   workflow_name=workflow_name),
                                                        workflow_execution_id
                                                        =self.generate_workflow_execution_id(dag_id, dagrun.run_id),
                                                        status=status_,
                                                        start_date=str(datetime_to_int64(dagrun.start_date)),
                                                        end_date=str(datetime_to_int64(dagrun.end_date)),
                                                        ))
                return result

    def start_job_execution(self, job_name: Text, workflow_execution_id: Text) -> Optional[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            if dagrun.state != State.RUNNING:
                raise Exception('execution: {} state: {} can not trigger job.'.format(workflow_execution_id,
                                                                                      dagrun.state))
            task = dagrun.get_task_instance(job_name, session)
            if task is None:
                return None
            if task.state in State.unfinished:
                raise Exception('job:{} state: {} can not start!'.format(job_name, task.state))
            self.airflow_client.schedule_task(dag_id=dagrun.dag_id,
                                              task_id=job_name,
                                              action=SchedulingAction.START,
                                              context=ExecutionContext(dagrun_id=dagrun.run_id),
                                              timeout=TIMEOUT)
            project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
            return JobExecutionInfo(job_name=job_name,
                                    status=self.airflow_state_to_status(task.state),
                                    workflow_execution
                                    =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                      workflow_name=workflow_name),
                                                           workflow_execution_id=workflow_execution_id,
                                                           status=self.airflow_state_to_status(dagrun.state)))

    def stop_job_execution(self, job_name: Text, workflow_execution_id: Text) -> Optional[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            task = dagrun.get_task_instance(job_name, session)
            if task is None:
                return None
            if task.state in State.finished:
                raise Exception('job:{} state: {} can not stop!'.format(job_name, task.state))
            else:
                self.airflow_client.schedule_task(dag_id=dagrun.dag_id,
                                                  task_id=job_name,
                                                  action=SchedulingAction.STOP,
                                                  context=ExecutionContext(dagrun_id=dagrun.run_id),
                                                  timeout=TIMEOUT)
            project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
            return JobExecutionInfo(job_name=job_name,
                                    status=self.airflow_state_to_status(task.state),
                                    workflow_execution
                                    =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                      workflow_name=workflow_name),
                                                           workflow_execution_id=workflow_execution_id,
                                                           status=self.airflow_state_to_status(dagrun.state)))

    def restart_job_execution(self, job_name: Text, workflow_execution_id: Text) -> Optional[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            if dagrun.state != State.RUNNING:
                raise Exception('execution: {} state: {} can not trigger job.'.format(workflow_execution_id,
                                                                                      dagrun.state))
            task = dagrun.get_task_instance(job_name, session)
            if task is None:
                return None
            self.airflow_client.schedule_task(dag_id=dagrun.dag_id,
                                              task_id=job_name,
                                              action=SchedulingAction.RESTART,
                                              context=ExecutionContext(dagrun_id=dagrun.run_id),
                                              timeout=TIMEOUT)
            project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
            return JobExecutionInfo(job_name=job_name,
                                    status=self.airflow_state_to_status(task.state),
                                    workflow_execution
                                    =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                      workflow_name=workflow_name),
                                                           workflow_execution_id=workflow_execution_id,
                                                           status=self.airflow_state_to_status(dagrun.state)))

    def get_job_executions(self, job_name: Text, workflow_execution_id: Text) -> List[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            task_list = session.query(TaskExecution).filter(TaskExecution.dag_id == dagrun.dag_id,
                                                            TaskExecution.execution_date == dagrun.execution_date,
                                                            TaskExecution.task_id == job_name).all()
            if task_list is None:
                return []
            else:
                result = self.build_job_execution_info_list(dagrun, task_list)
                return result

    def build_job_execution_info_list(self, dagrun, task_list):
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.dag_id)
        result = []
        for task in task_list:
            job = JobExecutionInfo(job_execution_id=self.generate_job_execution_id(dag_id=dagrun.dag_id,
                                                                                   run_id=dagrun.run_id,
                                                                                   task_id=task.task_id,
                                                                                   seq_num=task.seq_num),
                                   job_name=task.task_id,
                                   status=self.airflow_state_to_status(task.state),
                                   start_date=str(datetime_to_int64(task.start_date)),
                                   end_date=str(datetime_to_int64(task.end_date)),
                                   execution_label=task.execution_label,
                                   workflow_execution
                                   =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                     workflow_name=workflow_name),
                                                          workflow_execution_id=self.generate_workflow_execution_id(
                                                              dag_id=dagrun.dag_id,
                                                              run_id=dagrun.run_id),
                                                          status=self.airflow_state_to_status(dagrun.state)))
            result.append(job)
        return result

    def list_job_executions(self, workflow_execution_id: Text) -> List[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        with create_session() as session:
            dagrun = DagRun.get_run_by_id(dag_id=dag_id, run_id=run_id, session=session)
            if dagrun is None:
                return None
            task_list = session.query(TaskExecution).filter(TaskExecution.dag_id == dagrun.dag_id,
                                                            TaskExecution.execution_date == dagrun.execution_date).all()
            if task_list is None:
                return []
            else:
                result = self.build_job_execution_info_list(dagrun, task_list)
                return result


class AirFlowSchedulerRestful(AirFlowSchedulerBase):

    def __init__(self, config: Dict):
        super().__init__(config)
        if 'endpoint_url' not in config:
            raise Exception('`endpoint_url` option of scheduler config is not configured. '
                            'Please add the `endpoint_url` option under `scheduler_config` option!')
        if 'username' not in config:
            raise Exception('`username` option of scheduler config is not configured. '
                            'Please add the `username` option under `scheduler_config` option!')
        if 'password' not in config:
            raise Exception('`password` option of scheduler config is not configured. '
                            'Please add the `password` option under `scheduler_config` option!')
        self.restful_util: AirFlowRestfulUtil = AirFlowRestfulUtil(endpoint_url=config.get("endpoint_url"),
                                                                   user_name=config.get("username"),
                                                                   password=config.get("password"))

    @classmethod
    def datetime_str_to_int64_str(cls, datetime_str):
        if datetime_str is None:
            return '0'
        else:
            return str(datetime_to_int64(parse_execution_date(datetime_str)))

    def pause_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        self.restful_util.set_dag_is_paused(dag_id=dag_id, is_paused=True)
        return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)

    def resume_workflow_scheduling(self, project_name: Text, workflow_name: Text) -> WorkflowInfo:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        self.restful_util.set_dag_is_paused(dag_id=dag_id, is_paused=False)
        return WorkflowInfo(namespace=project_name, workflow_name=workflow_name)

    def start_new_workflow_execution(self, project_name: Text,
                                     workflow_name: Text,
                                     context: Text = None) -> Optional[WorkflowExecutionInfo]:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        deploy_path = self.config.get('airflow_deploy_path')
        if deploy_path is None:
            raise Exception("airflow_deploy_path config not set!")
        if not self.restful_util.dag_exist(dag_id):
            return None
        context: ExecutionContext = self.airflow_client.schedule_dag(dag_id, context, timeout=TIMEOUT)
        dagrun = self.restful_util.get_dagrun(dag_id=dag_id, run_id=context.dagrun_id)
        if dagrun is None:
            return None
        else:
            return WorkflowExecutionInfo(
                workflow_info=WorkflowInfo(namespace=project_name, workflow_name=workflow_name),
                workflow_execution_id=self.generate_workflow_execution_id(dag_id, context.dagrun_id),
                status=status.Status.INIT)

    def stop_all_workflow_execution(self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        dagrun_list = self.restful_util.list_dagruns(dag_id)
        result = []
        for dagrun in dagrun_list:
            status_ = self.airflow_state_to_status(dagrun.get('state'))
            workflow_execution_id = self.generate_workflow_execution_id(dag_id, dagrun.get('dag_run_id'))
            if status_ == State.RUNNING:
                self.stop_workflow_execution(workflow_execution_id)
            result.append(WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                           workflow_name=workflow_name),
                                                workflow_execution_id=workflow_execution_id,
                                                status=status_,
                                                start_date=self.datetime_str_to_int64_str(dagrun.get('start_date')),
                                                end_date=self.datetime_str_to_int64_str(dagrun.get('end_date'))
                                                ))
        return result

    def stop_workflow_execution(self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id=dag_id, run_id=run_id)
        if dagrun is None:
            return None
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dag_id)
        context: ExecutionContext = ExecutionContext(dagrun_id=run_id)
        current_context = self.airflow_client.stop_dag_run(dag_id, context, timeout=TIMEOUT)
        return WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                workflow_name=workflow_name),
                                     workflow_execution_id=workflow_execution_id,
                                     status=status.Status.KILLED)

    def stop_workflow_execution_by_context(self, workflow_name: Text, context: Text) -> Optional[WorkflowExecutionInfo]:
        # todo Need to implement the function
        raise NotImplementedError('Does not implement the top_workflow_execution_by_context function.')

    def get_workflow_execution(self, workflow_execution_id: Text) -> Optional[WorkflowExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id=dag_id, run_id=run_id)
        print(dagrun)
        if dagrun is None:
            return None
        else:
            status_ = self.airflow_state_to_status(dagrun.get('state'))
            project_name, workflow_name = self.dag_id_to_namespace_workflow(dag_id)
            return WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                    workflow_name=workflow_name),
                                         workflow_execution_id=workflow_execution_id,
                                         status=status_,
                                         start_date=self.datetime_str_to_int64_str(dagrun.get('start_date')),
                                         end_date=self.datetime_str_to_int64_str(dagrun.get('end_date'))
                                         )

    def list_workflow_executions(self, project_name: Text, workflow_name: Text) -> List[WorkflowExecutionInfo]:
        dag_id = self.airflow_dag_id(project_name, workflow_name)
        dagrun_list = self.restful_util.list_dagruns(dag_id)
        result = []
        for dagrun in dagrun_list:
            status_ = self.airflow_state_to_status(dagrun.get('state'))
            workflow_execution_id = self.generate_workflow_execution_id(dag_id, dagrun.get('dag_run_id'))
            result.append(WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                           workflow_name=workflow_name),
                                                workflow_execution_id=workflow_execution_id,
                                                status=status_,
                                                start_date=self.datetime_str_to_int64_str(dagrun.get('start_date')),
                                                end_date=self.datetime_str_to_int64_str(dagrun.get('end_date')),
                                                ))
        return result

    def start_job_execution(self, job_name: Text, workflow_execution_id: Text) -> JobExecutionInfo:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id, run_id)
        if dagrun is None:
            return None
        if dagrun.get('state') != State.RUNNING:
            raise Exception('execution: {} state: {} can not trigger job.'.format(workflow_execution_id,
                                                                                  dagrun.get('state')))
        task = self.restful_util.get_task_instance(dag_id, run_id, job_name)
        if task is None:
            return None
        if task.get('state') in State.running:
            raise Exception('job:{} state: {} can not start!'.format(job_name, task.get('state')))
        self.airflow_client.schedule_task(dag_id=dag_id,
                                          task_id=job_name,
                                          action=SchedulingAction.START,
                                          context=ExecutionContext(dagrun_id=run_id),
                                          timeout=TIMEOUT)
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dag_id)
        return JobExecutionInfo(job_name=job_name,
                                status=self.airflow_state_to_status(task.get('state')),
                                workflow_execution
                                =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                  workflow_name=workflow_name),
                                                       workflow_execution_id=workflow_execution_id,
                                                       status=self.airflow_state_to_status(dagrun.get('state'))))

    def stop_job_execution(self, job_name: Text, workflow_execution_id: Text) -> JobExecutionInfo:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id, run_id)
        if dagrun is None:
            return None
        task = self.restful_util.get_task_instance(dag_id, run_id, job_name)
        if task is None:
            return None
        if task.get('state') in State.finished:
            raise Exception('job:{} state: {} can not stop!'.format(job_name, task.get('state')))
        else:
            self.airflow_client.schedule_task(dag_id=dag_id,
                                              task_id=job_name,
                                              action=SchedulingAction.STOP,
                                              context=ExecutionContext(dagrun_id=run_id),
                                              timeout=TIMEOUT)
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dag_id)
        return JobExecutionInfo(job_name=job_name,
                                status=self.airflow_state_to_status(task.get('state')),
                                workflow_execution
                                =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                  workflow_name=workflow_name),
                                                       workflow_execution_id=workflow_execution_id,
                                                       status=self.airflow_state_to_status(dagrun.get('state'))))

    def restart_job_execution(self, job_name: Text, workflow_execution_id: Text) -> JobExecutionInfo:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id, run_id)
        if dagrun is None:
            return None
        if dagrun.get('state') != State.RUNNING:
            raise Exception('execution: {} state: {} can not trigger job.'.format(workflow_execution_id,
                                                                                  dagrun.get('state')))
        task = self.restful_util.get_task_instance(dag_id, run_id, job_name)
        if task is None:
            return None
        self.airflow_client.schedule_task(dag_id=dag_id,
                                          task_id=job_name,
                                          action=SchedulingAction.RESTART,
                                          context=ExecutionContext(dagrun_id=run_id),
                                          timeout=TIMEOUT)
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dag_id)
        return JobExecutionInfo(job_name=job_name,
                                status=self.airflow_state_to_status(task.get('state')),
                                workflow_execution
                                =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                  workflow_name=workflow_name),
                                                       workflow_execution_id=workflow_execution_id,
                                                       status=self.airflow_state_to_status(dagrun.get('state'))))

    def build_job_execution_info_list(self, dagrun, task_list):
        project_name, workflow_name = self.dag_id_to_namespace_workflow(dagrun.get('dag_id'))
        workflow_execution_id = self.generate_workflow_execution_id(dagrun.get('dag_id'), dagrun.get('dag_run_id'))
        result = []
        for task in task_list:
            job = JobExecutionInfo(job_execution_id=self.generate_job_execution_id(dag_id=dagrun.get('dag_id'),
                                                                                   run_id=dagrun.get('dag_run_id'),
                                                                                   task_id=task.get('task_id'),
                                                                                   seq_num=task.get('seq_num')),
                                   job_name=task.get('task_id'),
                                   status=self.airflow_state_to_status(task.get('state')),
                                   start_date=self.datetime_str_to_int64_str(task.get('start_date')),
                                   end_date=self.datetime_str_to_int64_str(task.get('end_date')),
                                   execution_label=task.execution_label,
                                   workflow_execution
                                   =WorkflowExecutionInfo(workflow_info=WorkflowInfo(namespace=project_name,
                                                                                     workflow_name=workflow_name),
                                                          workflow_execution_id=workflow_execution_id,
                                                          status=self.airflow_state_to_status(dagrun.get('state'))))
            result.append(job)
        return result

    def get_job_executions(self, job_name: Text, workflow_execution_id: Text) -> List[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id, run_id)
        if dagrun is None:
            return None
        tasks = self.restful_util.list_task_execution_by_task_id(dag_id, run_id, job_name)
        if tasks is None:
            return []
        else:
            return self.build_job_execution_info_list(dagrun, tasks)

    def list_job_executions(self, workflow_execution_id: Text) -> List[JobExecutionInfo]:
        dag_id, run_id = self.parse_dag_id_and_run_id(workflow_execution_id)
        dagrun = self.restful_util.get_dagrun(dag_id, run_id)
        if dagrun is None:
            return None
        task_list = self.restful_util.list_task_execution(dag_id, run_id)
        if task_list is None:
            return []
        else:
            result = self.build_job_execution_info_list(dagrun, task_list)
            return result
