import copy
import inspect
import os
import subprocess

import sys

import time

import psutil
from psutil import STATUS_RUNNING, STATUS_STOPPED

from torii import ToriiException
from .stream import Stream

ALL_PROCESS_FILEDS = ['cmdline', 'connections', 'cpu_affinity', 'cpu_percent', 'cpu_times', 'create_time', 'cwd',
                      'environ', 'exe', 'gids', 'io_counters', 'ionice', 'memory_full_info', 'memory_info',
                      'memory_info_ex', 'memory_maps', 'memory_percent', 'name', 'nice', 'num_ctx_switches', 'num_fds',
                      'num_threads', 'open_files', 'parent', 'pid', 'ppid', 'status', 'terminal', 'threads', 'uids',
                      'username']

STATS_KEYS = ['cpu_percent', 'cpu_times', 'create_time', 'cwd', 'exe', 'gids', 'io_counters', 'memory_info',
                'memory_percent', 'name', 'num_fds', 'num_threads',
                'pid', 'ppid', 'status', 'uids', 'username']

MAX_STATS_KEYS = ['cpu_percent', 'memory_info', 'memory_percent', 'num_threads']

SUM_STATS_KEYS = ['cpu_times', 'cpu_percent', 'memory_info', 'memory_percent', 'num_threads']

FIELD_MAP = {}

CHILDREN_CPU_OK = False

class Process(object):

    def __init__(self, command, name, manager, logdir=None, merge_outputs=True, mandatory=False, parse_streams=True):
        self.command = command
        self.name = name
        self.manager = manager
        self.logger = manager.logger
        self.merge_outputs = merge_outputs
        self.mandatory = mandatory
        self.parse_stream = parse_streams

        self.p = None
        self.t = None
        self.dt = None

        self.logdir = logdir if logdir else self.manager.logdir

        self.before_start_handler = None
        self.on_stop_handler = None
        self.on_start_handler = None
        self.stdout_handlers = []
        self.stderr_handlers = []
        self.stats = {}
        self.sum_stats = {}
        self.children_stats = {}

        self.stdout_stream = Stream(process=self, type='stdout')
        
        if not self.merge_outputs:
            self.stderr_stream = Stream(process=self, type='stderr')

    @property
    def desc(self):
        desc ={
            'pid': self.p.pid if self.p else None,
            'name': self.name,
            'command': self.command,
            'mandatory': self.mandatory,
            'status': self.p.poll() if self.p else None
        }

        self.update_stats()
        if self.stats:
            desc.update(self.stats)

        return desc

    def pid_exists(self):
        exist = False
        if self.pid is not None and psutil.pid_exists(self.pid):
            exist = True
        return exist

    def update_stats(self):

        if self.active:
            # update the process stats
            # try:
            #     process = psutil.Process(self.pid)
            # except:
            #     return
            process = self.p

            self._init_field_map(process)

            stats = process.as_dict(attrs=STATS_KEYS)
            t = time.time()
            self.dt = t - self.t
            self.t = t

            self.stats = self._update_stats_fields(stats, self.stats)

            # update the children stats
            current_child_pids = []
            for child in process.children(recursive=True):
                if child.is_running():
                    child_pid = str(child.pid)
                    current_child_pids.append(child_pid)
                    child_stats = child.as_dict(attrs=STATS_KEYS)

                    previous_stats = self.children_stats[child_pid] if child_pid in self.children_stats else None
                    self.children_stats[child_pid] = self._update_stats_fields(child_stats, previous_stats)

                else:
                    pass

            for child_pid, child in self.children_stats.items():
                if child_pid not in current_child_pids:
                    child['status'] = STATUS_STOPPED

            # clear the children hierarchy
            self.stats['children'] = {}

            for child_pid, child_stats in self.children_stats.items():
                child_stats['children'] = {}

            # ... and rebuild it
            for child_pid, child_stats in self.children_stats.items():
                parent_pid = child_stats['ppid']
                parent_stats = self.stats if parent_pid == self.p.pid else self.children_stats[str(parent_pid)]
                parent_stats['children'][child_pid] = child_stats

            # compute the max stats
            self._max_stats(self.stats)
            for child_pid, child_stats in self.children_stats.items():
                self._max_stats(child_stats)

            self._sum_stats()
            # self._aggregate_cpu(self.stats)



    def _init_field_map(self, process):
        if not FIELD_MAP:
            for key in STATS_KEYS:
                try:
                    func = getattr(process, key)
                    stat = func()
                    if hasattr(stat, '_fields'):
                        FIELD_MAP[key] = stat._fields
                except:
                    pass

    def _update_stats_fields(self, stats, previous_stats = None):
        # transform the raw psutils objects in dictionary
        for key, fields in FIELD_MAP.items():
            if key not in stats:
                continue

            raw_stat = stats[key]
            stat = {field: value for field, value in zip(fields, raw_stat)}
            stats[key] = stat

        # compute the cumulative cpu time
        if 'cpu_times' in stats :
            stats['cpu_total_time'] = sum(stats['cpu_times'].values())

        # re-compute the cpu_percent
        if self.dt > 0 and previous_stats is not None and 'cpu_total_time' in previous_stats \
                and 'cpu_total_time' in stats :
            stats['cpu_percent'] = (stats['cpu_total_time'] - previous_stats['cpu_total_time']) / self.dt * 100.


        if previous_stats is not None:
            previous_stats.update(stats)
            return previous_stats
        else:
            return stats

    def _max_stats(self, stats):

        for key in MAX_STATS_KEYS:
            if key not in stats:
                continue

            max_key = 'max_' + key
            stat = stats[key]

            if max_key in stats:
                max_stat = stats[max_key]
                if isinstance(max_stat, dict):
                    for sub_key in max_stat:
                        max_stat[sub_key] = max(stat[sub_key], max_stat[sub_key])
                else:
                    stats[max_key] = max(stat, max_stat)
            else:
                stats[max_key] =  stats[key]

    def _sum_stats(self):

        for key in SUM_STATS_KEYS:

            if key not in self.stats:
                continue

            sum_stat = copy.deepcopy(self.stats[key])

            for child_pid, child_stats in self.children_stats.items():
                if child_stats['status'] != STATUS_RUNNING:
                    continue

                child_stat = child_stats[key]

                if isinstance(sum_stat, dict):
                    for sub_key in sum_stat:
                        sum_stat[sub_key] = max(child_stat[sub_key], sum_stat[sub_key])
                else:
                    sum_stat += child_stat

            self.sum_stats[key] = sum_stat

        # we also apply the max on the sums
        self._max_stats(self.sum_stats)

    def _compute_cpu_percent(self, stats):
        sum_user   = stats['cpu_times']['user']
        sum_system = stats['cpu_times']['system']
        sum_percent = stats['cpu_percent']

        for child_pid, child_stats in stats['children'].items():
            user, system, percent = self._aggregate_cpu(child_stats)
            sum_user += user
            sum_system += system
            sum_percent += percent

        stats['cpu_times']['sum_user'] = sum_user
        stats['cpu_times']['sum_system'] = sum_system
        stats['sum_cpu_percent'] = sum_percent


    def _aggregate_cpu(self, stats):
        sum_user   = stats['cpu_times']['user']
        sum_system = stats['cpu_times']['system']
        sum_percent = stats['cpu_percent']

        for child_pid, child_stats in stats['children'].items():
            user, system, percent = self._aggregate_cpu(child_stats)
            sum_user += user
            sum_system += system
            sum_percent += percent

        stats['cpu_times']['sum_user'] = sum_user
        stats['cpu_times']['sum_system'] = sum_system
        stats['sum_cpu_percent'] = sum_percent


        return sum_user, sum_system, sum_percent

    def get_abstract(self):
        abstract = '{} ({}) : cpu = {:.1f}%, mem = {:.1f}%, threads = {}'\
            .format(self.name, self.command_abstract,
                    self.sum_stats['cpu_percent'], self.sum_stats['memory_percent'], self.sum_stats['num_threads'])
        return abstract

    @property
    def command_str(self):
        return ' '.join(self.command)

    @property
    def command_abstract(self):
        return self.command_str if len(self.command_str) < 25 else self.command_str[:21] + ' ...'

    @property
    def active(self):
        return self.p.poll() is None if self.p else False

    @property
    def dead(self):
        return self.p.poll() is not None if self.p else False

    @property
    def pid(self):
        return self.p.pid if self.p else None

    def start(self, check=None, check_timeout=None):
        """
        Start the process.
        If a check handler is provided, the function waits for the check to be satisfied.

        :param check:         check handler
        :param check_timeout: timeout for the checking
        :return:
        """

        # execute the treatement before the start
        try:
            if self.before_start_handler:
                before_start, kwargs = self.before_start_handler
                before_start(**kwargs)
        except Exception as e:
            raise ToriiException('before_start event handling failed : {}'.format(e))\
                .with_traceback(sys.exc_info()[2])
        
        # create the subprocess
        try:
            self.p = psutil.Popen(self.command,
                                  stdout=subprocess.PIPE,
                                  stderr=subprocess.STDOUT if self.merge_outputs else subprocess.PIPE)
            self.t = time.time()
            self.dt = 0

        except Exception as e:
            raise ToriiException('Cannot start process \'{}\': {}'.format(self.name, e))\
                .with_traceback(sys.exc_info()[2])

        if self.parse_stream:
            tee_path = os.path.join(self.logdir, '{}_{}'.format(self.name, self.p.pid))

            if not self.merge_outputs:
                # start the streaming of stderr
                self.stderr_stream.start(self.p.stderr, tee=tee_path + '.err')

            # start the streaming of stdout
            try:
                return self.stdout_stream.start(self.p.stdout, tee=tee_path + '.out',
                                                check=check, check_timeout=check_timeout)
            except Exception as e:
                self.p.kill()
                self.logger.error('Process {} failed to start properly : {}'.format(self.name, e))


    def set_before_start(self, handler, **kwparams):
        # add the process as argument if needed
        if 'process' in inspect.getargspec(handler).args:
            kwparams['process'] = self

        self.before_start_handler = (handler, kwparams)

    def set_on_start(self, handler, **kwparams):
        # add the process as argument if needed
        if 'process' in inspect.getargspec(handler).args:
            kwparams['process'] = self

        self.on_start_handler = (handler, kwparams)

    def on_start(self):
        self.logger.info('Process {} ({}) has started'.format(self.name, self.pid))

        # call the on_start method of the manager
        self.manager.on_start(process=self)

        try:
            if self.on_start_handler:
                on_start, kwargs = self.on_start_handler
                on_start(**kwargs)
        except Exception as e:
            self.logger.exception('on_start event handling failed: {}'.format(e))
            self.manager.shutdown()

    def set_on_stop(self, handler, **kwparams):
        # add the process as argument if needed
        if 'process' in inspect.getargspec(handler).args:
            kwparams['process'] = self

        self.on_stop_handler = (handler, kwparams)

    def on_stop(self):
        # make sure the process is stopped
        i = 0
        while self.p.poll() is None:
            time.sleep(0.1)
            i += 1
            if i > 50:
                self.logger.warn('Process {} ({}) has been detected closed but continue running'
                                 .format(self.name, self.pid, self.poll()))
                self.p.kill()
                return

        self.logger.info('Process {} ({}) has finished, status = {}'.format(self.name, self.pid, self.poll()))
        try:
            if self.on_stop_handler:
                on_stop, kwargs = self.on_stop_handler
                on_stop(**kwargs)
        except Exception as e:
            self.logger.exception('on_stop event handling failed : {}'.format(e))
            self.manager.shutdown()

        # call the on_stop method of the manager
        self.manager.on_stop(process=self)

    def add_stream_handler(self, handler, stream='stdout', **kwparams):
        # add the process as argument if needed
        if 'process' in inspect.getargspec(handler).args:
            kwparams['process'] = self

        handlers = self.stderr_handlers if stream == 'stderr' else self.stdout_handlers
        handlers.append((handler, kwparams))

    def on_new_line(self, line, stream='stdout'):
        try:
            handlers = self.stderr_handlers if stream == 'stderr' else self.stdout_handlers

            # call the handlers on the line
            for handler, kwparam in handlers:
                handler(line, **kwparam)
        except Exception as e:
            self.logger.warning('on_new_line event handling failed : {}'.format(e))
            self.kill()

    def poll(self):
        return self.p.poll() if self.p else None

    def send_signal(self, signal):
        self.logger.info('Signal {} order on {} ({})'.format(signal, self.name, self.pid))
        return self.p.send_signal(signal) if self.p else None

    def kill(self):
        self.logger.info('Kill order on {} ({})'.format(self.name, self.pid))
        return self.p.kill() if self.p else None

    def clean_shared_mem(self):
        """
        Clean the shared memory segments created by the process
        """

        if not self.p:
            return

        # list the shared memory segments created by VNC
        p = subprocess.Popen(['ipcs', '-mp'], stdout=subprocess.PIPE)
        # we get something like :

        #   ------ Shared Memory Creator/Last-op PIDs --------
        #   shmid      owner      cpid       lpid
        #   0          postgres   1625       29528
        #   32769      postgres   1625       29528
        #   65538      root       1851       1859
        #   98307      root       1851       1859
        #   ...
        #
        # the first column is the shared memory ID and the third is the ID of the process that created it
        #
        mids = []
        pid_str = str(self.p.pid)
        for line in p.stdout.readlines():
            cols = line.split()
            if len(cols) >= 4 and cols[2] == pid_str:
                mids.append(cols[0])

        if mids:
            # clean all the shared memory segments generated by the process
            command = 'ipcrm shm ' + ' '.join(mids)
            self.logger.info('Remove shared memory segments created by {} ({})'.format(self.name, self.p.pid))
            self.logger.debug('... with command : {}'.format(command))
            p = subprocess.Popen(command.split(), stdout=subprocess.PIPE)
            output = p.stdout.read()
            p.wait()
            status = p.poll()
            if status != 0:
                self.logger.warn('Failed to remove shared memory segments (status={}) : {}'.format(status, output))


