import girder_worker
import traceback as tb
from celery import Celery
from celery.signals import (task_prerun, task_postrun,
                            task_failure, task_success)
from .utils import JobStatus


def deserialize_job_info_spec(**kwargs):
    return girder_worker.utils.JobManager(**kwargs)


@task_prerun.connect
def gw_task_prerun(task=None, sender=None, task_id=None,
                   args=None, kwargs=None, **rest):
    """Deserialize the jobInfoSpec passed in through the headers.

    This provides the a JobManager class as an attribute of the
    task before task execution.  decorated functions may bind to
    their task and have access to the job_manager for logging and
    updating their status in girder.
    """
    try:
        try:
            # Celery 4.x API
            jobSpec = task.request.jobInfoSpec
        except AttributeError:
            # Celery 3.X API
            jobSpec = task.request.headers['jobInfoSpec']

        task.job_manager = deserialize_job_info_spec(**jobSpec)

        # For now,  only automatically update status if this is
        # not a child task. Otherwise child tasks completion will
        # update the parent task's jobModel in girder.
        if task.request.parent_id is None:
            task.job_manager.updateStatus(JobStatus.RUNNING)

    except (KeyError, TypeError):
        task.job_manager = None
        print('Warning: No jobInfoSpec. Setting job_manager to None.')


@task_success.connect
def gw_task_success(sender=None, **rest):
    try:
        if sender.request.parent_id is None:
            sender.job_manager.updateStatus(JobStatus.SUCCESS)
    except AttributeError:
        pass


@task_failure.connect
def gw_task_failure(sender=None, exception=None,
                    traceback=None, **rest):
    try:

        msg = '%s: %s\n%s' % (
            exception.__class__.__name__, exception,
            ''.join(tb.format_tb(traceback)))

        sender.job_manager.write(msg)

        if sender.request.parent_id is None:
            sender.job_manager.updateStatus(JobStatus.ERROR)
    except AttributeError:
        pass


@task_postrun.connect
def gw_task_postrun(task=None, sender=None, task_id=None,
                    args=None, kwargs=None,
                    retval=None, state=None, **rest):
    try:
        task.job_manager._flush()
        task.job_manager._redirectPipes(False)
    except AttributeError:
        pass


class _CeleryConfig:
    CELERY_ACCEPT_CONTENT = ['json', 'pickle', 'yaml']


app = Celery(
    main=girder_worker.config.get('celery', 'app_main'),
    backend=girder_worker.config.get('celery', 'broker'),
    broker=girder_worker.config.get('celery', 'broker'))

app.config_from_object(_CeleryConfig)

@@ -296,24 +296,69 @@
 **kwargs)%0A%0A%0A
+class JobSpecNotFound(Exception):%0A    pass%0A%0A%0A
 @task_prerun
@@ -795,33 +795,138 @@
 ry:%0A        
-try:%0A    
+# Celery 4.x API%0A        if hasattr(task.request, 'jobInfoSpec'):%0A            jobSpec = task.request.jobInfoSpec%0A%0A
         # Ce
@@ -926,32 +926,139 @@
    # Celery 
-4.x
+3.X
  API%0A
+        elif task.request.headers is not None and %5C%0A                'jobInfoSpec' in task.request.headers:%0A
             
@@ -1072,32 +1072,41 @@
  = task.request.
+headers%5B'
 jobInfoSpec%0A    
@@ -1100,16 +1100,19 @@
 InfoSpec
+'%5D%0A
 %0A       
@@ -1116,117 +1116,371 @@
     
-except AttributeError:%0A            # Celery 3.X API%0A            jobSpec = task.request.headers%5B'jobInfoSpec'%5D
+# Deprecated: This method of passing job information%0A        # to girder_worker is depricated. Newer versions of girder%0A        # pass this information automatically as apart of the%0A        # header metadata in the worker scheduler.%0A        elif 'jobInfo' in kwargs:%0A            jobSpec = kwargs.pop('jobInfo', %7B%7D)%0A%0A        else:%0A            raise JobSpecNotFound
 %0A%0A  
@@ -1847,29 +1847,23 @@
 ept 
-(KeyError, TypeError)
+JobSpecNotFound
 :%0A  

Support instantiation of jobManager through jobInfo passed in kwargs