#   Copyright (c) 2019 PaddlePaddle Authors. 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 cloudpickle
import multiprocessing
import os
import psutil
import signal
import socket
import subprocess
import sys
import tempfile
import time
import threading
import warnings
import zmq
from datetime import datetime
import parl
from parl.utils import get_ip_address, to_byte, to_str, logger, _IS_WINDOWS, kill_process
from parl.remote import remote_constants
from parl.remote.message import InitializedWorker
from parl.remote.status import WorkerStatus
from parl.remote.zmq_utils import create_server_socket, create_client_socket
from parl.remote.grpc_heartbeat import HeartbeatServerThread, HeartbeatClientThread
from six.moves import queue
from parl.remote.utils import get_version


class Worker(object):
    """Worker provides the cpu computation resources for the cluster.

    A worker node is connected to the master node and will send its
    computation resources information to the master node. When a worker
    node is created, it will start `cpu_num` empty jobs and these jobs'
    ip addresses will be send to the master node. Further, when an old
    job is killed, worker will start a new job and send the new job ip
    address to the master node.

    To start a worker, we use the following xparl command line api:

    .. code-block:: python

        xparl connect --address localhost:1234 --cpu_num 8

    Attributes:
        master_address (str): Master's ip address.
        request_master_socket (zmq.Context.socket): A socket which sends job
                                                    address to the master node.
        reply_job_socket (zmq.Context.socket): A socket which receives
                                               job_address from the job.
        remove_job_socket (zmq.Context.socket): A socket that receives commands to remove the job from jobs immediately.
                                                Used to remove the dead job immediately and allocate a new job
                                                instead of waiting for the heartbeat failure.
        job_buffer (str): A buffer that stores initialized jobs for providing new jobs in a short time.

    Args:
        master_address (str): IP address of the master node.
        cpu_num (int): Number of cpu to be used on the worker.
    """

    def __init__(self, master_address, cpu_num=None, log_server_port=None):
        self.lock = threading.Lock()
        self.ctx = zmq.Context.instance()
        self.master_address = master_address
        self.master_is_alive = True
        self.worker_is_alive = True
        self.worker_status = None  # initialized at `self._create_jobs`
        self._set_cpu_num(cpu_num)
        self.job_buffer = queue.Queue(maxsize=self.cpu_num)
        self._create_sockets()
        self.check_env_consistency()
        # create log server
        self.log_server_proc, self.log_server_address = self._create_log_server(
            port=log_server_port)

        # create a thread that waits commands from the job to kill the job.
        self.remove_job_thread = threading.Thread(
            target=self._reply_remove_job)
        self.remove_job_thread.setDaemon(True)
        self.remove_job_thread.start()

        self._create_jobs()

        # create a thread that initializes jobs and adds them into the job_buffer
        job_thread = threading.Thread(target=self._fill_job_buffer)
        job_thread.setDaemon(True)
        job_thread.start()

        thread = threading.Thread(target=self._update_worker_status_to_master)
        thread.setDaemon(True)
        thread.start()

    def _set_cpu_num(self, cpu_num=None):
        """set useable cpu number for worker"""
        if cpu_num is not None:
            assert isinstance(
                cpu_num, int
            ), "cpu_num should be INT type, please check the input type."
            self.cpu_num = cpu_num
        else:
            self.cpu_num = multiprocessing.cpu_count()

    def check_env_consistency(self):
        '''Verify that the parl & python version as well as some other packages in 'worker' process
            matches that of the 'master' process'''
        self.request_master_socket.send_multipart(
            [remote_constants.CHECK_VERSION_TAG])
        message = self.request_master_socket.recv_multipart()
        tag = message[0]
        if tag == remote_constants.NORMAL_TAG:
            worker_parl_version = parl.__version__
            worker_python_version_major = str(sys.version_info.major)
            worker_python_version_minor = str(sys.version_info.minor)
            assert worker_parl_version == to_str(message[1]) and worker_python_version_major == to_str(message[2])\
                and worker_python_version_minor == to_str(message[3]),\
                '''Version mismatch: the "master" is of version "parl={}, python={}.{}". However, \
                "parl={}, python={}.{}"is provided in your environment.'''.format(
                        to_str(message[1]), to_str(message[2]), to_str(message[3]),
                        worker_parl_version, worker_python_version_major, worker_python_version_minor
                    )
            worker_pyarrow_version = str(get_version('pyarrow'))
            master_pyarrow_version = to_str(message[4])
            if worker_pyarrow_version != master_pyarrow_version:
                if master_pyarrow_version == 'None':
                    error_message = """"pyarrow" is provided in your current enviroment, however, it is not \
found in "master"'s environment. To use "pyarrow" for serialization, please install \
"pyarrow={}" in "master"'s environment!""".format(worker_pyarrow_version)
                elif worker_pyarrow_version == 'None':
                    error_message = """"pyarrow" is provided in "master"'s enviroment, however, it is not \
found in your current environment. To use "pyarrow" for serialization, please install \
"pyarrow={}" in your current environment!""".format(master_pyarrow_version)
                else:
                    error_message = '''Version mismatch: the 'master' is of version 'pyarrow={}'. However, \
'pyarrow={}'is provided in your current environment.'''.format(
                        master_pyarrow_version, worker_pyarrow_version)
                raise Exception(error_message)
        else:
            raise NotImplementedError

    def _create_sockets(self):
        """Each worker maintains four sockets:

        (1) request_master_socket: sends job address to master node.
        (2) reply_job_socket: receives job_address from subprocess.
        (3) remove_job_socket : receives commands to remove the job from jobs immediately.
                                Used to remove the dead job immediately and allocate a new job
                                instead of waiting for the heartbeat failure.
        (4) reply_log_server_socket: receives log_server_heartbeat_address from subprocess.

        When a job starts, a new heartbeat socket is created to receive
        heartbeat signals from the job.

        """
        self.worker_ip = get_ip_address()

        # request_master_socket: sends job address to master
        self.request_master_socket = self.ctx.socket(zmq.REQ)
        self.request_master_socket.linger = 0

        # wait for 0.5 second to check whether master is started
        self.request_master_socket.setsockopt(zmq.RCVTIMEO, 500)
        self.request_master_socket.connect("tcp://" + self.master_address)

        # reply_job_socket: receives job_address from subprocess
        self.reply_job_socket = self.ctx.socket(zmq.REP)
        self.reply_job_socket.linger = 0
        reply_job_port = self.reply_job_socket.bind_to_random_port("tcp://*")
        self.reply_job_address = "{}:{}".format(self.worker_ip, reply_job_port)

        # remove_job_socket
        self.remove_job_socket = self.ctx.socket(zmq.REP)
        self.remove_job_socket.linger = 0
        remove_job_port = self.remove_job_socket.bind_to_random_port("tcp://*")
        self.remove_job_address = "{}:{}".format(self.worker_ip,
                                                 remove_job_port)

        # reply_log_server_socket: receives log_server_heartbeat_address from subprocess
        self.reply_log_server_socket, reply_log_server_port = create_server_socket(
            self.ctx)
        self.reply_log_server_address = "{}:{}".format(self.worker_ip,
                                                       reply_log_server_port)

    def _create_jobs(self):
        """Create jobs and send a instance of ``InitializedWorker`` that contains the worker information to the master."""
        try:
            self.request_master_socket.send_multipart(
                [remote_constants.WORKER_CONNECT_TAG])
            _ = self.request_master_socket.recv_multipart()
        except zmq.error.Again as e:
            logger.error("Can not connect to the master, "
                         "please check if master is started.")
            self.master_is_alive = False
            return

        initialized_jobs = self._init_jobs(job_num=self.cpu_num)
        self.request_master_socket.setsockopt(
            zmq.RCVTIMEO, remote_constants.HEARTBEAT_TIMEOUT_S * 1000)

        def master_heartbeat_exit_callback_func():
            logger.warning(
                "[Worker] lost connection with the master, will exit reply heartbeat for master."
            )
            if self.worker_status is not None:
                self.worker_status.clear()
            self.log_server_proc.kill()
            self.log_server_proc.wait()
            # exit the worker
            self.exit()

        self.master_heartbeat_thread = HeartbeatServerThread(
            heartbeat_exit_callback_func=master_heartbeat_exit_callback_func)
        self.master_heartbeat_thread.setDaemon(True)
        self.master_heartbeat_thread.start()
        self.master_heartbeat_address = self.master_heartbeat_thread.get_address(
        )

        logger.set_dir(
            os.path.expanduser('~/.parl_data/worker/{}'.format(
                self.master_heartbeat_address.replace(':', '_'))))
        logger.info("[Worker] Connect to the master node successfully. "
                    "({} CPUs)".format(self.cpu_num))

        for job in initialized_jobs:
            job.worker_address = self.master_heartbeat_address

        initialized_worker = InitializedWorker(self.master_heartbeat_address,
                                               initialized_jobs, self.cpu_num,
                                               socket.gethostname())
        self.request_master_socket.send_multipart([
            remote_constants.WORKER_INITIALIZED_TAG,
            cloudpickle.dumps(initialized_worker)
        ])

        _ = self.request_master_socket.recv_multipart()
        self.worker_status = WorkerStatus(self.master_heartbeat_address,
                                          initialized_jobs, self.cpu_num)

    def _fill_job_buffer(self):
        """An endless loop that adds initialized job into the job buffer"""
        initialized_jobs = []
        while self.worker_is_alive:
            if self.job_buffer.full() is False:
                job_num = self.cpu_num - self.job_buffer.qsize()
                if job_num > 0:
                    initialized_jobs = self._init_jobs(job_num=job_num)
                    for job in initialized_jobs:
                        self.job_buffer.put(job)

            time.sleep(0.02)
        self.exit()

    def _init_jobs(self, job_num):
        """Create jobs.

        Args:
            job_num(int): the number of jobs to create.
        """
        job_file = __file__.replace('worker.pyc', 'job.py')
        job_file = job_file.replace('worker.py', 'job.py')
        command = [
            sys.executable, job_file, "--worker_address",
            self.reply_job_address, "--log_server_address",
            self.log_server_address
        ]

        if sys.version_info.major == 3:
            warnings.simplefilter("ignore", ResourceWarning)

        # avoid that many jobs are killed and restarted at the same time.
        self.lock.acquire()

        # Redirect the output to DEVNULL
        FNULL = open(os.devnull, 'w')
        for _ in range(job_num):
            subprocess.Popen(command, stdout=FNULL, close_fds=True)
        FNULL.close()

        new_jobs = []
        for _ in range(job_num):
            job_message = self.reply_job_socket.recv_multipart()
            self.reply_job_socket.send_multipart([
                remote_constants.NORMAL_TAG,
                to_byte(self.remove_job_address)
            ])
            initialized_job = cloudpickle.loads(job_message[1])
            new_jobs.append(initialized_job)

            def heartbeat_exit_callback_func(job):
                job.is_alive = False
                logger.warning(
                    "[Worker] lost connection with the job:{}".format(
                        job.job_address))
                if self.master_is_alive and self.worker_is_alive:
                    self._remove_job(job.job_address)

            # a thread for sending heartbeat signals to job
            thread = HeartbeatClientThread(
                initialized_job.worker_heartbeat_address,
                heartbeat_exit_callback_func=heartbeat_exit_callback_func,
                exit_func_args=(initialized_job, ))
            thread.setDaemon(True)
            thread.start()

        self.lock.release()
        assert len(new_jobs) > 0, "init jobs failed"
        return new_jobs

    def _remove_job(self, job_address):
        """Kill a job process and update worker information"""
        success = self.worker_status.remove_job(job_address)
        if success:
            while True:
                initialized_job = self.job_buffer.get()
                initialized_job.worker_address = self.master_heartbeat_address
                if initialized_job.is_alive:
                    self.worker_status.add_job(initialized_job)
                    if not initialized_job.is_alive:  # make sure that the job is still alive.
                        self.worker_status.remove_job(
                            initialized_job.job_address)
                        continue
                else:
                    logger.warning(
                        "[Worker] a dead job found. The job buffer will not accept this one."
                    )
                if initialized_job.is_alive:
                    break

            self.lock.acquire()
            self.request_master_socket.send_multipart([
                remote_constants.NEW_JOB_TAG,
                cloudpickle.dumps(initialized_job),
                to_byte(job_address)
            ])
            _ = self.request_master_socket.recv_multipart()
            self.lock.release()

    def _reply_remove_job(self):
        """Worker starts a thread to wait jobs' commands to remove the job immediately"""
        self.remove_job_socket.linger = 0
        self.remove_job_socket.setsockopt(
            zmq.RCVTIMEO, remote_constants.HEARTBEAT_RCVTIMEO_S * 1000)
        while self.worker_is_alive and self.master_is_alive:
            try:
                message = self.remove_job_socket.recv_multipart()
                tag = message[0]
                assert tag == remote_constants.KILLJOB_TAG
                to_remove_job_address = to_str(message[1])
                self._remove_job(to_remove_job_address)
                self.remove_job_socket.send_multipart(
                    [remote_constants.NORMAL_TAG])
            except zmq.error.Again as e:
                #detect whether `self.worker_is_alive` is True periodically
                pass

    def _get_worker_status(self):
        now = datetime.strftime(datetime.now(), '%H:%M:%S')
        virtual_memory = psutil.virtual_memory()
        total_memory = round(virtual_memory[0] / (1024**3), 2)
        used_memory = round(virtual_memory[3] / (1024**3), 2)
        vacant_memory = round(total_memory - used_memory, 2)
        if _IS_WINDOWS:
            load_average = round(psutil.getloadavg()[0], 2)
        else:
            load_average = round(os.getloadavg()[0], 2)

        update_status = {
            "vacant_memory": vacant_memory,
            "used_memory": used_memory,
            "load_time": now,
            "load_value": load_average
        }
        return update_status

    def _update_worker_status_to_master(self):
        while self.master_is_alive and self.worker_is_alive:
            worker_status = self._get_worker_status()

            self.lock.acquire()
            try:
                self.request_master_socket.send_multipart([
                    remote_constants.WORKER_STATUS_UPDATE_TAG,
                    to_byte(self.master_heartbeat_address),
                    cloudpickle.dumps(worker_status)
                ])
                message = self.request_master_socket.recv_multipart()
            except zmq.error.Again as e:
                self.master_is_alive = False
            finally:
                self.lock.release()

            time.sleep(remote_constants.HEARTBEAT_INTERVAL_S)

    def _create_log_server(self, port):
        log_server_file = __file__.replace('worker.pyc', 'log_server.py')
        log_server_file = log_server_file.replace('worker.py', 'log_server.py')

        if port is None:
            port = "0"  # `0` means using a random port in flask
        command = [
            sys.executable,
            log_server_file,
            "--port",
            str(port),
            "--log_dir",
            "~/.parl_data/job/",
            "--line_num",
            "500",
            "--worker_address",
            self.reply_log_server_address,
        ]

        if sys.version_info.major == 3:
            warnings.simplefilter("ignore", ResourceWarning)

        if _IS_WINDOWS:
            FNULL = tempfile.TemporaryFile()
        else:
            FNULL = open(os.devnull, 'w')
        log_server_proc = subprocess.Popen(
            command, stdout=FNULL, close_fds=True)
        FNULL.close()

        log_server_address = "{}:{}".format(self.worker_ip, port)

        message = self.reply_log_server_socket.recv_multipart()
        log_server_heartbeat_addr = to_str(message[1])
        self.reply_log_server_socket.send_multipart(
            [remote_constants.NORMAL_TAG])

        def heartbeat_exit_callback_func():
            # only output warning
            logger.warning("[Worker] lost connection with the log_server.")

        # a thread for sending heartbeat signals to log_server
        thread = HeartbeatClientThread(
            log_server_heartbeat_addr,
            heartbeat_exit_callback_func=heartbeat_exit_callback_func)
        thread.setDaemon(True)
        thread.start()

        return log_server_proc, log_server_address

    def exit(self):
        """close the worker"""
        self.worker_is_alive = False
        kill_process('remote/job.py.*{}'.format(self.reply_job_address))
        if self.master_heartbeat_thread.is_alive():
            self.master_heartbeat_thread.exit()

    def run(self):
        """Keep running until it lost connection with the master.
        """
        if self.worker_is_alive:
            self.master_heartbeat_thread.join()
