import pandas as pd

# code to insert into first and last cell of generated jobs
init_cell_code = """
## generated by omegaml-nbtasks
if '{userid}' != '':
    from omegaml.client.auth import AuthenticationEnv
    auth_env = AuthenticationEnv.secure()
    om = auth_env.get_omega_from_apikey(userid='{userid}', apikey='{apikey}')
else:
    import omegaml as om 
job = {job}
def mark(status):
    task_name = job.get('task_name', '__testing__')
    if task_name == '__testing__':
        print("warning: test run, not recording status")
    meta = om.jobs.metadata(task_name)
    job_meta = meta.attributes['job']
    job_meta['status'] = status
    meta.save()
mark('pending')
"""
done_cell_code = """
## generated by omegaml-nbtasks
mark('finished')
"""


class JobTasks:
    """
    example notebook task runner using omegaml runtime

    Parallel execution of notebooks with parameters on the omegaml runtime

    Usage::

      # submit tasks
      # -- creates 10 tasks as copies of the main notebook in om.jobs.list('tasks')
      # -- runs each task using omegaml runtime
      # -- every .map() call generates a new task group, ensuring unique notebook names
      # -- the original notebook is not changed
      job = om.runtime.job('main')
      job.map(range(10))

      # check status
      job.status()

      # restart tasks that did not produce a result yet
      job.restart()

      # get the list of all notebooks created in one .map() call
      job.list()

      # get the list of all notebooks in any .map() call
      job.list(task_group='*')

      # example notebook
      job = globals().get('job', dict(param=my_default_value))

      def calculation(job):
         # job is a dict with job['param'] set to one value taken from .map(<jobs>)
         # job contains other keys to identify the task:
         #   job_id: a sequence number
         #   task_group: the group of tasks submitted in one .map() call
         #   task_name: the om.jobs name of the running notebook
         #   status: the status, pending or finished
         #   task_id: the celery task id, if available
         ... <insert your code here>

      # run the calculation
      # -- in each task, job will be one of the variables given in nbtasks.map(<jobs>)
      #    <jobs> is an iterable, returning one object for each job
      #    note the job must be serializable. if you need something more complex, pass
      #    the name of an omegaml dataset and possibly some query criteria
      calculation(job)
    """

    def _init_mixin(self, *args, **kwargs):
        self.task_group = None

    def map(self, jobs, job_ids=None, require=None, reset=False, task_group=None):
        """
        Generate any number of parallel jobs executed through om.runtime

        Args:
            jobs (iterable): an iterable to yield job parameters, each parameter
              can be any native object or container (e.g. int, float, dict, list)
            job_ids (str): optional, list of job ids. If passed, the
               job id will be used to name the task id, else
               it is the current count
            require (str): optional, passed on to om.runtime.require()
            reset (bool): optional, if True will resubmit finished jobs
            task_group (str): optional, a specified task group, if not given
               will generate a unique id

        Returns:
            list of started tasks
        """
        nbname = self.jobname
        # generate metadata
        self._generate_jobs(nbname, jobs, job_ids=job_ids, task_group=task_group)
        tasks = self.restart(require=require, reset=reset, task_group=task_group)
        return tasks

    def list(self, task_group=None, **kwargs):
        """
        List all generated tasks for this map call

        Args:
            task_group (str): optional, the task group
            **kwargs (kwargs): kwargs to om.jobs.list()

        Returns:
            list of tasks generated for map call
        """
        om = self.runtime.omega
        task_group = task_group or self.task_group
        nbname = self.jobname
        nbname += f'/{task_group}' if task_group else ''
        return om.jobs.list(f'tasks/{nbname}*', **kwargs)

    def restart(self, task_group=None, reset=False, require=None):
        """
        Run notebook for every entry in tasks/ with no result

        For every entry in om.jobs.list('tasks/*') will check if
        there is a result yet. If not, will call om.runtime.job().run()
        for the given notebook.

        Usage::

            generate_jobs(...)
            restart()

            This will start all jobs that do not have a result yet. To
            start the jobs even if there is a result already, set reset=True.

        Notes:

            ``metadata.attributes['job']`` will record the task status::

            * ``task_id``: the celery task id
            * ``status``: task status, initialized to PENDING
        """
        om = self.runtime.omega
        nbname = self.jobname
        nbname += f'/{task_group}' if task_group else ''
        tasks_nb = om.jobs.list(f'tasks/{nbname}*')
        tasks = []
        for nb in tasks_nb:
            results = om.jobs.list(f'results/{nb}*')
            if not results or reset:
                # generate meta data before running
                meta = om.jobs.metadata(nb)
                job = meta.attributes['job']
                job['status'] = 'pending'
                meta.save()
                # run
                task_rt = om.runtime.require(require).job(nb).run()
                print(f'started {nb} => {task_rt}')
                # update metadata to keep track
                meta = om.jobs.metadata(nb)
                job = meta.attributes['job']
                job['task_id'] = task_rt.id
                meta.save()
                tasks.append(meta)
            else:
                print(f"{nb} has already got results")
        return tasks

    def status(self, task_group=None):
        """
        get the status and the celery id of of each task

        Args:
            task_group (str): the task group id, defaults to None

        Returns::

            pd.DataFrame
                name: task name
                task_id: runtime task id
                status: final status, PENDING, RUNNING, SUCCESS, FAILURE
                run_status: current status, PENDING, RUNNING, SUCCESS, FAILURE
        """
        om = self.runtime.omega
        AsyncResult = self.runtime.celeryapp.AsyncResult
        task_group = task_group or self.task_group
        nbname = self.jobname
        nbname += f'/{task_group}' if task_group else ''
        tasks_nb = om.jobs.list(f'tasks/{nbname}*')
        stats = []
        for nb in tasks_nb:
            meta = om.jobs.metadata(nb)
            task_id = meta.attributes['job'].get('task_id')
            status = AsyncResult(task_id).status if task_id else 'invalid'
            job_runs = meta.attributes.get('job_runs')
            if job_runs:
                run_status = job_runs[-1]['status'] if job_runs else '(waiting)'
            else:
                run_status = 'unknown'
            stats.append((nb, task_id, status, run_status))
        if not tasks_nb:
            print("there are no tasks")
        return pd.DataFrame(stats, columns=['name', 'task_id', 'status', 'run_status'])

    def _generate_jobs(self, nb, jobs, job_ids=None, task_group=None):
        """
        From a notebook, generate a task notebook, parametrized to the jobs arg

        For every job spec in jobs, will generate a copy of the given notebook
        as tasks/{nb}/{task_group}-{id}. The task_group can be set arbitrary or
        be generated as a unique id.

        Once a task is executed, the results will be stored in a new notebook,
        in results/tasks/{nb}/{task_group}-{id}. This contains the full execution
        trace, including graphs and logs of the notebook, and can be accessed
        from Python or Jupyter.

        Usage::

            generate_jobs('main', range(10))

            This will create 10 copies of main, storing each
            in the tasks folder.

            Use restart() to start the tasks on the om runtime

        Args:
            nb (str): the name of the notebook
            jobs (iter): an interable of objects to pass as a job
            job_ids (list): optional, list of job ids. If passed, the
               job id will be used to name the task id, else
               it is the current count
        """
        from nbformat import v4 as nbv4

        om = self.runtime.omega
        job_ids = list(job_ids) if job_ids is not None else None
        self.task_group = task_group = task_group or self._make_task_group()
        tasks = []
        om.jobs.create('#do not delete', 'results/.placeholder')
        om.jobs.create('#do not delete', 'tasks/.placeholder')
        om.jobs.create('#do not delete', 'results/tasks/.placeholder')
        for i, job in enumerate(jobs):
            main_nb = om.jobs.get(nb)
            job_id = i if job_ids is None else job_ids[i]
            task_name = f'tasks/{nb}/{task_group}-{job_id}'
            print(f'generating task {task_name}')
            if om.jobs.metadata(task_name):
                om.jobs.drop(task_name, force=True)
            # store setup of omegaml in main.nb
            job = dict(param=job,
                       job_id=job_id,
                       task_group=task_group,
                       task_name=task_name)
            valid_auth = om.runtime.auth and om.runtime.auth.userid
            auth = (om.runtime.auth.__dict__
                    if valid_auth else dict(userid='', apikey=''))
            code = init_cell_code.format(job=job,
                                         **auth)
            init_cell = nbv4.new_code_cell(source=code)
            done_cell = nbv4.new_code_cell(source=done_cell_code)
            main_nb['cells'].insert(0, init_cell)
            main_nb['cells'].append(done_cell)
            task_meta = om.jobs.put(main_nb, task_name, attributes={'job': job})
            tasks.append(task_meta)

    def _make_task_group(self, max_idlen=8):
        from hashlib import md5
        from uuid import uuid4

        om = self.runtime.omega
        nbname = self.jobname
        value = md5(uuid4().bytes).hexdigest()
        existing = ','.join(om.jobs.list(f'tasks/{nbname}*'))
        while True:
            candidate = value[0:max_idlen]
            if candidate not in existing:
                break
        return candidate


JobTasks.map.__doc__ = JobTasks.__doc__
