from __future__ import unicode_literals

import atexit
import logging
import os
import re
import subprocess
from copy import copy
from os.path import join
from subprocess import list2cmdline

from humanfriendly import format_size, parse_size
from pickle_blosc import pickle, unpickle
from pickle_mixin import SlotPickleMixin
from tqdm import tqdm

from . import config, job, util
from ._elapsed import BeginEnd
from ._path import make_sure_path_exists, touch
from ._string import make_sure_unicode

try:
    from concurrent.futures import ProcessPoolExecutor
except ImportError:
    from futures import ProcessPoolExecutor



_cluster_runs = dict()


def load(runid):
    if runid not in _cluster_runs:
        folder = join(config.stdoe_folder(), runid)
        if os.path.exists(join(folder, '.deleted')):
            return None
        cr = unpickle(join(folder, 'cluster_run.pkl'))
        _cluster_runs[runid] = cr
    return _cluster_runs[runid]


def rm(runid):
    folder = join(config.stdoe_folder(), runid)
    touch(join(folder, '.deleted'))


def exists(runid):
    if runid not in _cluster_runs:
        root = join(config.stdoe_folder(), runid)
        ok = os.path.exists(join(root, 'cluster_run.pkl'))
        ok = ok and not os.path.exists(join(root, '.deleted'))
        return ok
    return True


def _submit_job(job):
    fcmd = job.full_cmd
    p = subprocess.Popen(
        list2cmdline(fcmd),
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True)
    (job.odata, job.edata) = p.communicate()
    job.odata = make_sure_unicode(job.odata)
    job.edata = make_sure_unicode(job.edata)
    p.stdout.close()
    p.stderr.close()
    return job


class ClusterRunBase(SlotPickleMixin):
    __slots__ = [
        'requests', 'megabytes', 'jobs', 'group', 'runid', 'title', 'nprocs',
        'mkl_nthreads', 'queue'
    ]

    def __init__(self, title):
        self.requests = []
        self.megabytes = 4096
        self.nprocs = 1
        self.jobs = []
        self.runid = None
        self.title = title
        self.mkl_nthreads = 1
        self.queue = None


class ClusterRun(ClusterRunBase):
    __slots__ = []

    def __init__(self, title='notitle'):
        super(ClusterRun, self).__init__(title)

    def store(self):
        fp = join(config.stdoe_folder(), self.runid, 'cluster_run.pkl')
        pickle(self, fp)

    def request(self, req):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        if req == 'panfs':
            self.requests.append('panfs_nobackup_production')
        elif req == 'fasttmp':
            self.requests.append('fasttmp')
        elif req == 'gpfs':
            self.requests.append('gpfs')
        else:
            raise Exception("Don't know about %s." % req)

    def kill(self, block=True):
        grp = '/cluster/%s' % self.runid
        util.kill_group(grp, block=block)

    @property
    def memory(self):
        nbytes = int(round(self.megabytes / 1024. / 1024.))
        return format_size(nbytes)

    @memory.setter
    def memory(self, siz):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        nbytes = parse_size(siz)
        self.megabytes = int(round(nbytes / 1024. / 1024.))

    def add(self, cmd):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        cmd = cmd if isinstance(cmd, list) else [cmd]
        cmd = [str(c) for c in cmd]

        j = job.Job(len(self.jobs), cmd)
        self.jobs.append(j)

    def _dryrun_job(self, job):
        print(list2cmdline(job.full_cmd))

    def _sequential_run_job(self, job):
        os.system(list2cmdline(job.cmd))

    def resubmit(self, jobid):
        raise NotImplementedError('me implementa ai porra')

    def run(self, dryrun=False):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        self.runid = _generate_runid()
        if not dryrun:
            make_sure_path_exists(join(config.stdoe_folder(), self.runid))
        bcmd = self._init_parse()
        self._parse_requests(bcmd)
        self._parse_memory(bcmd)
        self._parse_nprocs(bcmd)
        bcmd += ['-g', '/cluster/%s' % self.runid]
        if self.queue:
            bcmd += ['-q', self.queue]
        self._end_parse(bcmd)

        ujobs = []
        max_workers = 500

        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            for job in tqdm(
                    executor.map(_submit_job, self.jobs),
                    desc='Submitting jobs'):
                job.runid = self.runid
                ujobs.append(job)

        self.jobs = ujobs
        print('   %d jobs have been submitted   ' % len(self.jobs))

        print("Run ID: %s" % self.runid)

        self.store()

        return self.runid

    def _init_parse(self):
        return ['bsub']

    def _end_parse(self, bcmd):
        for job in self.jobs:
            ofile, efile = self._get_output_files(job.jobid)
            job.bcmd = copy(bcmd) + ['-o', ofile, '-e', efile]
            job.mkl_nthreads = self.mkl_nthreads

    def _get_output_files(self, jobid):
        _max_nfiles = 1000
        pr = str(int(jobid / _max_nfiles))
        base = join(config.stdoe_folder(), self.runid, pr)
        make_sure_path_exists(base)
        ofile = join(base, 'out_%d.txt' % jobid)
        efile = join(base, 'err_%d.txt' % jobid)
        return (ofile, efile)

    def _parse_requests(self, bcmd):
        bcmd.append("-R")
        bcmd.append("select[%s]" % (','.join(self.requests)))

    def _parse_memory(self, bcmd):
        m = self.megabytes
        bcmd.append("-M")
        bcmd.append("%d" % m)
        bcmd.append("-R")
        bcmd.append("rusage[mem=1]")

    def _parse_nprocs(self, bcmd):
        bcmd.append('-n')
        bcmd.append('%d' % self.nprocs)

    def __str__(self):
        njobs = len(self.jobs)
        if njobs == 0:
            return self.runid + ': no job'
        jsum = list2cmdline(self.jobs[0].cmd)
        return self.runid + ' (%d jobs): %s' % (njobs, jsum)

    @property
    def number_jobs(self):
        return len(self.jobs)

    @property
    def number_jobs_pending(self):
        return sum([1 for job in self.jobs if job.ispending()])

    @property
    def number_jobs_running(self):
        return sum([1 for job in self.jobs if job.isrunning()])

    @property
    def number_jobs_finished(self):
        return sum([1 for job in self.jobs if job.hasfinished()])

    @property
    def number_jobs_succeed(self):
        return sum([1 for job in self.jobs if job.hasfinished() and\
                                              job.exit_status() == 0])

    @property
    def number_jobs_failed(self):
        return sum([1 for job in self.jobs if job.hasfinished() and\
                                              job.exit_status() != 0])


def _generate_runid():
    from time import gmtime, strftime
    return strftime('%Y-%m-%d-%H-%M-%S', gmtime())


_registered_cluster_runs = dict()


def _register_cluster_run(cr):
    if id(cr) not in _registered_cluster_runs:
        _registered_cluster_runs[id(cr)] = cr


def _update_storage():
    for cr in list(_registered_cluster_runs.values()):
        cr.store()


atexit.register(_update_storage)


def get_bjob(runid, jobid):
    cr = load(runid)
    _register_cluster_run(cr)
    return cr.jobs[jobid]


def get_groups_summary(nlast=10):

    awk = "awk -F\" \" '{print $1, $2, $3, $4, $5, $6, $7}'"
    cmd = "bjgroup | grep -E \".*`whoami`$\" | %s" % awk
    msg = subprocess.check_output(
        cmd, shell=True, stderr=subprocess.STDOUT, universal_newlines=True)
    msg = make_sure_unicode(msg)
    msg = msg.strip()
    lines = msg.split('\n')
    table = [line.split(' ') for line in lines]

    table = _clear_groups_summary(table)
    table.sort(key=lambda x: x[0], reverse=True)

    table_out = []
    logger = logging.getLogger(__file__)
    for row in table:
        runid = row[0].strip('/').split('/')[1]
        try:
            cr = load(runid)
            if cr is None:
                continue
        except ImportError as e:
            logger.warn('Could not load cluster run %s. Reason: %s.', runid,
                        str(e))
            row.append('UNK')
            row.append('UNK')
            row.append('UNK')
        except IOError as e:
            pass
        else:
            try:
                row.append(cr.number_jobs_failed)
            except TypeError:
                row.append('UNK')

            try:
                row.append(cr.number_jobs_succeed)
            except TypeError:
                row.append('UNK')

            try:
                row.append(cr.title)
            except TypeError:
                row.append('UNK')

            table_out.insert(0, row)
            if len(table_out) >= nlast:
                break

    header = [
        'group_name', 'njobs', 'pend', 'run', 'ssusp', 'ususp', 'finish',
        'failed', 'succeed', 'title'
    ]
    table_out.insert(0, header)
    return table_out


def _clear_groups_summary(table):
    ntable = []
    for row in table:
        if not _isrun_group(row[0]):
            continue
        row[1:] = [int(c) for c in row[1:]]
        ntable.append(row)
    return ntable


def _isrun_group(name):
    name = name.strip('/')
    names = name.split('/')
    if len(names) != 2:
        return False
    return names[0] == 'cluster' and _isrunid(names[1])


_runid_matcher = re.compile(r'^\d\d\d\d-\d\d-\d\d-\d\d-\d\d-\d\d$')


def _isrunid(runid):
    return _runid_matcher.match(runid) is not None

from __future__ import unicode_literals

import atexit
import logging
import os
import re
import subprocess
from copy import copy
from os.path import join
from subprocess import list2cmdline

from humanfriendly import format_size, parse_size
from pickle_blosc import pickle, unpickle
from pickle_mixin import SlotPickleMixin
from tqdm import tqdm

from . import config, job, util
from ._elapsed import BeginEnd
from ._path import make_sure_path_exists, touch
from ._string import make_sure_unicode

try:
    from concurrent.futures import ProcessPoolExecutor
except ImportError:
    from futures import ProcessPoolExecutor



_cluster_runs = dict()


def load(runid):
    if runid not in _cluster_runs:
        folder = join(config.stdoe_folder(), runid)
        if os.path.exists(join(folder, '.deleted')):
            return None
        cr = unpickle(join(folder, 'cluster_run.pkl'))
        _cluster_runs[runid] = cr
    return _cluster_runs[runid]


def rm(runid):
    folder = join(config.stdoe_folder(), runid)
    touch(join(folder, '.deleted'))


def exists(runid):
    if runid not in _cluster_runs:
        root = join(config.stdoe_folder(), runid)
        ok = os.path.exists(join(root, 'cluster_run.pkl'))
        ok = ok and not os.path.exists(join(root, '.deleted'))
        return ok
    return True


def _submit_job(job):
    fcmd = job.full_cmd
    p = subprocess.Popen(
        list2cmdline(fcmd),
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True)
    (job.odata, job.edata) = p.communicate()
    job.odata = make_sure_unicode(job.odata)
    job.edata = make_sure_unicode(job.edata)
    p.stdout.close()
    p.stderr.close()
    return job


class ClusterRunBase(SlotPickleMixin):
    __slots__ = [
        'requests', 'megabytes', 'jobs', 'group', 'runid', 'title', 'nprocs',
        'mkl_nthreads', 'queue'
    ]

    def __init__(self, title):
        self.requests = []
        self.megabytes = 4096
        self.nprocs = 1
        self.jobs = []
        self.runid = None
        self.title = title
        self.mkl_nthreads = 1
        self.queue = None


class ClusterRun(ClusterRunBase):
    __slots__ = []

    def __init__(self, title='notitle'):
        super(ClusterRun, self).__init__(title)

    def store(self):
        fp = join(config.stdoe_folder(), self.runid, 'cluster_run.pkl')
        pickle(self, fp)

    def request(self, req):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        if req == 'panfs':
            self.requests.append('panfs_nobackup_production')
        elif req == 'fasttmp':
            self.requests.append('fasttmp')
        elif req == 'gpfs':
            self.requests.append('gpfs')
        else:
            raise Exception("Don't know about %s." % req)

    def kill(self, block=True):
        grp = '/cluster/%s' % self.runid
        util.kill_group(grp, block=block)

    @property
    def memory(self):
        nbytes = int(round(self.megabytes / 1024. / 1024.))
        return format_size(nbytes)

    @memory.setter
    def memory(self, siz):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        nbytes = parse_size(siz)
        self.megabytes = int(round(nbytes / 1024. / 1024.))

    def add(self, cmd):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        cmd = cmd if isinstance(cmd, list) else [cmd]
        cmd = [str(c) for c in cmd]

        j = job.Job(len(self.jobs), cmd)
        self.jobs.append(j)

    def _dryrun_job(self, job):
        print(list2cmdline(job.full_cmd))

    def _sequential_run_job(self, job):
        os.system(list2cmdline(job.cmd))

    def resubmit(self, jobid):
        raise NotImplementedError('me implementa ai porra')

    def run(self, dryrun=False):
        if self.runid is not None:
            raise Exception('This command set has already been sent.')

        self.runid = _generate_runid()
        if not dryrun:
            make_sure_path_exists(join(config.stdoe_folder(), self.runid))
        bcmd = self._init_parse()
        self._parse_requests(bcmd)
        self._parse_memory(bcmd)
        self._parse_nprocs(bcmd)
        bcmd += ['-g', '/cluster/%s' % self.runid]
        if self.queue:
            bcmd += ['-q', self.queue]
        self._end_parse(bcmd)

        ujobs = []
        max_workers = 500

        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            for job in tqdm(
                    executor.map(_submit_job, self.jobs),
                    desc='Submitting jobs'):
                job.runid = self.runid
                ujobs.append(job)

        self.jobs = ujobs
        print('   %d jobs have been submitted   ' % len(self.jobs))

        print("Run ID: %s" % self.runid)

        self.store()

        return self.runid

    def _init_parse(self):
        return ['bsub']

    def _end_parse(self, bcmd):
        for job in self.jobs:
            ofile, efile = self._get_output_files(job.jobid)
            job.bcmd = copy(bcmd) + ['-o', ofile, '-e', efile]
            job.mkl_nthreads = self.mkl_nthreads

    def _get_output_files(self, jobid):
        _max_nfiles = 1000
        pr = str(int(jobid / _max_nfiles))
        base = join(config.stdoe_folder(), self.runid, pr)
        make_sure_path_exists(base)
        ofile = join(base, 'out_%d.txt' % jobid)
        efile = join(base, 'err_%d.txt' % jobid)
        return (ofile, efile)

    def _parse_requests(self, bcmd):
        bcmd.append("-R")
        bcmd.append("select[%s]" % (','.join(self.requests)))

    def _parse_memory(self, bcmd):
        m = self.megabytes
        bcmd.append("-M")
        bcmd.append("%d" % m)
        bcmd.append("-R")
        bcmd.append("rusage[mem=1]")

    def _parse_nprocs(self, bcmd):
        bcmd.append('-n')
        bcmd.append('%d' % self.nprocs)

    def __str__(self):
        njobs = len(self.jobs)
        if njobs == 0:
            return self.runid + ': no job'
        jsum = list2cmdline(self.jobs[0].cmd)
        return self.runid + ' (%d jobs): %s' % (njobs, jsum)

    @property
    def number_jobs(self):
        return len(self.jobs)

    @property
    def number_jobs_pending(self):
        return sum([1 for job in self.jobs if job.ispending()])

    @property
    def number_jobs_running(self):
        return sum([1 for job in self.jobs if job.isrunning()])

    @property
    def number_jobs_finished(self):
        return sum([1 for job in self.jobs if job.hasfinished()])

    @property
    def number_jobs_succeed(self):
        return sum([1 for job in self.jobs if job.hasfinished() and\
                                              job.exit_status() == 0])

    @property
    def number_jobs_failed(self):
        return sum([1 for job in self.jobs if job.hasfinished() and\
                                              job.exit_status() != 0])


def _generate_runid():
    from time import gmtime, strftime
    return strftime('%Y-%m-%d-%H-%M-%S', gmtime())


_registered_cluster_runs = dict()


def _register_cluster_run(cr):
    if id(cr) not in _registered_cluster_runs:
        _registered_cluster_runs[id(cr)] = cr


def _update_storage():
    for cr in list(_registered_cluster_runs.values()):
        cr.store()


atexit.register(_update_storage)


def get_bjob(runid, jobid):
    cr = load(runid)
    _register_cluster_run(cr)
    return cr.jobs[jobid]


def get_groups_summary(nlast=10):

    awk = "awk -F\" \" '{print $1, $2, $3, $4, $5, $6, $7}'"
    cmd = "bjgroup | grep -E \".*`whoami`$\" | %s" % awk
    msg = subprocess.check_output(
        cmd, shell=True, stderr=subprocess.STDOUT, universal_newlines=True)
    msg = make_sure_unicode(msg)
    msg = msg.strip()
    lines = msg.split('\n')
    table = [line.split(' ') for line in lines]

    table = _clear_groups_summary(table)
    table.sort(key=lambda x: x[0], reverse=True)

    table_out = []
    logger = logging.getLogger(__file__)
    for row in table:
        runid = row[0].strip('/').split('/')[1]
        try:
            cr = load(runid)
            if cr is None:
                continue
        except ImportError as e:
            logger.warn('Could not load cluster run %s. Reason: %s.', runid,
                        str(e))
            row.append('UNK')
            row.append('UNK')
            row.append('UNK')
        except IOError as e:
            pass
        else:
            try:
                row.append(cr.number_jobs_failed)
            except TypeError:
                row.append('UNK')

            try:
                row.append(cr.number_jobs_succeed)
            except TypeError:
                row.append('UNK')

            try:
                row.append(cr.title)
            except TypeError:
                row.append('UNK')

            table_out.insert(0, row)
            if len(table_out) >= nlast:
                break

    header = [
        'group_name', 'njobs', 'pend', 'run', 'ssusp', 'ususp', 'finish',
        'failed', 'succeed', 'title'
    ]
    table_out.insert(0, header)
    return table_out


def _clear_groups_summary(table):
    ntable = []
    for row in table:
        if not _isrun_group(row[0]):
            continue
        row[1:] = [int(c) for c in row[1:]]
        ntable.append(row)
    return ntable


def _isrun_group(name):
    name = name.strip('/')
    names = name.split('/')
    if len(names) != 2:
        return False
    return names[0] == 'cluster' and _isrunid(names[1])


_runid_matcher = re.compile(r'^\d\d\d\d-\d\d-\d\d-\d\d-\d\d-\d\d$')


def _isrunid(runid):
    return _runid_matcher.match(runid) is not None
