import time
import sys
import uuid
import math
import random
import keyword
import resource
import os
import datetime
import traceback
import warnings
import platform
import subprocess
#import pickle5

import pandas as pd
from pandas_appender import DF_Appender

from . import stats
from . import pslogger
from . import carbon


class MapProgress(object):
    '''Class to track progress of a map()'''
    # would be perfect as a dataclass, once python 3.7 is our minimum
    def __init__(self, name, d={}, verbose=1, progress_dt=None,
                 carbon_server=None, carbon_port=None, carbon_prefix=None):
        self.name = name
        self.print_last = time.time()
        if progress_dt is not None:
            self.print_dt = progress_dt
        else:
            self.print_dt = self.pick_dt(verbose)
        self.print_log_last = time.time()
        self.print_log_dt = 30  # hardwired
        self.carbon_server = carbon_server
        self.carbon_port = carbon_port
        if carbon_prefix:
            self.carbon_prefix = carbon_prefix.rstrip('.')
        else:
            self.carbon_prefix = None
        if carbon_server and carbon_port:
            carbon.init()

        self.total = d.get('total', 0)
        self.active = d.get('active', 0)
        self.finished = d.get('finished', 0)
        self.failures = d.get('failures', 0)
        self.exceptions = d.get('exceptions', 0)

    def __str__(self):
        attrs = ('total', 'active', 'finished', 'failures', 'exceptions')
        return ', '.join([k+': '+str(getattr(self, k)) for k in attrs])

    def pick_dt(self, verbose):
        if verbose >= 3:
            return 0
        elif verbose >= 2:
            return 1
        elif verbose == 1:
            return 30
        else:
            return 1000000

    def finalize(self, verbose, missing):
        self.active = 0
        failures = self.failures
        actual_failures = len(missing)

        # fixup wrapper failures
        if actual_failures > failures:
            # this seems to happen occasionally in the azure pipeline, disappears on a rerun. always on 38verbose
            pslogger.log('correcting failure count from {} to {}'.format(failures, actual_failures))
            self.failures = actual_failures
        elif actual_failures < failures:
            pslogger.log('can\'t happen! missing pset_ids {} less than failures {}'.format(actual_failures, failures))
        else:
            if verbose > 1 and failures > 0:
                print('failures equal to actual failures, hurrah', file=sys.stderr)

    def report(self, final=None):
        t = time.time()
        last = t - self.print_last > self.print_dt
        log_last = t - self.print_log_last > self.print_log_dt

        if last:
            self.print_last = t
        if log_last:
            self.print_log_last = t

        if final or last or log_last:
            pslogger.log(self.name, 'progress:', str(self), stderr=final or last)
            if self.carbon_server and self.carbon_port:
                carbon.carbon_push_pickle(self.construct_tuples(),
                                          carbon_server=self.carbon_server,
                                          carbon_port=self.carbon_port)

    def construct_tuples(self):
        # [(path, (timestamp, value)), ...]
        timestamp = time.time()

        attrs = ('total', 'active', 'finished', 'failures', 'exceptions')
        return [(self.carbon_prefix+'.'+k, (timestamp, getattr(self, k))) for k in attrs]


class MapResults(pd.DataFrame):
    '''
    A container object for the outcome of paramsurvey.map()

    Properties
    ----------

    to_df : DataFrame

    '''
    def __init__(self, results, missing, progress, stats, verbose=0):
        super().__init__(results)
        self._results_as_listdict = None
        with warnings.catch_warnings():
            # suppress "Pandas doesn't allow columns to be created via a new attribute name"
            warnings.simplefilter('ignore', category=UserWarning)
            self._missing = DFIterDictsWrapper(missing)
        self._progress = progress
        self._stats = stats
        self._verbose = verbose

    @property
    def verbose(self):
        return self._verbose

    @property
    def missing(self):
        return self._missing

    @property
    def progress(self):
        return self._progress

    @property
    def stats(self):
        return self._stats

    def to_df(self):
        return pd.DataFrame(self)

    def to_dict(self, orient=None):
        if orient != 'records':
            raise ValueError('only orient=records is supported')

        df = self.to_df()
        if not self._results_as_listdict:
            vmem0 = vmem()
            size = df.memory_usage().sum()
            if self._verbose > 1 or self._verbose > 0 and size > 1000000:
                pslogger.log('converting Pandas DataFrame to listdict, pandas size was {} bytes'.format(size))
            self._results_as_listdict = df.to_dict(orient='records')
            delta = vmem() - vmem0
            if self._verbose > 1 or self._verbose > 0 and delta > 1.:
                pslogger.log(' vmem increased by {} gigabytes'.format(delta))
        return self._results_as_listdict

    def iterdicts(self):
        return DFListDictIter(self.to_df())


class DFIterDictsWrapper(pd.DataFrame):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def iterdicts(self):
        return DFListDictIter(self)

    def to_df(self):
        return pd.DataFrame(self)

    def to_psets(self):
        return pd.DataFrame(self).drop(columns=['_exception', '_traceback'])


class DFListDictIter:
    def __init__(self, df):
        self._it = df.itertuples(index=False)
        self._colmap = [(c, '_{}'.format(i)) for i, c in enumerate(df.columns)]

    def __iter__(self):
        return self

    def __next__(self):
        d = next(self._it)._asdict()
        for k1, k2 in self._colmap:
            if k2 in d:
                d[k1] = d.pop(k2)
        return d


def report_missing(missing, verbose):
    if missing:
        pslogger.log('missing psets:', stderr=verbose > 1)
        for pset in missing:
            pslogger.log(' ', pset, stderr=verbose > 1)


def get_pset_group(psets, pset_index, group_size):
    group = psets.iloc[pset_index:pset_index+group_size]
    pset_index += group_size

    return group, pset_index


def psets_prep(psets):
    if not isinstance(psets, pd.DataFrame):
        psets = pd.DataFrame(psets)

    for bad in ('_pset_id', '_exception', '_traceback'):
        if bad in psets.columns:
            raise ValueError('disallowed column name {} appears in pset'.format(bad))
    for c in psets.columns.values:
        if not c.isidentifier():
            raise ValueError('pset key "{}" is not a valid Python identifier; this cases problems with pandas'.format(c))
        if keyword.iskeyword(c):
            raise ValueError('pset key "{}" is a Python keyword; this cases problems with pandas'.format(c))
        if c.startswith('_'):
            raise ValueError('pset key "{}" starts with an underscore; this cases problems with pandas'.format(c))
    return psets


def map_prep(psets, name, system_kwargs, chdir, out_subdirs, keep_results=True, progress_dt=None, raise_in_wrapper=None):
    verbose = system_kwargs['verbose']
    vstats = system_kwargs['vstats']

    pslogger.log('paramsurvey.map start time', datetime.datetime.utcnow().strftime('%Y%m%d-%H%M%S'))
    pslogger.log('paramsurvey.map starting work on '+name, stderr=verbose)
    pslogger.log('paramsurvey.map system_kwargs '+repr(system_kwargs), stderr=verbose > 1)

    psets = psets_prep(psets)
    if system_kwargs['limit'] is not None and system_kwargs['limit'] >= 0:
        psets = psets.iloc[:system_kwargs['limit']]

    carbon_server = system_kwargs['carbon_server']
    carbon_port = system_kwargs['carbon_port']
    carbon_prefix = system_kwargs['carbon_prefix']
    system_kwargs['progress'] = MapProgress(name, {'total': len(psets)},
                                            verbose=verbose, progress_dt=progress_dt,
                                            carbon_server=carbon_server, carbon_port=carbon_port, carbon_prefix=carbon_prefix)

    if keep_results:
        system_kwargs['results'] = DF_Appender(ignore_index=True)
    if chdir:
        system_kwargs['chdir'] = chdir
    if out_subdirs:
        system_kwargs['out_subdirs'] = out_subdirs
    if name:
        system_kwargs['name'] = name
    if raise_in_wrapper:
        system_kwargs['raise_in_wrapper'] = raise_in_wrapper

    system_kwargs['pset_ids'] = {}

    system_stats = stats.PerfStats(vstats=vstats)

    pslogger.log('paramsurvey.map pset count {}, pset columns {}'.format(len(psets), list(psets.columns.values)))

    return psets, system_stats, system_kwargs


def map_finalize(name, system_kwargs, system_stats):
    verbose = system_kwargs['verbose']
    vstats = system_kwargs['vstats']

    pslogger.log('finished getting results', stderr=verbose)

    progress = system_kwargs['progress']
    progress.finalize(verbose, system_kwargs['pset_ids'])
    progress.report(final=True)
    system_stats.report(final=True)

    missing = list(system_kwargs['pset_ids'].values())
    report_missing(missing, verbose)
    missing = pd.DataFrame(missing)

    if 'results' in system_kwargs:
        results = system_kwargs['results'].finalize()
    else:
        results = None

    pslogger.log('paramsurvey.map end time', datetime.datetime.utcnow().strftime('%Y%m%d-%H%M%S'))
    pslogger.log('paramsurvey.map returning results')

    return MapResults(results, missing, system_kwargs['progress'], system_stats)


def flatten_results(results):
    seen_pset_keys = set()
    seen_result_keys = set()
    ret = []

    for r in results:
        if 'pset' in r:
            [seen_pset_keys.add(k) for k in r['pset'].keys()]
        if 'result' in r:
            [seen_result_keys.add(k) for k in r['result'].keys()]
        rr = r.get('pset', {}).copy()
        rr.update(r.get('result', {}))
        ret.append(rr)

    conflict = seen_pset_keys.intersection(seen_result_keys)
    if conflict:
        raise ValueError('conflicting key(s) seen in both pset and result: '+repr(conflict))

    return ret


def make_pset_ids(psets):
    pset_ids = {}
    ret = []
    for pset in psets.itertuples(index=False):
        pset_sans_id = pset
        pset_id = str(uuid.uuid4())  # flatten object because of serialization problems downstream
        pset_ids[pset_id] = pset_sans_id._asdict().copy()
        new_pset = pset._asdict().copy()
        new_pset.update({'_pset_id': pset_id})
        ret.append(new_pset)

    return ret, pset_ids


def check_and_exit_if_nested_pool(user_ret):
    # File ".../lib/python3.8/multiprocessing/process.py", line 118, in start
    # assert not _current_process._config.get('daemon'), \
    # AssertionError: daemonic processes are not allowed to have children

    if 'exception' in user_ret:
        e = user_ret['exception']
        tb = user_ret.get('traceback')
        if 'AssertionError' in e and 'daemonic processes are not allowed to have children' in e:
            # this is multiprocessing only
            # after this exception, pool.join() will hang.
            # if we exit without calling pool.join(), in _run_finalizers this will happen:
            # AssertionError: Cannot have cache with result_hander not alive

            if tb:
                pslogger.log(tb, stderr=True)
            pslogger.log('Your code created its own multiprocessing pool, nesting them causes an exception', stderr=True)
            sys.stderr.flush()
            os._exit(1)


def handle_return_common(out_func, ret, system_stats, system_kwargs, user_kwargs):
    progress = system_kwargs['progress']
    verbose = system_kwargs['verbose']
    vstats = system_kwargs['vstats']

    for user_ret, system_ret in ret:
        if 'result' in user_ret and not isinstance(user_ret['result'], dict) and user_ret['result'] is not None:
            # fake an exception, make this case look like other failures
            pslogger.log('user function did not return a dict. faking an exception that says that.', stderr=verbose > 1)
            user_ret['exception'] = "ValueError('user function did not return a dict: {}')".format(
                repr(user_ret['result']))
            user_ret['result'] = {}
        if 'raw_stats' in system_ret:
            system_stats.combine_stats(system_ret['raw_stats'])

        pset_id = user_ret['pset']['_pset_id']

        if 'exception' in user_ret:
            check_and_exit_if_nested_pool(user_ret)

            progress.failures += 1
            progress.exceptions += 1
            progress.active -= 1
            pslogger.log('saw exception in worker: ' + user_ret['exception'], stderr=verbose)
            pslogger.log('pset: '+repr(system_kwargs['pset_ids'][pset_id]), stderr=verbose > 1)

            system_kwargs['pset_ids'][pset_id]['_exception'] = user_ret['exception']

            if 'traceback' in user_ret:
                system_kwargs['pset_ids'][pset_id]['_traceback'] = user_ret['traceback']
                pslogger.log('traceback:\n' + user_ret['traceback'], stderr=verbose > 1)
                pslogger.traceback(True)
                if verbose > 0:
                    print('(traceback is in', pslogger.logger_filename+')', file=sys.stderr)
        else:
            del system_kwargs['pset_ids'][pset_id]
            user_ret['pset'].pop('_pset_id', None)

            new_row = user_ret['pset'].copy()
            if user_ret['result']:  # allowed to be None
                new_row.update(user_ret['result'])

            if 'results' in system_kwargs:
                system_kwargs['results'].append(new_row)

            progress.finished += 1
            progress.active -= 1
            if verbose > 2:
                pslogger.log('finished: pset {} result {}'.format(repr(user_ret['pset']), repr(user_ret['result'])))

        if out_func:
            try:
                out_func(user_ret, system_kwargs, user_kwargs)
            except Exception as e:
                progress.exceptions += 1
                progress.failures += 1
                pslogger.log('saw exception in a user-supplied out_func:', repr(e), stderr=verbose)
                pslogger.log('traceback:\n' + traceback.format_exc(), stderr=verbose)
                pslogger.traceback(True)

    progress.report()
    system_stats.report()


def initialize_kwargs(global_kwargs, kwargs):
    for k, v in global_kwargs.items():
        value = None
        try_env = 'env' in v and v['env'] in os.environ
        if try_env and os.environ[v['env']] == '':
            pslogger.log('empty value for env variable {} ignored'.format(v['env']), stderr=True)

        if try_env and os.environ[v['env']] != '':
            value = os.environ[v['env']]
            v['strong'] = True
        elif kwargs.get(k) is not None:
            value = kwargs.get(k)
        else:
            value = v.get('default')
        type = v.get('type', int)
        if value is not None:
            try:
                v['value'] = type(value)
            except (TypeError, ValueError) as e:
                pslogger.log('exception {} converting type of {} {} {}'.format(str(e), k, repr(value), str(type)), stderr=True)
                raise
        else:
            v['value'] = value


def resolve_kwargs(global_kwargs, kwargs, backend, backends):
    # not quite perfect hack, so we can be verbose here
    verbose = kwargs.get('verbose', 0) or global_kwargs['verbose']['value']

    system_kwargs = {}
    other_kwargs = {}

    for k in kwargs:
        if k in global_kwargs:
            gkw = global_kwargs[k]
            if gkw.get('strong'):
                pslogger.log('environment variable overrides passed in value for '+k, verbose > 0)
                system_kwargs[k] = gkw['value']
            else:
                system_kwargs[k] = kwargs[k]
        else:
            other_kwargs[k] = kwargs[k]
    for k in global_kwargs:
        if k not in system_kwargs:
            system_kwargs[k] = global_kwargs[k]['value']

    backend_kwargs = other_kwargs.pop(backend, {})

    for b in backends:
        if b in other_kwargs:
            pslogger.log('discarding kwarg {} due to our backend choice'.format(b), stderr=verbose > 1)
            other_kwargs.pop(b)

    return system_kwargs, backend_kwargs, other_kwargs


def make_subdir_name(count, prefix='ps'):
    try:
        digits = math.ceil(math.log10(count))
    except Exception:
        print('count argument must be a number greater than 0', file=sys.stderr)
        raise

    return prefix + str(random.randint(0, count-1)).zfill(digits)


def psets_empty(psets):
    if isinstance(psets, pd.DataFrame):
        return psets.empty
    if not psets:
        return True


def vmem():
    ru = resource.getrusage(resource.RUSAGE_SELF)
    gigs = ru[2]/(1024*1024)  # gigabytes
    if platform.system() == 'Darwin':
        gigs = gigs / 1024.
    return gigs


def subprocess_run_worker(pset, system_kwargs, user_kwargs):
    if 'run_args' not in pset:
        raise ValueError('run_args key not found in pset')

    run_kwargs = {}
    if 'run_kwargs' in pset:
        run_kwargs = pset['run_kwargs']
    elif user_kwargs:
        run_kwargs = user_kwargs.get('run_kwargs', {})

    try:
        ret = subprocess.run(pset['run_args'], **run_kwargs)
    except FileNotFoundError:
        if not run_kwargs.get('shell', False) and isinstance(pset['run_args'], str):
            if ' ' in pset['run_args']:
                # we are in the worker so we can't use pslogger()
                print('utils.subprocess_run_worker: space in command string, did you need to set shell=True?', file=sys.stderr)
        raise

    return {'cli': ret}


def pick_factor(args):
    try:
        # pickle5 causes ray to fail https://github.com/ray-project/ray/issues/19938
        #size = len(pickle5.dumps(args))
        size = 0
    except Exception as e:
        if str(e) != "cannot pickle '_io.FileIO' object":  # a test causes this one
            pslogger.log('pick_factor: could not estimate arguments size due to pickle5 exception '+str(e), stderr=True)
        size = 0

    million = 1000000
    if size > million:
        mbytes = size / million
        pslogger.log('worker args looks awfully large ({:.1f} megabytes)'.format(mbytes), stderr=True)

    # also needs: memory, group size, ncores
    # TODO if memory usage is too large, and group size was automatic, decrease group size
    # TODO if after that memory usage is still big, reduce factor

    return 2.4
