#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
#
# run_info.py: information about a run (kept on the controller)
import os
import sys
import copy 
import rpyc
import time 
import json
import psutil
import logging
from threading import Lock

from xtlib import mirror_worker, utils
from xtlib import errors
from xtlib import file_utils
from xtlib.utils import log_info, log_title

from xtlib.console import console

logger = logging.getLogger(__name__)

class RunInfo():
    def __init__(self, run_name, workspace, cmd_parts, run_script, repeat, context, status, show_output=True, parent_name=None, 
        parent_prep_needed=False, mirror_close_func=None, node_id=None, run_index=None, max_delay=0,
        store=None):

        if isinstance(repeat, str):
            repeat = int(repeat)

        self.run_name = run_name
        self.parent_name = parent_name
        self.workspace = workspace
        self.cmd_parts = cmd_parts
        self.run_script = run_script
        self.repeat = repeat
        self.is_parent = bool(repeat)
        self.repeats_remaining = repeat if context.repeats_remaining is None else context.repeats_remaining 
        self.context = context
        self.status = status    
        self.is_wrapped_up = False     # this is set to True by controller when all wrapup processing has completed
        self.show_output = show_output
        self.callbacks = [] 
        self.acallbacks = []
        self.exit_code = None
        self.killing = False
        self.lock = Lock()   
        self.process = None     
        self.store = store    
        self.recent_output = []
        self.max_recent = 10
        self.rundir = None      # assigned at start time
        self.started = time.time()      # time started with current status
        self.elapsed = None
        self.parent_prep_needed = parent_prep_needed      # when True, must run parent prep script before running first child
        self.console_fn = None
        self.process_was_created = False
        self.mirror_worker = None
        self.mirror_close_func = mirror_close_func
        self.job_id = context.job_id
        self.exper_name = context.exper_name
        self.username = context.username
        self.run_as_parent = context.search_style != "single"
        self.node_id = node_id
        self.run_index = run_index
        self.killed_for_restart = False
        self.max_delay = max_delay
        self.delay_started = None

    def set_console_fn(self, console_fn):
        console_fn = os.path.expanduser(console_fn)
        self.console_fn = console_fn

        if os.path.exists(console_fn):
            os.remove(console_fn)

        file_utils.ensure_dir_exists(file=console_fn)


    def process_run_output(self, text_msg, run_info_is_locked=False):

        # if self.context.scrape:
        #     # scrape output line for XT log records
        #     if self.scrape_output(text_msg):
        #         # don't show scraped output to user
        #         return

        # let run_info keep recent output for new clients
        self.update_recent_output(text_msg)

        # append to output file
        if self.console_fn:
            with open(self.console_fn, "a") as tfile:
                tfile.write(text_msg)

        # send output to attached clients
        if self.show_output:
            # console.print output on CONTROLLER console
            sys.stdout.write(self.run_name + ": ")
            sys.stdout.write(text_msg)

            # if run_info_is_locked:
            #     list2 = list(self.acallbacks)
            # else:
            #     with self.lock:
            #         # make a copy of list for safe enumeration 
            #         list2 = list(self.acallbacks)

            # semi-thread safe (avoid creating list under a lock for every output line)
            for i in range(len(self.acallbacks)):
                try:
                    callback = self.acallbacks[i]
                    #print("sending text to callback: " + text_msg)
                    callback(self.run_name, text_msg)
                except BaseException as ex:
                    # if stream is closed unexpectly, treat as non-fatal error and just log
                    logger.exception("Error in process_run_output, ex={}".format(ex))

    def get_core_properties(self):
        dd = {"run_name": self.run_name, "workspace": self.workspace, "cmd_parts": self.cmd_parts, "run_script": self.run_script, 
            "repeat": self.repeat,  "status": self.status, "show_output": self.show_output, "repeats_remaining": self.repeats_remaining,
            "parent_name": self.parent_name}

        context = copy.copy(self.context.__dict__)
        dd["context"] = context

        return dd

    def attach(self, callback):
        self.callbacks.append(callback)       
        acallback = rpyc.async_(callback)

        # first, send recent output 
        #lines = "\n<recent output>\n\n" + "".join(self.recent_output)
        lines = "".join(self.recent_output)
        acallback(self.run_name, lines)

        # now, hook to our list of callback for next output line
        with self.lock: 
            self.acallbacks.append(acallback)           

    def detach(self, callback):
        index = None

        with self.lock:
            index = self.callbacks.index(callback)
            if index > -1:
                del self.callbacks[index]
                del self.acallbacks[index]
            else:
                errors.internal_error("could not find callback to detach")

        return index

    def update_recent_output(self, msg):
        self.recent_output.append(msg)
        if len(self.recent_output) > self.max_recent:
            self.recent_output = self.recent_output[10:]        # drop oldest 10 lines

    def wrapup_parent_prep_run(self):
        context = self.context
        self.store.log_run_event(context.ws, self.run_name, "parent_script_completed", {"status": self.status})  

    def run_wrapup(self):
        '''wrap-up the run (logging, capture)'''
        #console.print("run_wrapup, self.killed_for_restart={}".format(self.killed_for_restart))
        
        log_info("run_wrapup: mirror_worker", mirror_worker)
        if self.mirror_worker:
            self.mirror_close_func(self)

        context = self.context
        run_name = self.run_name
        store = self.store

        self.check_for_completed(True)

        if not self.killed_for_restart:
            # for "xt restart controller", do not log "cancelled" for runs
            # in spirit of a box that has just been preempted

            log_info("runinfo.is_parent", self.is_parent)
            log_info("runinfo.repeat", self.repeat)

            log_info("calling WRAPUP_RUN", self.run_name)

            store.wrapup_run(context.ws, run_name, context.aggregate_dest, context.dest_name, 
                status=self.status, exit_code=self.exit_code, primary_metric=context.primary_metric, 
                maximize_metric=context.maximize_metric, report_rollup=context.report_rollup, rundir=self.rundir, 
                after_files_list=context.after_files_list, after_omit_list=context.after_omit_list, log_events=context.log, 
                capture_files=context.after_upload, job_id=self.job_id, is_parent=self.is_parent, node_id=self.node_id, 
                run_index=self.run_index, start_time=self.started)

    def set_status(self):
        if self.killing:
            self.status = "cancelled"
        elif self.exit_code:
            self.status = "error"
        else:
            self.status = "completed"

        self.elapsed = time.time() - self.started
        #console.print("set_status: status=", self.status)

    def check_for_completed(self, wait_if_needed=False):
        result = None
        if self.process:

            # wait for process to completely terminate
            if wait_if_needed:
                log_info("waiting for terminate", self.run_name)
                self.process.wait()

            presult = self.process.poll()
            exited = (presult != None)
            log_info("process.poll() result=", presult)
            log_info("exited=", exited)

            if exited or wait_if_needed:
                self.exit_code = self.process.wait()
                self.process = None
                self.set_status()
                result = True  # "completed (exit code=" + str(self.exit_code) + ")"
            elif self.killing:
                self.status = "dying"   
        return result

    def get_summary_stats(self):
        self.check_for_completed()
        elapsed = self.elapsed if self.elapsed else time.time() - self.started
        return "{}^{}^{}^{}".format(self.workspace, self.run_name, self.status, elapsed)

    def set_process(self, process):
        self.process = process
        self.status = "running"
        self.process_was_created = True
        #console.print("set_process for {}, pid={}".format(self.run_name, self.process.pid))

    def kill(self):

        result = False
        console.print("run_info: self.status=", self.status, ", self.process=", self.process)
        before_status = self.status

        if self.process:
            result = self.check_for_completed(False)

            if not result and self.process and self.process.pid:
                self.killing = True

                # convert popen object to real Process object
                console.print("processing kill request for {}, pid={}".format(self.run_name, self.process.pid))

                # allow for "no such process" exception due to timing errors
                try:
                    p = psutil.Process(self.process.pid)

                    # since we run job in a batch file, we need to enumerate all kill
                    # all child processes
                    kids = p.children(recursive=True)
                    for kid in kids:
                        console.print("  killing CHILD process, pid=", kid.pid)
                        kid.kill()
                except psutil.NoSuchProcess as ex:
                    console.print("run={}, exception while killing process: {}".format(self.run_name, ex))

                # it may have changed async since above check 
                if self.process and self.process.pid:
                    console.print("  killing MAIN process, pid=", self.process.pid)
                    self.process.kill()
                
                result = self.check_for_completed(False)
                #self.exit_code = self.process.returncode
                #self.process = None
                #self.status = "cancelled"
                result = True   # "cancelled! (exit code=" + str(self.exit_code) + ")"
        elif self.status in ["queued", "spawning", "running"]:
            self.status = "cancelled"
            self.elapsed = time.time() - self.started

            # must call manually since it doesn't have a process exit_handler()
            self.run_wrapup()          
            result = True   
        elif self.status in ["completed", "error", "cancelled", "aborted"]:
            result = False   # nothing for us to do
        else:
            console.print("run_info.kill: unexpected status=", self.status)

        return result, self.status, before_status

