# Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0
# ----------------------------------
# SGE.py
#       _JobData Class
# ----------------------------------
import logging
import pwd
import os
import re
import subprocess
import time

#
from qsubopts import qsubOptions
sgeJob_t = qsubOptions()


class _JobData:
    """
    Internal helper class to manage job data from qstat
    """
    def __init__(self, qstat_job_line):
        # The format of the line goes like
        # job-ID prior name  user  state submit/start at queue slots ja-task-ID

        tokens = qstat_job_line.split()

        assert len(tokens) >= 8, 'Not a valid qstat line: ' + qstat_job_line
        # Job line must have at least 8 tokens
        try:
            self.id = int(tokens[0])
        except ValueError:
            error_msg = "Could not convert data to an integer."
            raise ValueError(error_msg)
        self.priority = float(tokens[1])
        self.name = tokens[2]
        self.user = tokens[3]
        self.state = tokens[4]
        self.time = ' '.join(tokens[5:7])

        if '@' in qstat_job_line:
            # Has queue assigned, e.g. core2.q@q2.caspian.mit.edu
            self.queue = tokens[7]
            self.slots = tokens[8]
            if len(tokens) == 9:
                self.ja_task_id = None
            elif len(tokens) == 10:
                self.ja_task_id = tokens[-1]
            else:
                raise ValueError(f"Could not parse line: {qstat_job_line}")
        else:
            # No queue assigned
            self.slots = tokens[7]
            if len(tokens) == 8:
                self.ja_task_id = None
            elif len(tokens) == 9:
                self.ja_task_id = tokens[-1]
            else:
                raise ValueError(f"Could not parse line: {qstat_job_line}")

        # Convert array indices ja_task_id into python list format
        ja_task_id = []
        if self.ja_task_id is not None:
            for blob in self.ja_task_id.split(','):
                # Parse data of form '193-349:1' or '1934'
                x = blob.split(':')
                if len(x) == 1:
                    ja_task_id += x
                else:
                    subjobs, step = x
                    begin, end = subjobs.split('-')
                    ja_task_id += range(int(begin), int(end) + 1, int(step))

        self._ja_tasklist = ja_task_id

    def __repr__(self):
        repr = ['{']
        for key, value in self.__dict__.items():
            if key[0] != '_':
                repr.append(key + '=' + str(value))
        repr.append('}')
        return '\n'.join(repr)


class JobList:
    """
    Internal helper class to manage job lists
    """

    def __init__(self, qstat_output=None):
        self._joblist = []
        for line in qstat_output.split('\n')[2:-1]:
            self._joblist.append(_JobData(line))

    def __iter__(self):
        for job in self._joblist:
            yield job

    def __repr__(self):
        return '\n'.join([str(job) for job in self._joblist])


class SGE:
    """External system call handler for Sun Grid Engine environment."""

    def __init__(self, q=None, path='', ):

        logger = logging.getLogger('SGE.__init__')

        if q is None:
            # No queue specified. By default, submit to all available queues.
            self.cmd_qconf = os.path.join(path, 'qconf')

            try:
                qliststr = _exec(self.cmd_qconf + ' -sql')
            except IOError:
                error_msg = 'Error querying queue configuration'
                logger.error(error_msg)
                raise IOError(error_msg)

            self.q = qliststr.replace('\n', ',')[:-1]

            logger.info("""Sun Grid Engine handler initialized
Queues detected: %s""", self.q)

        else:
            self.q = q

        self.cmd_qsub = os.path.join(path, 'qsub')
        self.cmd_qstat = os.path.join(path, 'qstat')

    def wait(self, jobid, interval=10, name=None, pbar=None,
             pbar_mode=None):
        """Waits for job running on SGE Grid Engine environment to finish.

        If you are just waiting for one job, this becomes a dumb substitute for
        the -sync y option which can be specified to qsub.

        Inputs:

        jobid
        interval - Polling interval of SGE queue, in seconds. (Default: 10)
        """

        logger = logging.getLogger('SGE.wait')

        dowait = True
        while dowait:
            p = subprocess.Popen(self.cmd_qstat, shell=True,
                                 stdout=subprocess.PIPE)
            pout, _ = p.communicate()

            if pbar is not None:
                logger.error('Progress bar handling not implemented')

            dowait = False
            for line in pout.split('\n'):
                t = line.split()
                if len(t) >= 5 and t[0] == str(jobid):
                    # Find a line with useful info
                    if re.search(t[4], 'qwrt'):
                        # Job must be queued, running or being transferred
                        dowait = True
                        break
                    if re.search(t[4], 'acuE'):
                        # Job or host in error state
                        logger.warning('Job %d in error state', str(jobid))

            if dowait:
                time.sleep(interval)
                if name is None:
                    logger.info("Time %s: waiting for jobid %s to finish",
                                time.ctime(), str(jobid))
                else:
                    logger.info("Time %s: waiting for job '%s' (jobid %s) to \
finish", time.ctime(), name, str(jobid))

    def submit(self, job, array=False, useenvironment=True, usecwd=True,
               name=None, stdin=None, stdout=None, stderr=None,
               joinstdouterr=True, nproc=1, wait=True, lammpi=True):
        """
        Submits a job to SGE
        Returns jobid as a number
        """

        logger = logging.getLogger('SGE.submit')

        logger.info("Submitting job:   " + str(job) + " stdout: %s \
Stderr: %s", stdout, stderr)

        # Parameters to qsub specified as the header of the job specified on
        # STDIN
        lamstring = lammpi and f" -pe lammpi {nproc}" or ""
        qsuboptslist = ['-cwd -V ', lamstring]

        if name is not None:
            qsuboptslist.append('-N ' + name)
        if stdin is not None:
            qsuboptslist.append('-i ' + stdin)
        if stdout is not None:
            qsuboptslist.append('-o ' + stdout)
        if stderr is not None:
            qsuboptslist.append('-e ' + stderr)
        if joinstdouterr:
            qsuboptslist.append('-j')
        if wait:
            qsuboptslist.append('-sync y')
        if usecwd:
            qsuboptslist.append('-cwd')
        if useenvironment:
            qsuboptslist.append('-V')
        if array is not False:
            try:
                n = int(array[0])
            except IndexError:
                n = int(array)
                raise IndexError("List is empty!")
            except ValueError:
                error_msg = "array[0] being an out of bounds access."
                logger.error(error_msg)
                raise ValueError(error_msg)
            try:
                m = int(array[1])
            except ValueError:
                m = None
                raise ValueError("Could not convert data to an integer.")
            except IndexError:
                m = None
                raise IndexError("array[1] being an out of bounds access.")
            try:
                s = int(array[2])
            except IndexError:
                s = None
                raise IndexError("array[2] being an out of bounds access.")
            except ValueError:
                s = None
                raise ValueError("Could not convert data to an integer.")
            if m == s is None:
                qsuboptslist.append('-t %d' % n)
            elif s is None:
                qsuboptslist.append('-t %d-%d' % (n, m))
            else:
                qsuboptslist.append('-t %d-%d:%d' % (n, m, s))

        qsubopts = ' '.join(qsuboptslist)

        pout = _exec(self.cmd_qsub, stdin=qsubopts + '\n' + job,
                     print_command=False)

        try:
            # Next to last line should be
            # "Your job 1389 (name) has been submitted"
            # parse for job id
            jobid = int(pout.split('\n')[-2].split()[2])
            return jobid
        # except (ValueErrorValueError, IndexError, AttributeError) (e):
        except (ValueError, IndexError, AttributeError):
            error_msg = """Error submitting SGE job:
%s
%s

Output was:
%s""" % (qsubopts, job, pout)
            logger.error(error_msg)
            raise IOError(error_msg)

    def getuserjobs(self, user=pwd.getpwuid(os.getuid())[0]):
        """Returns a list of SGE jobids run by a specific user

        Inputs
            user - SGE user to poll (Default = '', i.e. current user)
            qstat - path to qstat binary (Default = 'qstat')
        """

        p = subprocess.Popen(self.cmd_qstat + " -u " + user, shell=True,
                             stdout=subprocess.PIPE)
        qstat_output, _ = p.communicate()
        joblist = JobList(qstat_output)
        return [job for job in joblist if job.user == user]

    def run_job(self, command, name='default', logfnm='default.log',
                wait=True):
        """Run job on SGE with piped logging."""
        jobid = self.submit(command, name=name, stdout=logfnm,
                            stderr=logfnm, wait=wait)
        return jobid

    def get_queue_instance_status(self):
        """
        Get loads for each queue instance
        """
        output = _exec(' '.join([self.cmd_qstat, '-f']))

        data = []
        for line in output.split('\n')[1:]:
            t = line.split()
            if len(t) != 5:
                continue

            nodename = t[0].split('@')[1].split('.')[0]
            maxslots = int(t[2].split('/')[2])
            load = float(t[3])

            data.append({'name': nodename, 'maxslots': maxslots, 'load': load})

        return data


def _exec(command, print_to_screen=False, logfnm=None, stdin='',
          print_command=False):
    """
    Runs command line using subprocess, optionally returning stdout
    """
    def _call_cmd(command, stdin=''):
        p = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT)
        output, _ = p.communicate(stdin)
        return output

    logger = logging.getLogger('_exec')

    if print_command:
        logger.info("Executing process: \x1b[1;92m%-50s\x1b[0m Logfile: %s",
                    command, logfnm)

    output = ""
    if logfnm is not None:
        try:
            with open(logfnm, 'a') as f:
                if print_command:
                    print(f, "Executing process: %s" % command)
                output = _call_cmd(command, stdin)
                f.write(output)
        except IOError:
            error_msg = 'Error: File: ' + str(logfnm) + ' does not appear to exist.'
            logger.error(error_msg)
            raise IOError(error_msg)
    else:
        output = _call_cmd(command, stdin)

    logger.info('Output of command is:\n%s', output)

    if print_to_screen:
        print(output)

    return output