import os
import subprocess
from subprocess import Popen, PIPE
from pathlib import Path
from mstk import logger
from .pbsjob import PbsJob
from .scheduler import Scheduler


class Slurm(Scheduler):
    '''
    Slurm job scheduler with support of GPU allocation and MPI/OpenMP hybrid parallelization.

    Slurm is a rather powerful and *complicated* job scheduler with tons of configurations and options.
    It is not the goal of `mstk` to provided a comprehensive wrapper for Slurm.
    Therefore, it's very likely that the job script generated by this class doesn't fully fit the requirement of a specific computing center.
    In that case, it's viable to do some process on the generated job script before submitting it.

    Parameters
    ----------
    queue : str
        The jobs will be submitted to this partition.
    n_proc : int
        The CPU cores a job can use.
    n_gpu : int
        The GPU card a job can use.
    n_node : int
        The nodes a job can use. If 0, then will be decided by slurm.
    exclude: str, Optional
        The nodes to be excluded, in Slurm format
    env_cmd : str, Optional
        The commands for setting up the environment before running real calculations.
        It will be inserted on the top of job scripts.

    Attributes
    ----------
    queue : str
        The jobs will be submitted on this queue.
    n_proc : int
        The CPU cores a job can use.
    n_gpu : int
        The GPU card a job can use.
    n_node : int
        The nodes a job can use. If 0, then will be decided by slurm.
    exclude: str
        The nodes to be excluded, in Slurm format
    env_cmd : str
        The commands for setting up the environment before running real calculations.
    sh : str
        The default name of the job script
    max_running_hour: int
        The wall time limit for a job in hours.
    username : str
        The current user
    cached_jobs_expire : int
        The lifetime of cached jobs in seconds.
    submit_cmd : str
        The command for submitting the job script.
        If is `sbatch` by default. But extra argument can be provided, e.g. `sbatch --qos=debug`.
    '''

    #: Whether or not this is a remote job scheduler
    is_remote = False

    def __init__(self, queue, n_proc, n_gpu, n_node=0, exclude=None, env_cmd=None):
        super().__init__(queue=queue, n_proc=n_proc, n_gpu=n_gpu, env_cmd=env_cmd)
        self.n_node = n_node
        self.exclude = exclude or ''
        self.sh = '_job_slurm.sh'
        self.submit_cmd = 'sbatch'

    def is_working(self) -> bool:
        '''
        Check whether or not Slurm is working normally on this machine.

        It calls `sinfo --version` and check the output.

        Returns
        -------
        is : bool
        '''
        cmd = 'sinfo --version'
        sp = Popen(cmd.split(), stdout=PIPE, stderr=PIPE)
        stdout, stderr = sp.communicate()
        return stdout.decode().startswith('slurm')

    def generate_sh(self, commands, name, workdir=None, sh=None, id_prior=None, **kwargs):
        '''
        Generate a shell script for commands to be executed by the job scheduler on compute nodes.

        Because of the complexity of Slurm configurations, it's probable that the job script generated here is not fully valid.
        In that case, it's viable to do some process on the generated job script before submitting it.

        Parameters
        ----------
        commands : list of str
            The commands to be executed step by step
        name : str
            The name of the job to be submitted
        workdir : str, Optional
            The working directory of this job
        id_prior: int, Optional
            The id of prior job this one depends on
        sh : str, Optional
            The name (path) of shell script to be generated.
            If not set, will use the default :attr:`sh`.
        '''
        workdir = workdir or (self.remote_dir if self.is_remote else os.getcwd())
        workdir = Path(workdir).absolute().as_posix()
        sh = sh or self.sh
        sh_basename = Path(sh).stem
        out = sh_basename + '.out'
        err = sh_basename + '.err'
        node_cmd = f'#SBATCH --nodes={self.n_node}\n' if self.n_node > 0 else ''
        gpu_cmd = f'#SBATCH --gres=gpu:{self.n_gpu}\n' if self.n_gpu > 0 else ''
        dep_cmd = f'#SBATCH --dependency=afterok:{id_prior}\n' if id_prior is not None else ''  # id_prior can be 0
        exclude_cmd = f'#SBATCH --exclude={self.exclude}\n' if self.exclude else ''

        with open(sh, 'w') as f:
            f.write(f'#!/bin/bash\n'
                    f'#SBATCH -J {name}\n'
                    f'#SBATCH -D {workdir}\n'
                    f'#SBATCH -o {out}\n'
                    f'#SBATCH -e {err}\n'
                    f'{dep_cmd}'
                    f'#SBATCH -p {self.queue}\n'
                    f'#SBATCH --time={self.max_running_hour}:00:00\n'
                    f'{node_cmd}'
                    f'#SBATCH --ntasks={self.n_proc}\n'
                    f'{gpu_cmd}'
                    f'{exclude_cmd}'
                    f'\n'
                    f'{self.env_cmd}\n\n'
                    )
            for cmd in commands:
                f.write(cmd + '\n')

    def submit(self, sh=None):
        sh = sh or self.sh
        cmd = self.submit_cmd + ' ' + sh
        p = Popen(cmd.split(), stdout=PIPE, stderr=PIPE)
        b_out, b_err = p.communicate()
        out, err = b_out.strip().decode(), b_err.strip().decode()
        if p.returncode == 0 and out.startswith('Submitted batch job'):
            logger.info(f'{out}')
            return int(out.split()[-1])
        else:
            if out:
                logger.error(f'{out}')
            if err:
                logger.error(f'{err}')
            if not out and not err:
                logger.error(f'sbatch failed')
            return -1

    def kill_job(self, name) -> bool:
        job = self.get_job_from_name(name)
        if job is None:
            return False

        cmd = f'scancel {job.id}'
        return subprocess.call(cmd.split()) == 0

    def get_all_jobs(self):
        # Show all jobs. Then check the user
        cmd = 'scontrol show job'
        sp = Popen(cmd.split(), stdout=PIPE, stderr=PIPE)
        stdout, stderr = sp.communicate()
        if sp.returncode != 0:
            logger.error(stderr.decode())
            return []

        jobs = []
        for job_str in stdout.decode().split('\n\n'):  # split jobs
            if job_str.startswith('JobId'):
                job = self._get_job_from_str(job_str)
                # Show all jobs. Then check the user
                if job.user == self.username:
                    jobs.append(job)
        return jobs

    def _get_job_from_str(self, job_str) -> PbsJob:
        workdir = None
        for line in job_str.split():  # split properties
            try:
                key, val = line.split('=')[0:2]
            except:
                continue
            if key == 'JobId':
                id = int(val)
            elif key == 'UserId':
                user = val.split('(')[0]  # UserId=username(uid)
            elif key == 'JobName' or key == 'Name':
                name = val
            elif key == 'Partition':
                queue = val
            elif key == 'JobState':
                state_str = val
                if val in ('PENDING', 'RESV_DEL_HOLD'):
                    state = PbsJob.State.PENDING
                elif val in ('CONFIGURING', 'RUNNING', 'COMPLETING', 'STOPPED', 'SUSPENDED'):
                    state = PbsJob.State.RUNNING
                else:
                    state = PbsJob.State.DONE
            elif key == 'WorkDir':
                workdir = val
        job = PbsJob(id=id, name=name, state=state, workdir=workdir, user=user, queue=queue)
        job.state_str = state_str
        return job
