import os
import dill
import lz4.frame
import logging
from tqdm.auto import tqdm
from .area import WorkingArea
from .submitter import SGETaskSubmitter, MPTaskSubmitter
from .monitor import JobMonitor

logger = logging.getLogger(__name__)

def validate_tasks(tasks):
    for task in tasks:
        if any(key not in task for key in ("task", "args", "kwargs")):
            return False
        if not callable(task["task"]):
            return False
    return True

def sge_submit(
    tasks, label, tmpdir, options="-q hep.q", dryrun=False, quiet=False,
    sleep=5, request_resubmission_options=True, return_files=False,
):
    """
    Submit jobs to an SGE batch system. Return a list of the results of each
    job (i.e. the return values of the function calls)

    Parameters
    ----------
    tasks : list
        A list of dictrionaries with the keys: task, args and kwargs. Each
        element is run on a node as task(*args, **kwargs).

    label : str
        Label given to the qsub submission script through -N.

    tmpdir : str
        Path to temporary directory (doesn't have to exist) where pysge stores
        job infomation. Each call will have a unique identifier in the form
        tpd_YYYYMMDD_hhmmss_xxxxxxxx. Within this directory exists all tasks in
        separate directories with a dilled file, stdout and stderr for that
        particular job.

    options : str (default = "-q hep.q")
        Additional options to pass to the qsub command. Take care since the
        following options are already in use: -wd, -V, -e, -o and -t.

    dryrun : bool (default = False)
        Create directories and files but don't submit the jobs.

    quiet : bool (default = False)
        Don't print tqdm progress bars. Other prints are controlled by logging.

    sleep : float (default = 5)
        Minimum time between queries to the batch system.

    request_resubmission_options : bool (default = True)
        When a job fails the master process will expect a stdin from the user
        to alter the submission options (e.g. to increase walltime or memory
        requested). If False it will use the original options.

    return_files : bool (default = False)
        Instead of opening the output files and loading them into python, just
        send the paths to the output files and let the user deal with them.
    """
    if not validate_tasks(tasks):
        logger.error(
            "Invalid tasks. Ensure tasks=[{'task': .., 'args': [..], "
            "'kwargs': {..}}, ...], where 'task' is callable."
        )
        return []
    area = WorkingArea(os.path.abspath(tmpdir))
    submitter = SGETaskSubmitter(" ".join(['-N {}'.format(label), options]))
    monitor = JobMonitor(submitter)

    results = []
    area.create_areas(tasks, quiet=quiet)
    try:
        submitter.submit_tasks(area.task_paths, dryrun=dryrun, quiet=quiet)
        if not dryrun:
            results = monitor.monitor_jobs(
                sleep=sleep, request_user_input=request_resubmission_options,
            )
    except KeyboardInterrupt as e:
        submitter.killall()

    if return_files:
        return results

    results_not_files = []
    for path in results:
        with lz4.frame.open(path, 'rb') as f:
            results_not_files.append(dill.load(f))
    return results_not_files

def sge_submit_yield(
    name, path, tasks=[], options="-q hep.q", quiet=False, sleep=5,
    request_resubmission_options=True,
):
    if not validate_tasks(tasks):
        logger.error(
            "Invalid tasks. Ensure tasks=[{'task': .., 'args': [..], "
            "'kwargs': {..}}, ...], where 'task' is callable."
        )
        return []
    area = WorkingArea(os.path.abspath(path))
    submitter = SGETaskSubmitter(" ".join(['-N {}'.format(name), options]))
    monitor = JobMonitor(submitter)

    area.create_areas(tasks, quiet=quiet)
    submitter.submit_tasks(area.task_paths, quiet=quiet)
    return monitor.request_jobs(
        sleep=sleep, request_user_input=request_resubmission_options,
    )

def sge_resume(
    name, path, options="-q hep.q", quiet=False, sleep=5,
    request_resubmission_options=True,
):
    area = WorkingArea(os.path.abspath(path), resume=True)
    submitter = SGETaskSubmitter(" ".join(['-N {}'.format(name), options]))
    for idx in range(len(area.task_paths)):
        submitter.jobid_tasks['{}'.format(idx)] = area.task_paths[idx]
    monitor = JobMonitor(submitter)

    results = []
    try:
        results = monitor.monitor_jobs(
            sleep=sleep, request_user_input=request_resubmission_options,
        )
    except KeyboardInterrupt as e:
        submitter.killall()
    return results

def mp_submit(tasks, ncores=4, quiet=False):
    if not validate_tasks(tasks):
        logger.error(
            "Invalid tasks. Ensure tasks=[{'task': .., 'args': [..], "
            "'kwargs': {..}}, ...], where 'task' is callable."
        )
        return []
    submitter = MPTaskSubmitter()
    return submitter.submit_tasks(tasks, ncores=ncores, quiet=quiet)

def local_submit(tasks, quiet=False):
    if not validate_tasks(tasks):
        logger.error(
            "Invalid tasks. Ensure tasks=[{'task': .., 'args': [..], "
            "'kwargs': {..}}, ...], where 'task' is callable."
        )
        return []

    results = []
    pbar = tqdm(total=len(tasks), desc="Finished", disable=quiet)

    try:
        for t in tasks:
            results.append(t["task"](*t["args"], **t["kwargs"]))
            pbar.update()
    except KeyboardInterrupt:
        results = []

    pbar.close()
    return results
