#!/usr/bin/env python3
#
# Copyright (c) Bo Peng and the University of Texas MD Anderson Cancer Center
# Distributed under the terms of the 3-clause BSD License.
import logging
import multiprocessing as mp
import os
import re
import shlex
import sys
import tempfile
from threading import Event

import psutil
import zmq
from sos.__main__ import get_run_parser
from sos.controller import Controller, connect_controllers, disconnect_controllers
from sos.parser import SoS_Script
from sos.section_analyzer import analyze_section
from sos.syntax import SOS_SECTION_HEADER
from sos.targets import RemovedTarget, UnknownTarget, textMD5, sos_targets

from sos.utils import _parse_error, env, get_traceback, pexpect_run

from .step_executor import Interactive_Step_Executor


class NotebookLoggingHandler(logging.Handler):

    def __init__(self, level, kernel=None, title="Log Messages"):
        super(NotebookLoggingHandler, self).__init__(level)
        self.kernel = kernel
        self.title = title

    def setTitle(self, title):
        self.title = title

    def emit(self, record):
        msg = re.sub(r'``([^`]*)``', r'<span class="sos_highlight">\1</span>',
                     record.msg)
        self.kernel.send_frontend_msg(
            'display_data', {
                'metadata': {},
                'data': {
                    'text/html':
                        f'<div class="sos_logging sos_{record.levelname.lower()}">{record.levelname}: {msg}</div>'
                }
            })
        #self.kernel.send_response(self.kernel.iopub_socket, 'stream',
        #                          {'name': 'stdout', 'text': record.msg})


def start_controller(kernel):
    env.zmq_context = zmq.Context()
    # ready to monitor other workflows
    env.config['exec_mode'] = 'master'

    ready = Event()
    controller = Controller(ready, kernel)
    controller.start()
    # wait for the thread to start with a signature_req saved to env.config
    ready.wait()
    connect_controllers(env.zmq_context)
    return controller


def stop_controller(controller):
    if not controller:
        return
    env.master_request_socket.send_pyobj(['done'])
    env.master_request_socket.recv()
    disconnect_controllers()
    controller.join()


last_cell_id = None
def execute_scratch_cell(code, raw_args, kernel):
    # we then have to change the parse to disable args.workflow when
    # there is no workflow option.
    raw_args = shlex.split(raw_args) if isinstance(raw_args, str) else raw_args
    if code is None or '-h' in raw_args:
        parser = get_run_parser(interactive=True, with_workflow=True)
        parser.print_help()
        return
    if raw_args and raw_args[0].lstrip().startswith('-'):
        parser = get_run_parser(interactive=True, with_workflow=False)
        parser.error = _parse_error
        args, workflow_args = parser.parse_known_args(raw_args)
        args.workflow = None
    else:
        parser = get_run_parser(interactive=True, with_workflow=True)
        parser.error = _parse_error
        args, workflow_args = parser.parse_known_args(raw_args)

    if not code.strip():
        return

    # for reporting purpose
    sys.argv = ['%run'] + raw_args
    env.verbosity = args.verbosity

    if not any(
            isinstance(x, NotebookLoggingHandler) for x in env.logger.handlers):
        env.logger.handlers = [
            x for x in env.logger.handlers
            if type(x) is not logging.StreamHandler
        ]
        levels = {
            0: logging.ERROR,
            1: logging.WARNING,
            2: logging.INFO,
            3: logging.DEBUG,
            4: logging.DEBUG,
            None: logging.INFO
        }
        env.logger.addHandler(
            NotebookLoggingHandler(
                levels[env.verbosity], kernel, title=' '.join(sys.argv)))
    else:
        env.logger.handers[0].setTitle(' '.join(sys.argv))

    global last_cell_id
    # we retain step_input etc only when we step through a cell  #256
    if kernel and kernel.cell_id != last_cell_id:
        # clear __step_input__, __step_output__ etc because there is
        # no concept of passing input/outputs across cells.
        env.sos_dict.set('__step_output__', sos_targets([]))
        for k in [
                '__step_input__', '__default_output__', 'step_input', 'step_output',
                'step_depends', '_input', '_output', '_depends'
        ]:
            env.sos_dict.pop(k, None)

    last_cell_id = kernel.cell_id

    config = {
        'config_file': args.__config__,
        'default_queue': args.__queue__,
        'run_mode': 'dryrun' if args.dryrun else 'interactive',
        # issue 230, ignore sig mode in interactive mode
        'sig_mode': 'ignore',
        'verbosity': args.verbosity,
        # for backward compatibility, we try both args.__worker_procs__ and args.__max_procs__
        'worker_procs': args.__worker_procs__ if hasattr(args, '__worker_procs__') else args.__max_procs__,
        'max_running_jobs': args.__max_running_jobs__,
        # for infomration and resume only
        'workdir': os.getcwd(),
        'workflow': args.workflow,
        'targets': args.__targets__,
        'workflow_args': workflow_args,
        'workflow_id': textMD5(code),

        # interactive work is also a slave of the controller
        'slave_id': kernel.cell_id,
    }

    env.sos_dict.set('workflow_id', config['workflow_id'])
    env.config.update(config)

    try:
        if not any([
                SOS_SECTION_HEADER.match(line) or line.startswith('%from') or
                line.startswith('%include') for line in code.splitlines()
        ]):
            code = '[scratch_0]\n' + code
            script = SoS_Script(content=code)
        else:
            return
        workflow = script.workflow(args.workflow)
        section = workflow.sections[0]
        res = analyze_section(section)
        env.sos_dict.quick_update({
            '__signature_vars__': res['signature_vars'],
            '__environ_vars__': res['environ_vars'],
            '__changed_vars__': res['changed_vars']
        })
        executor = Interactive_Step_Executor(section, mode='interactive')
        ret = executor.run()
        try:
            return ret['__last_res__']
        except Exception as e:
            raise RuntimeError(f'Unknown result returned from executor {ret}: {e}')
    except (UnknownTarget, RemovedTarget) as e:
        raise RuntimeError(f'Unavailable target {e.target}')
    except SystemExit:
        # this happens because the executor is in resume mode but nothing
        # needs to be resumed, we simply pass
        return
    except Exception:
        env.log_to_file('PROCESS', get_traceback())
        raise


class Tapped_Executor(mp.Process):

    def __init__(self, code, args, config):
        # the worker process knows configuration file, command line argument etc
        super(Tapped_Executor, self).__init__()
        self.code = code
        self.args = args
        self.config = config

    def run(self):
        env.config.update(self.config)
        # start a socket?
        context = zmq.Context()
        stdout_socket = context.socket(zmq.PUSH)
        stdout_socket.connect(self.config["sockets"]["tapping_logging"])

        informer_socket = context.socket(zmq.PUSH)
        informer_socket.connect(self.config["sockets"]["tapping_listener"])

        try:
            filename = os.path.join(tempfile.gettempdir(),
                                    f'interactive_{os.getpid()}.sos')
            with open(filename, 'w') as script_file:
                script_file.write(self.code)

            cmd = f'sos run {filename} {self.args} -m tapping slave {self.config["slave_id"]} {self.config["sockets"]["tapping_logging"]} {self.config["sockets"]["tapping_listener"]} {self.config["sockets"]["tapping_controller"]}'
            ret_code = pexpect_run(cmd, shell=True, stdout_socket=stdout_socket)
            # status will not trigger frontend update if it was not
            # started with a pending status
            informer_socket.send_pyobj({
                'msg_type': 'workflow_status',
                'data': {
                    'cell_id': env.config['slave_id'],
                    'status': 'completed' if ret_code == 0 else 'failed'
                }
            })
        except Exception as e:
            stdout_socket.send_multipart([b'ERROR', str(e).encode()])
            informer_socket.send_pyobj({
                'msg_type': 'workflow_status',
                'data': {
                    'cell_id': env.config['slave_id'],
                    'status': 'failed',
                    'exception': str(e)
                }
            })
        finally:
            stdout_socket.LINGER = 0
            stdout_socket.close()
            informer_socket.LINGER = 0
            informer_socket.close()
            context.term()


# workflow queue that holds all workflow
g_workflow_queue = []


def run_next_workflow_in_queue():
    # execute the first available item
    global g_workflow_queue

    for idx, (cid, proc) in enumerate(g_workflow_queue):
        if proc is None:
            continue
        # this is ordered
        elif isinstance(proc, tuple):
            env.config['slave_id'] = cid
            executor = Tapped_Executor(*proc)
            executor.start()
            g_workflow_queue[idx][1] = executor
            break
        elif not (proc.is_alive() and psutil.pid_exists(proc.pid)):
            g_workflow_queue[idx][1] = None
            continue
        else:
            # if already running, do not submit new one
            break


def execute_pending_workflow(cell_ids, kernel):
    # we are giving a list of cell_ids because some cells might be removed
    # we use this list to clear workflow queue of removed cells
    for idx, (cid, proc) in enumerate(g_workflow_queue):
        if isinstance(proc, tuple) and cid not in cell_ids:
            g_workflow_queue[idx][1] = None
    run_next_workflow_in_queue()


def run_sos_workflow(code,
                     raw_args='',
                     kernel=None,
                     workflow_mode=False,
                     run_in_queue=False):
    # when user asks to execute a cell as workflow. We either
    # execute the workflow or put it in queue
    global g_workflow_queue

    # if a cell already exist, remove previous pending job (a tuple)
    # completed job (dead process), or running job.
    if kernel.cell_id in [
            cid for cid, proc in g_workflow_queue if proc is not None
    ]:
        cancel_workflow(kernel.cell_id, kernel)

    if run_in_queue:
        # put to the back
        g_workflow_queue.append([kernel.cell_id, (code, raw_args, env.config)])
        # in any case, we start with a pending status
        kernel.send_frontend_msg(
            'workflow_status', {
                'cell_id': kernel.cell_id,
                'status': 'pending',
                'index': len(g_workflow_queue)
            })
        run_next_workflow_in_queue()
    else:
        env.config['slave_id'] = kernel.cell_id
        executor = Tapped_Executor(code, raw_args, env.config)
        executor.start()
        executor.join()


def cancel_workflow(cell_id, kernel):
    global g_workflow_queue
    env.logger.info(f'A queued or running workflow in this cell is canceled')
    kernel.send_frontend_msg('workflow_status', {
        'cell_id': cell_id,
        'status': 'purged'
    })

    for idx, (cid, proc) in enumerate(g_workflow_queue):
        if cid != cell_id or proc is None:
            continue
        if not isinstance(proc, tuple) and \
            (proc.is_alive() and psutil.pid_exists(proc.pid)):
            from sos.executor_utils import kill_all_subprocesses
            kill_all_subprocesses(proc.pid, include_self=True)
            proc.terminate()
            if psutil.pid_exists(proc.pid):
                raise RuntimeError('Failed to kill workflow')
        g_workflow_queue[idx][1] = None
