'''
This module contains functionality for dependency resolution for constructing
the dependency graph of workflows.
'''

import luigi
from luigi.contrib.postgres import PostgresTarget
from luigi.contrib.s3 import S3Target
from luigi.six import iteritems

# ==============================================================================

class TargetInfo(object):
    '''
    Class to be used for sending specification of which target, from which
    task, to use, when stitching workflow tasks' outputs and inputs together.
    '''
    task = None
    path = None
    target = None

    def __init__(self, task, path, format=None, is_tmp=False):
        self.task = task
        self.path = path
        self.target = luigi.LocalTarget(path, format, is_tmp)

    def open(self, *args, **kwargs):
        '''
        Forward open method, from luigi's target class
        '''
        return self.target.open(*args, **kwargs)

# ==============================================================================

class S3TargetInfo(TargetInfo):
    def __init__(self, task, path, format=None, client=None):
        self.task = task
        self.path = path
        self.target = S3Target(path, format=format, client=client)

# ==============================================================================

class PostgresTargetInfo(TargetInfo):
    def __init__(self, task, host, database, user, password, update_id, table=None, port=None):
        self.task = task
        self.host = host
        self.database = database
        self.user = user
        self.password = password
        self.update_id = update_id
        self.table = table
        self.port = port
        self.target = PostgresTarget(host=host, database=database, user=user, password=password, table=table, update_id=update_id, port=port)

# ==============================================================================

class DependencyHelpers(object):
    '''
    Mixin implementing methods for supporting dynamic, and target-based
    workflow definition, as opposed to the task-based one in vanilla luigi.
    '''

    # --------------------------------------------------------
    # Handle inputs
    # --------------------------------------------------------

    def requires(self):
        '''
        Implement luigi API method by returning upstream tasks
        '''
        return self._upstream_tasks()

    def _upstream_tasks(self):
        '''
        Extract upstream tasks from the TargetInfo objects
        or functions returning those (or lists of both the earlier)
        for use in luigi's requires() method.
        '''
        upstream_tasks = set()
        for attrname, attrval in iteritems(self.__dict__):
            if attrname.startswith('in_'):
                upstream_tasks = self._parse_inputitem(attrval, upstream_tasks)

        return list(upstream_tasks)

    def _parse_inputitem(self, val, tasks):
        '''
        Recursively loop through lists of TargetInfos, or
        callables returning TargetInfos, or lists of ...
        (repeat recursively) ... and return all tasks.
        '''
        if callable(val):
            val = val()
        if isinstance(val, TargetInfo):
            tasks.add(val.task)
        elif isinstance(val, list):
            for valitem in val:
                tasks = self._parse_inputitem(valitem, tasks)
        elif isinstance(val, dict):
            for _, valitem in iteritems(val):
                tasks = self._parse_inputitem(valitem, tasks)
        else:
            raise Exception('Input item is neither callable, TargetInfo, nor list: %s' % val)
        return tasks

    # --------------------------------------------------------
    # Handle outputs
    # --------------------------------------------------------

    def output(self):
        '''
        Implement luigi API method
        '''
        return self._output_targets()

    def _output_targets(self):
        '''
        Extract output targets from the TargetInfo objects
        or functions returning those (or lists of both the earlier)
        for use in luigi's output() method.
        '''
        output_targets = []
        for attrname in dir(self):
            attrval = getattr(self, attrname)
            if attrname.startswith('out_'):
                output_targets = self._parse_outputitem(attrval, output_targets)

        return output_targets

    def _parse_outputitem(self, val, targets):
        '''
        Recursively loop through lists of TargetInfos, or
        callables returning TargetInfos, or lists of ...
        (repeat recursively) ... and return all targets.
        '''
        if callable(val):
            val = val()
        if isinstance(val, TargetInfo):
            targets.append(val.target)
        elif isinstance(val, list):
            for valitem in val:
                targets = self._parse_outputitem(valitem, targets)
        elif isinstance(val, dict):
            for _, valitem in iteritems(val):
                targets = self._parse_outputitem(valitem, targets)
        else:
            raise Exception('Input item is neither callable, TargetInfo, nor list: %s' % val)
        return targets

@@ -2781,33 +2781,27 @@
 lf._
-parse_inputitem(attrval, 
+add_upstream_tasks(
 upst
@@ -2806,24 +2806,33 @@
 stream_tasks
+, attrval
 )%0A%0A        r
@@ -2872,35 +2872,44 @@
 ef _
-parse_inputitem(self, val, 
+add_upstream_tasks(self, tasks, new_
 task
@@ -3118,35 +3118,41 @@
     if callable(
-val
+new_tasks
 ):%0A            v
@@ -3146,33 +3146,45 @@
             
-val = val
+new_tasks = new_tasks
 ()%0A        i
@@ -3188,35 +3188,41 @@
   if isinstance(
-val
+new_tasks
 , TargetInfo):%0A 
@@ -3242,19 +3242,25 @@
 sks.add(
-val
+new_tasks
 .task)%0A 
@@ -3274,35 +3274,41 @@
 elif isinstance(
-val
+new_tasks
 , list):%0A       
@@ -3312,38 +3312,45 @@
         for 
-valitem in val
+new_task in new_tasks
 :%0A          
@@ -3373,38 +3373,42 @@
 lf._
-parse_inputitem(valitem, 
+add_upstream_tasks(tasks, new_
 task
-s
 )%0A  
@@ -3421,35 +3421,41 @@
 elif isinstance(
-val
+new_tasks
 , dict):%0A       
@@ -3458,39 +3458,40 @@
          for _, 
-valitem
+new_task
  in iteritems(va
@@ -3480,35 +3480,41 @@
 sk in iteritems(
-val
+new_tasks
 ):%0A             
@@ -3534,38 +3534,42 @@
 lf._
-parse_inputitem(valitem, 
+add_upstream_tasks(tasks, new_
 task
-s
 )%0A  
@@ -3607,36 +3607,37 @@
 xception('Input 
-item
+value
  is neither call

Clarify code for adding upstream tasks