import shutil
import fnmatch
import os
from pathlib import Path
from subprocess import run

from ploomber.io._commander import Commander, CommanderStop, CommanderException
from jinja2 import Template, meta
import click

from soopervisor import abc
from soopervisor import commons
from soopervisor.shell.config import SlurmConfig
from soopervisor import validate


def _warn_on_exit_if_param(cmdr, param, name):
    if param:
        cmdr.warn_on_exit(f'{name!r} option has no effect when '
                          'using SLURM, ignoring...')


def _check_template_variables(env, source):
    return meta.find_undeclared_variables(env.parse(source))


def _validate_template(env, source):
    vars = _check_template_variables(env, source)
    validate.keys({'command', 'name'}, vars,
                  'Error validating template.sh, missing placeholders')


def _script_name_for_task_name(name, workspace):
    """
    Returns the path to the script to use based on the task name
    """
    name_sh = f'{name}.sh'
    exact_match = Path(workspace, name_sh)

    if exact_match.is_file():
        return exact_match

    files = sorted(os.listdir(workspace))

    for file in files:
        if fnmatch.fnmatch(name_sh, file.replace('__', '*')):
            return Path(workspace, file)

    return Path(workspace, 'template.sh')


# TODO: check ploomber version
# TODO: submit all tasks to a single node
class SlurmExporter(abc.AbstractExporter):
    CONFIG_CLASS = SlurmConfig

    @staticmethod
    def _add(cfg, env_name):
        """
        Add sample template
        """
        with Commander(workspace=env_name,
                       templates_path=('soopervisor', 'assets'),
                       environment_kwargs=dict(
                           variable_start_string='[[',
                           variable_end_string=']]',
                       )) as e:
            e.copy_template('slurm/template.sh')
            e.success('Done')

    @staticmethod
    def _validate(cfg, dag, env_name):
        pass

    @classmethod
    def validate(cls):
        """
        Override method from the abstract class, since it validates lock files
        exist, but we don't need those for SLURM
        """
        pass

    @staticmethod
    def _export(cfg, env_name, mode, until, skip_tests, skip_docker,
                ignore_git, lazy_import, task_name):
        """
        Export and submit jbs
        """
        with Commander(workspace=env_name,
                       templates_path=('soopervisor', 'assets')) as cmdr:

            # these do not apply when using SLURM, warn the user
            # TODO: test the warning is shown
            _warn_on_exit_if_param(cmdr, ignore_git, 'ignore_git')
            _warn_on_exit_if_param(cmdr, until, 'until')
            _warn_on_exit_if_param(cmdr, skip_tests, 'skip_tests')
            _warn_on_exit_if_param(cmdr, skip_docker, 'skip_docker')

            template = Path(env_name, 'template.sh').read_text()
            _validate_template(cmdr._env, template)

            tasks, args = commons.load_tasks(cmdr=cmdr,
                                             name=env_name,
                                             mode=mode,
                                             lazy_import=lazy_import,
                                             task_name=task_name)

            if not tasks:
                raise CommanderStop(f'Loaded DAG in {mode!r} mode has no '
                                    'tasks to submit. Try "--mode force" to '
                                    'submit all tasks regardless of status')

            # FIXME: add unit test for this
            if not shutil.which('sbatch'):
                raise CommanderException(
                    'sbatch is not installed, but it is '
                    'required to submit the jobs to the cluster, '
                    'please install it and try again.')

            _submit_to_slurm(tasks, args, env_name)


def _submit_to_slurm(tasks, args, workspace):
    """
    Script to submit tasks in a Ploomber pipeline as SLURM jobs.

    Parameters
    ----------
    tasks : dict

    args : list

    workspace : str

    """
    # maps task name to SLURM job id
    name2id = {}

    # iterate over tasks
    for name, upstream in tasks.items():

        # determine which script file to use
        script_sh = _script_name_for_task_name(name, workspace)
        # generate script and save
        job_sh = Template(script_sh.read_text())

        ploomber_command = ' '.join(['ploomber', 'task', name] + args)
        script = job_sh.render(name=name, command=ploomber_command)
        Path('_job.sh').write_text(script)

        # does the task have dependencies?
        if upstream:
            # if yes, then use --dependency=afterok:
            ids = ':'.join([name2id[task_name] for task_name in upstream])
            # docs: https://hpc.nih.gov/docs/job_dependencies.html
            # https://slurm.schedmd.com/sbatch.html
            # sbatch [options] script [args...]
            cmd = [
                'sbatch',
                f'--dependency=afterok:{ids}',
                '--parsable',
                # kill job if upstream dependencies fail
                '--kill-on-invalid-dep=yes',
                '_job.sh',
            ]
        else:
            # if no, just submit
            cmd = ['sbatch', '--parsable', '_job.sh']

        # print the submitted command
        click.echo(f'Running command: {" ".join(cmd)}')
        click.echo(f'_job.sh content:\n{script}')

        # submit job
        res = run(cmd, capture_output=True, check=True)

        # retrieve the job id, we'll use this to register --dependency
        name2id[name] = res.stdout.decode().strip()
