"""
Job
"""
import os
import time
import inspect
import dill
import subprocess

from . import core
from .scheduler import Scheduler
from .helpers import mkdir, serialize_kwargs, rmd

class JobFailure(Exception):
    """Exception raised when a failed job is found in the session"""
    pass

def _load_yaml(path):
    with open(path) as fh:
        db = {}
        for line in fh:
            key, value = line.split(':')
            db[key] = value
    return db


_stack_index = -1

class Job:

    """
    Batch execution on distibuted memory environment (ex. scheduler)
    """

    mode = core.mode

    def __init__(self, task, scheduler=None, cores=1, wall_time=None,
                 output=None, error=None,
                 **kwargs):
        """
        :param task: the `Task` to execute
        :param scheduler: the scheduler to usel if `None` it will be auto-detected
        :param cores: number of cores to use
        :param wall_time: the maximum wall time
        :param output: output file
        :param error: error file
        """
        self.task = task
        self.scheduler = scheduler
        self.cores = cores
        self.wall_time = wall_time
        self.output = output
        self.error = error
        # Arguments for task (as keyword arguments only)
        self.kwargs = kwargs

        # This will autodetect a scheduler on localhost
        if self.scheduler is None:
            self.scheduler = Scheduler(backend='auto')

    @property
    def path(self):
        """Path to the cache of the job"""
        return os.path.join(self.task.cache.path, self.qualified_name())

    def name(self):
        """General name of job (independent of arguments)"""
        return self.task.name(**self.kwargs)

    def qualified_name(self):
        """Unique name of job"""
        return self.task.qualified_name(**self.kwargs)

    def pretty_name(self):
        """Pretty printed name of job (with args)"""
        args = []
        for key in self.kwargs:
            if self.task.ignore is not None:
                if key in self.task.ignore:
                    continue
            if isinstance(self.kwargs[key], str):
                args.append(f'{key}="{self.kwargs[key]}"')
            else:
                args.append(f'{key}={self.kwargs[key]}')
        kwargs = ','.join(args)
        return f'{self.task.name()}({kwargs})'

    def clear(self, **kwargs):
        """Clear undelrying task and its artifacts"""
        self.task.clear(self.task, **kwargs)
        # Clear job artifacts too
        rmd(self.path)

    # TODO: drop?
    def done(self, **kwargs):
        """Return True is Job is done"""
        # TODO: add condition on job execution
        if hasattr(self.task, 'done'):
            return self.task.done(**kwargs)
        return False

    @property
    def duration(self):
        """Current duration of job. If job is ended, return the elapsed duration"""
        import datetime
        job_file = os.path.join(self.path, 'job.yaml')
        if not os.path.exists(job_file):
            return datetime.timedelta(seconds=0)
        else:
            db = _load_yaml(job_file)
            if 'job_start' in db and 'job_end' in db:
                delta = float(db['job_end']) - float(db['job_start'])
            elif 'job_start' in db:
                delta = time.time() - float(db['job_start'])
            else:
                delta = 0
            return datetime.timedelta(seconds=int(delta))

    @property
    def state(self):
        """
        State of job. 

        Possible values:
        - failed
        - ended
        - running
        - queued
        - unknown       
        """
        job_file = os.path.join(self.path, 'job.yaml')
        if not os.path.exists(job_file):
            return ''
        else:
            db = _load_yaml(job_file)
            if 'job_fail' in db:
                return 'failed'
            elif 'job_end' in db:
                return 'ended'
            elif 'job_start' in db:
                # return 'started'
                return 'running'
            elif 'job_queue' in db:
                return 'queued'
            elif len(db) == 0:
                return 'unknown'
            else:
                raise ValueError(f'wrong state {list(db.keys())} in {self.path}')

    def __call__(self, **kwargs):
        # If no arguments are passed than we assume they were passed
        # to the constructor
        if len(kwargs) > 0:
            self.kwargs = kwargs

        # Add this to global list, make a snapshot for transient jobs
        from copy import deepcopy
        core._jobs.append(deepcopy(self))
        core.jobs.append(self.qualified_name())

        # TODO: if job is running (ex. on a remote cluster) we should skip right away
        assert self.mode in ['safe', 'brave', 'timid', 'clean', 'veryclean', 'scratch']
        if self.mode == 'timid':
            pass
        elif self.mode == 'scratch':
            self.clear(**self.kwargs)
            self._submit()
        elif self.state == 'failed' and self.mode == 'safe':
            raise JobFailure(f'stop because of previous job failure {self.qualified_name()}')
        elif self.state == 'failed' and self.mode == 'brave':
            # TODO: this should only be done if there are no deps
            print(f'WARNING: skipping failed job {self.qualified_name()}')
        elif self.state == 'failed' and self.mode == 'clean':
            self.clear(**self.kwargs)
        elif self.mode == 'veryclean':
            self.clear(**self.kwargs)
        elif self.task.done(**self.kwargs):
            return self.task(**self.kwargs)
        else:
            self._submit()

        # We reset the arguments if they were passed to __call__
        if len(kwargs) > 0:
            self.kwargs = {}

    def _submit(self):
        # Do nothing if job is in queue
        if self.scheduler is not None:
            if self.scheduler.queued(self.qualified_name()):
                # print('JOB IS IN QUEUE', self.qualified_name())
                return

        # dirname = os.path.join(self.task.cache.path, self.qualified_name())
        mkdir(self.path)
        session_pkl = os.path.join(self.path, '.session.pkl')
        job_state = os.path.join(self.path, 'job.yaml')
        job_output = os.path.join(self.path, 'job.out')

        n = _stack_index
        # n=-2
        stacks = inspect.stack()
        s = stacks[n]
        kwargs = serialize_kwargs(self.kwargs)
        code = f"dill.load_module('{session_pkl}')"

        # print(s.frame.f_locals)
        # print('******')
        # for i, ss in enumerate(stacks):
        #     print(i)
        #     print(ss.frame.f_locals)
        #     dill.dump_module(f'/tmp/s{i}.pkl', module=inspect.getmodule(ss.frame.f_code), refimported=True)
        #     print()

        # print('---------------')
        # print(s.frame)
        # print('---------------')

        # fix issues with emacs python shell execution
        var = s.frame.f_locals
        if '__pyfile' in var:
            # s = stacks[-2]
            s = stacks[n - 1]
            garbage = var.pop('__pyfile')
            del garbage

        # TODO: at the line of the called there should be a Job instance
        # else it
        # for s in stacks[:0:-1]:
        #     print(s.frame, s.frame.f_lineno, s.frame.f_code.co_name)
        #     #print(s.frame.f_code)
        #     j = s.frame.f_lineno - s.frame.f_code.co_firstlineno
        #     src, _ = inspect.getsourcelines(s.frame.f_code)
        #     found = False
        #     for i, l in enumerate(src):
        #         if i == j:
        #             if l.strip().startswith('job('):
        #                 print(i, l.strip())
        #                 found = True
        #             break
        #     if found:
        #         break
        # print(s.frame.f_locals)
        # https://stackoverflow.com/questions/1253528/is-there-an-easy-way-to-pickle-a-python-function-or-otherwise-serialize-its-cod
        # print(s.frame.f_globals)
        dill.dump_module(session_pkl, module=inspect.getmodule(s.frame.f_code), refimported=True)
        # TODO: process the calling line to get the job object instead of hardcoding job
        script = f"""\
#!/usr/bin/env -S python -u
__name__ = '__main__'
import sys
import os
import time
import dill
sys.path.append('.')

while True:
    if os.path.exists('{job_state}'): break

hostname = 'unknown'
if 'HOSTNAME' in os.environ:
    hostname = os.environ['HOSTNAME']

fh = open('{job_state}', 'a')
print('job_node:', hostname, file=fh, flush=True)
print('job_start:', time.time(), file=fh, flush=True)
try:
    {code}
    # print(globals().keys())
    job.task({kwargs})
except:
    print('job_fail:', time.time(), file=fh, flush=True)
    raise
else:
    print('job_end:', time.time(), file=fh, flush=True)
finally:
    fh.close()
    # Remove session data when job is over
    os.remove('{session_pkl}')
"""
        # TODO: drop this, it was only for synchronous, scheduler should always be not None
        if self.scheduler is None:
            import time
            with open(f'{job_state}', 'w') as fh:
                print('job_queue:', time.time(), file=fh, flush=True)

        self.scheduler.submit(script, self.qualified_name(), job_output=job_output,
                              job_error=None)

        # Job is queued
        if self.scheduler is not None:
            import time
            with open(f'{job_state}', 'w') as fh:
                print('job_queue:', time.time(), file=fh, flush=True)
