#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Common interface to manage jobs created with launchgen.

Each job is described with an `~sciexp2.common.instance.Instance`, as generated
by `~sciexp2.launchgen.Launchgen.launcher`.

You can integrate job management into your application by instantiating the
`Launcher` class.

"""

from __future__ import print_function
from __future__ import absolute_import

__author__ = "Lluís Vilanova"
__copyright__ = "Copyright 2009-2015, Lluís Vilanova"
__license__ = "GPL version 3 or later"

__maintainer__ = "Lluís Vilanova"
__email__ = "vilanova@ac.upc.edu"


import hashlib
import os
import stat
import struct
import time
from six.moves import cPickle as pickle
import subprocess

from .common import instance
from .common import progress
from .common import utils
from .common import varref
import sciexp2.system


_DICTXLATOR = varref.DictXlator("@")


class Launcher:
    """Common interface to manage jobs.

    The `types` argument to all methods identifies a set of job states to
    narrow the selection of jobs.

    The `filters` argument to all methods narrows the selection of jobs to
    those matching the given filters.

    """

    def __init__(self, base_path, system, group, depends, submit_args):
        """
        Parameters
        ----------
        base_path : str
            Path to the base directory created by
            `~sciexp2.launchgen.Launchgen.launcher`.
        system : str or System subclass
            Execution system.
        group : `~sciexp2.common.instance.InstanceGroup`
            Job descriptors.
        depends : sequence of str
            Variable names to which jobs depend.
        submit_args : sequence of str
            Extra arguments to the job-submitting program.

        See also
        --------
        sciexp2.system.get

        """
        if not os.path.isdir(base_path):
            raise IOError("Not a directory: " + base_path)

        if isinstance(system, str):
            name = system
            system = sciexp2.system.get(system)
            if system is None:
                raise ValueError("Not a launcher system name: " + name)
        if not issubclass(system, sciexp2.system.System):
            raise TypeError("Not a System subclass: " +
                            system.__class__.__name__)

        if not isinstance(group, instance.InstanceGroup):
            raise TypeError("Not an InstanceGroup")

        self._system = system(base_path, group, depends, submit_args)

    def parse_filters(self, *filters):
        value_filters = []
        filters_ = []
        for f in filters:
            try:
                f = sciexp2.common.filter.Filter(f)
                filters_.append(f)
            except sciexp2.common.filter.ParseError:
                path_abs = os.path.abspath(f)
                path = self._system.get_relative_path(
                    path_abs, self._system._base_abs)
                found = None
                for var in self.variables():
                    values = self._system._launchers[var]
                    if f in values:
                        found = "%s == %r" % (var, f)
                        break
                    if path in values:
                        found = "%s == %r" % (var, path)
                        break
                    if path_abs in values:
                        found = "%s == %r" % (var, path_abs)
                        break
                if found is None:
                    raise
                value_filters.append(found)
        res = sciexp2.common.filter.Filter(*filters_)
        if len(value_filters) > 0:
            res &= sciexp2.common.filter.Filter(" || ".join(value_filters))
        res.validate(self.variables())
        return res

    def variables(self):
        """Return a list with the available variables."""
        return list(self._system._launchers.variables())

    def values(self, *filters):
        """Return a list with the available variables."""
        filter_ = sciexp2.common.filter.Filter(*filters)
        return list(self._system._launchers.select(filter_))

    def summary(self, types, *filters):
        """Print a summary of the state of selected jobs."""
        if len(types) == 0:
            types = list(sciexp2.system.Job.STATES)

        parts = [(sciexp2.system.Job.STATE_LONG[t],
                  len(list(self._system.build([t], *filters))))
                 for t in types]
        max_name = max(len(t) for t, _ in parts)
        max_count = max(len(str(n)) for _, n in parts)
        fmt = "%%-%ds: %%%dd" % (max_name, max_count)
        for t, n in parts:
            print(fmt % (t, n))

    def state(self, types, *filters, **kwargs):
        """Print the states for the selected jobs."""
        expression = kwargs.pop("expression", None)
        if len(kwargs) > 0:
            raise ValueError("Unknown arguments: " + " ".join(kwargs.keys()))

        for job in self._system.build(types, *filters):
            state, name = job.state()
            if expression is not None:
                name = _DICTXLATOR.xlate_rec(expression, job)
            print("(%s) %s" % (sciexp2.system.Job.STATE_SHORT[state], name))

    def submit(self, types, *filters, **kwargs):
        """Submit selected jobs to execution.

        Calls `kill` before submitting a job if it's already running.

        """
        submit_args = kwargs.pop("submit_args", [])
        keep_going = kwargs.pop("keep_going", False)
        if len(kwargs) > 0:
            raise ValueError("Unknown arguments: " + " ".join(kwargs.keys()))
        jobs = list(self._system.build(types, *filters))
        exit_code = 0
        with progress.progressable_simple(
                jobs, None, msg="Submitting jobs...") as pjobs:
            for job in pjobs:
                if job["_STATE"] == sciexp2.system.Job.RUNNING:
                    job.kill()
                job_path = self._system.get_relative_path(job["LAUNCHER"])
                progress.log(progress.LVL_INFO, "Submitting", job_path, "\n")
                try:
                    job.submit(*submit_args)
                except subprocess.CalledProcessError as e:
                    exit_code = e.returncode
                    msg = "ERROR: Exited with code %d" % exit_code
                    if progress.level() < progress.LVL_INFO:
                        msg += " (%s)" % self._system.get_relative_path(job["LAUNCHER"])
                    progress.log(progress.LVL_NONE, msg, "\n")
                    if not keep_going:
                        break
        return exit_code

    def kill(self, types, *filters, **kwargs):
        """Kill the execution of selected jobs."""
        kill_args = kwargs.pop("kill_args", [])
        if len(kwargs) > 0:
            raise ValueError("Unknown arguments: " + " ".join(kwargs.keys()))
        for job in self._system.build(types, *filters):
            if job["_STATE"] == sciexp2.system.Job.RUNNING:
                job.kill(*kill_args)

    def files(self, types, *filters, **kwargs):
        """List files matching an expression."""
        expression = kwargs.pop("expression")
        not_expanded = kwargs.pop("not_expanded", False)
        if len(kwargs) > 0:
            raise ValueError("Unknown arguments: " + " ".join(kwargs.keys()))

        expr_path = expression
        if not os.path.isabs(expr_path):
            expr_path = os.path.join(self._system._base_abs, expr_path)

        expr_files = list(f["__FILE"]
                          for f in utils.find_files("@", expr_path,
                                                   name="__FILE"))

        job_files = set()
        for job in self._system.build(types, *filters):
            path = _DICTXLATOR.xlate_rec(expression, job)
            abs_path = path
            if not os.path.isabs(abs_path):
                abs_path = os.path.join(self._system._base_abs, abs_path)
            if not os.path.exists(abs_path):
                continue
            job_files.add(abs_path)

        if not_expanded:
            res = set(expr_files) - job_files
        else:
            res = set(expr_files) & job_files
        for i in sorted(res, cmp=lambda i1, i2: cmp(expr_files.index(i1),
                                              expr_files.index(i2))):
            print(i)


def _header():
    return "#!/usr/bin/env launcher\n"


def _magic():
    version = "launcher"
    return hashlib.md5(version).hexdigest()


class LauncherLoadError (Exception):
    """Could not load given file."""

    def __init__(self, path):
        self.path = path

    def __str__(self):
        return "Not a job descriptor file: %s" % self.path


def save(file_name, base_to_file, system, group, export, depends, submit_args):
    """Save an InstanceGroup as a job descriptor into a file.

    Parameters
    ----------
    file_name : str
        Path to destination file.
    base_to_file : str
        Relative path from some base directory to the directory containing
        `file_name`.
    system : `~sciexp2.system.System`
        Class of the execution system to use.
    group : `~sciexp2.common.instance.InstanceGroup`
        Job descriptors.
    export : set
        Variable names to export into the description file.
    depends : sequence of str
        Variable names to which jobs depend.
    submit_args : sequence of str
        Extra arguments to the job-submitting program.

    See also
    --------
    load, sciexp2.system.get

    """
    assert system is not None
    if not issubclass(system, sciexp2.system.System):
        system = sciexp2.system.get(system)

    if len(group) > 0:
        for assumed in system.assumes():
            if assumed not in group:
                raise ValueError("'%s' must be defined" % assumed)

    file_dir = os.path.dirname(file_name)
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    with open(file_name, "wb") as file_obj:
        file_obj.write(_header())

        version = _magic()
        file_obj.write(struct.pack("I", len(version)))
        file_obj.write(version)

        pickle.dump((base_to_file, system.name, depends, submit_args),
                    file_obj, -1)
        group.dump(file_obj, export)
    fmode = stat.S_IRWXU | stat.S_IRWXG | stat.S_IROTH | stat.S_IXOTH
    os.chmod(file_name, fmode)


def load(file_name):
    """Load a job description from a file.

    Parameters
    ----------
    file_name : str
        Path to source file.

    Returns
    -------
    sciexp2.system.System
        Instance of the execution system with the job descriptions.

    Raises
    ------
    LauncherLoadError
        The given file cannot be loaded as a `Launcher`.

    See also
    --------
    save
    """
    with open(file_name, "rb") as file_obj:
        header = file_obj.readline()
        if header != _header():
            raise LauncherLoadError(file_name)
        version_size = struct.unpack("I", file_obj.read(struct.calcsize("I")))
        version_size = version_size[0]
        version = file_obj.read(version_size)
        if version != _magic():
            raise LauncherLoadError(file_name)

        base_to_file, system_name, depends, submit_args = pickle.load(file_obj)
        group = instance.InstanceGroup.load(file_obj)

    jd_base = os.path.dirname(os.path.abspath(file_name))
    if base_to_file != "":
        base = jd_base[:-len(base_to_file)]
    else:
        base = jd_base
    base = os.path.relpath(base)
    return Launcher(base, system_name, group, depends, submit_args)
