#
# 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 base64
import cloudpickle
import sys
import threading
import traceback
from concurrent.futures import Future
from threading import Timer

# Possible job states.
PENDING = 'PENDING'
RUNNING = 'RUNNING'
CANCELLED = 'CANCELLED'
FINISHED = 'FINISHED'
SENT = 'SENT'
QUEUED = 'QUEUED'


class JobHandle(Future):

    """A child class of concurrent.futures.Future. Allows for monitoring and
        controlling of the running remote job

    """
    # Poll intervals in seconds
    _JOB_INITIAL_POLL_INTERVAL = 0.1
    _JOB_MAX_POLL_INTERVAL = 5

    def __init__(self, conn, session_id, executor):
        Future.__init__(self)
        self._conn = conn
        self._session_id = session_id
        self._executor = executor
        self._cancelled = False
        self._job_id = -1
        self._done = False
        self._job_handle_condition = threading.Condition()
        self._running_callbacks = []
        self._queued_callbacks = []

    def _start(self, command, serialized_job):
        self._executor.submit(self._send_job_task, command, serialized_job)

    def _send_job_task(self, command, job):
        suffix_url = "/" + str(self._session_id) + "/" + command
        job_status = self._conn.send_request('POST', suffix_url,
            headers=self._conn._JSON_HEADERS, data=job)
        self._job_id = job_status.json()['id']
        self._poll_result()

    def queued(self):
        """
        Returns
        -------
        True if the job is currently queued.
        """
        with self._job_handle_condition:
            return self._state == QUEUED

    def _invoke_queued_callbacks(self):
        for callback in self._queued_callbacks:
            try:
                callback(self)
            except Exception:
                traceback.print_exc()
            self._queued_callbacks.remove(callback)

    def _invoke_running_callbacks(self):
        for callback in self._running_callbacks:
            try:
                callback(self)
            except Exception:
                traceback.print_exc()
            self._running_callbacks.remove(callback)

    def add_queued_callback(self, fn):
        """
        Attaches a callable that will be called when the job is queued.

        Parameters
        ----------
        fn : Function
            A callable that will be called with this future as its only
            argument when the job is queued.
        """
        with self._job_handle_condition:
            if self._state == PENDING:
                self._queued_callbacks.append(fn)
            elif self._state == QUEUED:
                fn(self)

    def add_running_callback(self, fn):
        """
        Attaches a callable that will be called when the job is running.

        Parameters
        ----------
        fn : Function
            A callable that will be called with this future as its only
            argument when the job is running.
        """
        with self._job_handle_condition:
            if self._state in [PENDING, QUEUED]:
                self._running_callbacks.append(fn)
            elif self._state == RUNNING:
                fn(self)

    def _update_state(self, state):
        with self._job_handle_condition:
            if state == 'STARTED':
                self._state = RUNNING
                self._invoke_running_callbacks()
            elif state == 'QUEUED':
                self._state = QUEUED
                self._invoke_queued_callbacks()
            elif state == 'CANCELLED':
                self._state = CANCELLED
                self._invoke_callbacks()
            else:
                raise RuntimeError('Future in unexpected state::', self._state)
            self._job_handle_condition.notifyAll()

    def cancel(self):
        """
        Cancel the job if possible.

        Returns
        -------
        True if the job was cancelled, False if the job has already finished.
        """
        with self._job_handle_condition:
            if self._state == FINISHED:
                return False
            if self._state == CANCELLED:
                return True
            if self._job_id > -1:
                self._executor.submit(self._send_cancel_request)
            self._job_handle_condition.notify_all()
        return True

    def _send_cancel_request(self):
        try:
            end_point = "/" + str(self._session_id) + "/jobs/" + \
                str(self._job_id) + "/cancel"
            self._conn.send_json(None, end_point)
        except Exception as err:
            self.set_job_exception(err, traceback.format_exc())
            traceback.print_exc()

    def _poll_result(self):
        def do_poll_result():
            try:
                suffix_url = "/" + str(self._session_id) + "/jobs/" + \
                    str(self._job_id)
                job_status = self._conn.send_request('GET', suffix_url,
                    headers={'Accept': 'application/json'}).json()
                job_state = job_status['state']
                job_result = None
                has_finished = False
                job_error = None
                if job_state == 'SUCCEEDED':
                    job_result = job_status['result']
                    has_finished = True
                elif job_state == 'FAILED':
                    job_error = job_status['error']
                    has_finished = True
                elif job_state == 'CANCELLED':
                    repeated_timer.stop()
                else:
                    pass
                if has_finished:
                    if job_result is not None:
                        b64_decoded = base64.b64decode(job_result)
                        b64_decoded_decoded = base64.b64decode(b64_decoded)
                        deserialized_object = cloudpickle.loads(
                            b64_decoded_decoded)
                        super(JobHandle, self).set_result(deserialized_object)
                    if job_error is not None:
                        self.set_job_exception(Exception(job_error))
                    repeated_timer.stop()
                else:
                    self._update_state(job_state)
            except Exception as err:
                repeated_timer.stop()
                traceback.print_exc()
                self.set_job_exception(err, traceback.format_exc())

        repeated_timer = self._RepeatedTimer(self._JOB_INITIAL_POLL_INTERVAL,
            do_poll_result, self._executor)
        repeated_timer.start()

    def set_running_or_notify_cancel(self):
        raise NotImplementedError("This operation is not supported.")

    def set_result(self, result):
        raise NotImplementedError("This operation is not supported.")

    def set_exception_info(self, exception, traceback):
        raise NotImplementedError("This operation is not supported.")

    def set_exception(self, exception):
        raise NotImplementedError("This operation is not supported.")

    def set_job_exception(self, exception, error_msg=None):
        if sys.version >= '3':
            super(JobHandle, self).set_exception(exception)
        else:
            super(JobHandle, self).set_exception_info(exception, error_msg)

    class _RepeatedTimer(object):
        def __init__(self, interval, polling_job, executor):
            self._timer = None
            self.polling_job = polling_job
            self.interval = interval
            self.is_running = False
            self.stop_called = False
            self.executor = executor

        def _run(self):
            self.is_running = False
            self.executor.submit(self.polling_job)
            self.start()

        def start(self):
            if not self.is_running and not self.stop_called:
                self._timer = Timer(self.interval, self._run)
                self._timer.start()
                self.interval = min(self.interval * 2,
                    JobHandle._JOB_MAX_POLL_INTERVAL)
                self.is_running = True

        def stop(self):
            self._timer.cancel()
            self.stop_called = True
            self.is_running = False