from dataclasses import dataclass
from os import chmod, getcwd, mkdir, path

from .utils import exec_shell_command


@dataclass
class SlurmResourcesConfiguration:
    partition: str = "gpu"
    time: str = "01:00:00"
    memory: str = "100GB"
    cpus_per_task: int = 1
    gres: str = "gpu:1"
    cores_per_socket: int = 1

    def format_cli(self):
        return (
            "--partition %s --time=%s --mem=%s -c %d --cores-per-socket %d --gres=%s"
            % (
                self.partition,
                self.time,
                self.memory,
                self.cpus_per_task,
                self.cores_per_socket,
                self.gres,
            )
        )


def submit_bash_script(
    slurm_resources,
    script_file,
    log_dir=None,
    workspace_dir=None,
    slurm_command_args=None,
    log_basename="slurm",
):
    """
    Submit a SLURM script.

    Parameters
    ----------
    slurm_resources: SlurmResourcesConfiguration
        Object describing the SLURM resources to use.
    script_source: str, optional
        String with the source code of the bash script.
        Mutually exclusive with 'script_file'.
    script_file: str, optional
        Script file name.
        Mutually exclusive with 'script_source'.
    log_dir: str, optional
        Directory where the SLURM ".out" file is written. Default is "logs".
    workspace_dir: str, optional
        directory where auxilliary files are stored. Default is .integrator_workspace.
    slurm_command_args: str or list, optional
        Parameters to be passed the bash script.
        If a list is passed, the same "sbatch" command will be looped with the parameters from the list.

    Returns
    -------
    ret: str or bytes or list
        Result from submission command.
        If 'slurm_command_args' is a list, ret will be the result of each individual sbatch command.
    """
    log_dir = log_dir or path.join(path.relpath(getcwd()), "logs")
    workspace_dir = workspace_dir or path.join(
        path.relpath(getcwd()),
        ".integrator_workspace",
    )

    for dir_ in [log_dir, workspace_dir]:
        if not path.isdir(dir_):
            mkdir(dir_)
            chmod(dir_, 0o755)  # nosec B103

    # don't surround slurm-%j.out with quotes, otherwise calls from within python will fail
    slurm_log_arg = "--error=%s/%s-%%j.out --output=%s/%s-%%j.out" % (
        log_dir,
        log_basename,
        log_dir,
        log_basename,
    )
    slurm_resources_arg = slurm_resources.format_cli()

    shell_cmd = "sbatch --export=NONE %s %s %s" % (
        slurm_log_arg,
        slurm_resources_arg,
        script_file,
    )

    if isinstance(slurm_command_args, list):
        res = []
        for arg in slurm_command_args:
            shell_cmd2 = shell_cmd + " %s" % arg
            r = exec_shell_command(shell_cmd2)
            res.append(r)
    else:
        if isinstance(slurm_command_args, str):
            shell_cmd += " %s" % slurm_command_args
        res = exec_shell_command(shell_cmd)
    return res
