from functools import wraps
from collections import OrderedDict
import petl
import copy
from enum import Enum
import logging
logger = logging.getLogger(__name__)

HOOK_EVENTS = {}


class OAAHookEvent(str, Enum):
    '''
    :code:`PRE_TRANSFORM` runs just before transforms. A particular source can be
    targeted with target=<source-name>. It is useful when you need to queue
    additional processing of the data stream. The initial connection, query,
    api request, etc will have been made at this point but no paginated
    results.

    :code:`POST_TRANSFORM` runs just after transforms have been queued. A particular
    source can be targeted with target=<source-name>. Additional tweaks the the
    stream of data can be queued here, prior to all data being read. The
    initial connection, query, api request, etc will have been made at this
    point but no paginated results.

    Keep in mind that field mapping happens during
    transforms, just after the `PRE_TRANSFORM`. This means that field mappings will
    have already happened during a POST_TRANSFORM hook execution.

    :code:`PRE_CREATE_OAA` runs after transforms. It is assumed that all paginated
    results of requests, queries, etc will be processed by this hook and read
    into memory. If no hooks listen to this event, the the objectst will be
    read into memory to create the OAA objects. A particular source can be
    targeted with target=<source-name>.

    :code:`PRE_PUSH_OAA` runs just before sending the data to Veza via the API. It
    does not have a specific source in mind. All data will have been read by
    this point and is in-memory.

    :code:`POST_PUSH_OAA` runs just after the data has been sent to Veza. Hooks that
    listen to this event have a `response` object received from Veza.

    '''

    # `RECORD_TRANSFORM_STREAM` runs just after transforms have been mapped to the
    # data stream. Functions that listen for this event should accept a parameter
    # `record`.

    # .. code-block:: python

    #     @record_transform(event=OAAHookEvent.RECORD_TRANSFORM_STREAM, add_field=['roles'])
    #     def manipulate_users_2(record, provider=None, sources=None, **kwargs):
    #         # whatever happens here is after whatever happened in first one
    #         # let's add roles to the user

    #         # TODO pickup here
    #         # import bpdb; bpdb.set_trace()  # noqa: E702
    #         record['roles'] = 'some list of roles'

    #         return record


    # RECORD_TRANSFORM, like RECORD_TRANSFORM_STREAM runs just after the
    # data stream has been processed. The difference between the two is that
    # these functions are called after the OAA objects have been created. The
    # stream of data has been read and it is safe to manipulate objects on the
    # `provider` in these function.
    # LCM_GET_METADATA runs outside of the normal OAA data ingestion. It is used
    # to create hooks that inform Veza of Object metadata. There should only be
    # one of these defined. If mutliple hooks listen to this event, the last one
    # defined is used.

    # LCM_CREATE_USER runs outside of the normal OAA data ingestion. It is used
    # to create a user object.

    # LCM_MODIFY_USER runs outside of the normal OAA data ingestion. It is used
    # to modify a user object.

    # LCM_GENERIC_ACTION  runs outside of the normal OAA data ingestion. It is used
    # to run generic actions. It should be passed an object and an action. The
    # hook function does whatever it wants to with the object and returns a
    # response.

    # TODO Define LCM_ response for LCM hooks.

    PRE_TRANSFORM = 'PRE_TRANSFORM'
    POST_TRANSFORM = 'POST_TRANSFORM'
    PRE_CREATE_OAA = 'PRE_CREATE_OAA'
    PRE_PUSH_OAA = 'PRE_PUSH_OAA'
    POST_PUSH_OAA = 'POST_PUSH_OAA'
    POST_DB_CONNECT = 'POST_DB_CONNECT'
    RECORD_TRANSFORM_STREAM = 'RECORD_TRANSFORM_STREAM'
    RECORD_TRANSFORM = 'RECORD_TRANSFORM'

    # define a hook for this event to get LCM metadata. There should only be
    # one of these defines for an application. Multiple results will 
    # [merge, last one wins, first wins, TBD]
    LCM_GET_METADATA = 'LCM_GET_METADATA'
    # Modify a user. There can be multiple hooks into this one
    LCM_CREATE_USER = 'LCM_CREATE_USER'
    LCM_MODIFY_USER = 'LCM_MODIFY_USER'

    # LCM_POST_CREATE_USER = 'LCM_POST_CREATE_USER'

    # send a payload in and pass an action
    LCM_GENERIC_ACTION = 'LCM_GENERIC_ACTION'

# These hooks should only be defined once per integration.  
HOOKS_THAT_REQUIRE_SINGLETON = [
    OAAHookEvent.LCM_GENERIC_ACTION,
    OAAHookEvent.LCM_GET_METADATA,
]

HOOKS_THAT_REQUIRE_RESULTS = [
    OAAHookEvent.PRE_PUSH_OAA,
    OAAHookEvent.PRE_CREATE_OAA,
]

RECORD_TRANSFORMS_STREAM = {}
RECORD_TRANSFORMS = {}
APP_TRANSFORM_CACHE = {}


def clear_hooks(app_name=None):
    '''If you pass in an app_name, this can be re-used for multiple sequential
    uses of the runner. Otherwise, the hook cache cleared out each time before
    running and hooks will not be retained.
    '''
    global HOOK_EVENTS
    global RECORD_TRANSFORM
    global RECORD_TRANSFORMS_STREAM
    if app_name and app_name in APP_TRANSFORM_CACHE:
        HOOK_EVENTS = APP_TRANSFORM_CACHE[app_name]['HOOK_EVENTS']
        RECORD_TRANSFORMS_STREAM = APP_TRANSFORM_CACHE[app_name]['RECORD_TRANSFORMS_STREAM']
        RECORD_TRANSFORMS = APP_TRANSFORM_CACHE[app_name]['RECORD_TRANSFORMS']
    else:

        HOOK_EVENTS = {}

        if not HOOK_EVENTS:
            # HOOK_EVENTS = {}
            for event in OAAHookEvent:
                HOOK_EVENTS[event] = {'hooks': [], 'results': []}

        RECORD_TRANSFORMS_STREAM = {}
        RECORD_TRANSFORMS = {}

        if app_name:
            APP_TRANSFORM_CACHE.setdefault(app_name, {})
            APP_TRANSFORM_CACHE[app_name]['HOOK_EVENTS'] = copy.copy(HOOK_EVENTS)
            APP_TRANSFORM_CACHE[app_name]['RECORD_TRANSFORMS_STREAM'] = copy.copy(RECORD_TRANSFORMS_STREAM)
            APP_TRANSFORM_CACHE[app_name]['RECORD_TRANSFORMS'] = copy.copy(RECORD_TRANSFORMS)


clear_hooks()


def record_transform(_func=None, *args, event=None,
                     keep_results=False, add_field=[], stream=True, **kwargs):

    '''
    Wrap a function with this decorator and it will add that function
    to the list of hooks that can be called on the source.  By default it is a
    streaming function that accepts a kwarg `source` which you can use to
    manipulate the stream in any way you need to.  If you add new fields to the
    records, pass that in as add_field=['list', 'of', 'new', 'fields']

    In this example, we manipulate the stream of data to de-duplicate the
    records on the `user_id` field.

    .. code-block:: yaml

      # (partial)config.yaml

      user-to-role:
        destination: None
        source_type: source
        source: users
        record_transforms:
          - user_to_role

    .. code-block:: python

        # in modules/appname_hooks.py
        @record_transform()
        def manipulate_rows_for_user(source=None, provider=None, **kwargs):
            # This is only used to verify that the role exists.
            source._stream = source.stream.distinct('user_id')
    
    Decorate a function with this (and set stream=False) and it becomes available as a
    :code:`record_transform_non_stream`. You can reference it in your config.yaml on
    any source. 

    .. code-block:: yaml

      # (partial)config.yaml
      user-to-role:
        destination: None
        source_type: source
        source: users
        record_transforms_non_stream:
          - user_to_role
    

    This can be set to be a non-streaming function so the record itself is passed
    into the function.
    
    In this example, we are looking up the user in the :code:`provider` (the
    OAAClient SDK object) for the user to be sent to Veza and adding a role to the
    user.

    .. code-block:: python

        @record_transform(stream=False)
        def user_to_role(record=None, provider=None, **kwargs):
            # This is only used to verify that the role exists.
            if provider.local_roles.get(record.role_id):
                user_ref = provider.local_users[record.user_id]
                resources = [provider.resources[r] for r in record.resources]
                user_ref.add_role(record.role_id,
                                  apply_to_application=False,
                                  resources=resources)
            else:
                logger.error('role %s does not exist for assignment', record.role_id)
    '''
    def inner(func):

        @wraps(func)
        def wrapper(source, *args, **kwargs):
            """A wrapper function"""

            # Extend some capabilities of func
            # source = kwargs.get('source')
            logger.debug('running %s', func.__name__)

            logger.debug('args %s', len(args))
            logger.debug('kwargs %s', kwargs.keys())
            # results = func(*args, event=event, **kwargs)

            # source._stream = source.stream.rowmap(modify_rows(source), headers)

            def modify_rows(row):

                temp_dict = OrderedDict({k:row[k] for k in row.flds})

                temp_dict = func(temp_dict, *args, **kwargs)
                # return row.values
                # import bpdb; bpdb.set_trace()  # noqa: E702

                return list(temp_dict.values())

            fieldnames = list(petl.header(source.stream))
            fieldnames.extend(add_field)
            source._stream = source.stream.rowmap(modify_rows, fieldnames)

            # return results

        @wraps(func)
        def wrapper_nonstream(*args, **kwargs):
            """A wrapper function"""

            # Extend some capabilities of func
            # source = kwargs.get('source')
            logger.debug('running %s', func.__name__)

            logger.debug('args %s', len(args))
            logger.debug('kwargs %s', kwargs.keys())
            # results = func(*args, event=event, **kwargs)

            # source._stream = source.stream.rowmap(modify_rows(source), headers)

            func(*args, **kwargs)

            # return results

        logger.debug('wrapping event %s', event)

        func_list = RECORD_TRANSFORMS
        _wrapper = wrapper_nonstream

        if stream is True:
            func_list = RECORD_TRANSFORMS_STREAM
            _wrapper = wrapper

        if _wrapper not in func_list:
            func_list[_wrapper.__name__] = _wrapper


        return wrapper

    if _func is None:
        return inner

    return inner(_func)


def record_transform_non_stream(*args, **kwargs):

    return record_transform(*args, **kwargs, stream=False)


def hook(_func=None, *args, event=None,
         target=None, keep_results=False, **kwargs):

    '''
    Wrap a function with this decorator and it will add that function
    to the hooks that are triggered at a certain event.

    An example of this is the following function that reads the `db_raw_input`
    source stream to add the users to roles. It assumes that the source has an
    `ROLE_ID` and `USER_ID` attribute.

    .. code-block:: python

       from oaa.hooks.decorators import hook

       @hook(event=OAAHookEvent.PRE_PUSH_OAA)
       def do_roles_mapping(event, sources=None, provider=None, **kwargs):

           for role in sources['db_raw_input'].records:
               user = provider.local_users[role.USER_ID]
               resource_name = provider.resources[role.ROLE_ID]
               user.add_role(role.ROLE,
                             # apply_to_application=True,
                             resources=[resource_name])

           return True


    Here is an example of the :code:`PRE_TRANSFORM` event. This hook transforms the
    stream of data to return one row for each unique value of the field `name`.

    .. code-block:: python

       from oaa.hooks.decorators import hook

       @hook(event=OAAHookEvent.PRE_TRANSFORM, target='databases')
       def gen_databases(source=None, **kwargs):

           source._stream = source.stream.distinct('name')

    
    The following hook transforms the :code:`roles` stream. Essentially we are
    completely reshaping the stream to have a different set of headers and
    values.

    .. code-block:: python

        @hook(event=OAAHookEvent.PRE_TRANSFORM, target='roles')
        def generate_roles_from_users(event, source=None, *args, **kwargs):

            def gen_roles(row):
                yield [
                   f"id_{row.role['id']}", 
                   row.role['name'], 
                   row.role['description'], 
                   row.role['policies'] 
                ]

            source._stream = source.stream.rowmapmany(gen_roles, header=['id', 'name', 'description', 'policies']) # \

    In this one, we are again transforming a source stream, this time called
    :code:`resources`.  Users in the original stream come in with two different
    fields that can list of resources the user has access to.  In this case,
    all we really want is a list of resources so we use the :code:`rowmapmany` function on the
    :code:`petl` stream to yield the values. Finally, we de-dup the resources by calling
    :code:`.distinct('resource_id')` at the end.

    .. code-block:: python
        
        @hook(event=OAAHookEvent.PRE_TRANSFORM, target='resources')
        def generate_resources_from_users(source=None, **kwargs):

            def gen_resources(row):

                for policy in row.role['policies']:
                    yield [
                       policy['resource'] 
                    ]

                for perm in row.resource_permissions:
                    yield [
                       perm['resource'] 
                    ]

            source._stream = source.stream.rowmapmany(gen_resources, header=['resource_id']) \
                             .distinct('resource_id')
    '''
    def inner(func):

        @wraps(func)
        def wrapper(*args, **kwargs):
            """A wrapper function"""

            # Extend some capabilities of func
            source = kwargs.get('source')
            logger.debug('running %s', func.__name__)

            if target and source:
                logger.debug('only run if target %s == %s',
                             target, source.source_name)

                if source and source.source_name != target:
                    return

            logger.debug('args %s', len(args))
            logger.debug('kwargs %s', kwargs.keys())
            results = func(*args, event=event, **kwargs)

            if keep_results or event in HOOKS_THAT_REQUIRE_RESULTS:
                HOOK_EVENTS[event]['results'].append(results)

            return results

        logger.debug('wrapping event %s', event)

        if event in HOOKS_THAT_REQUIRE_SINGLETON:
            HOOK_EVENTS[event]['hooks'] = [wrapper]
        else:
            if wrapper not in HOOK_EVENTS[event]['hooks']:
                HOOK_EVENTS[event]['hooks'].append(wrapper)
            # print('hello', HOOK_EVENTS[event]['hooks'])
            # import bpdb; bpdb.set_trace()  # noqa: E702

        return wrapper

    if _func is None:
        return inner

    return inner(_func)


def run_hooks(event, *args, **kwargs):
    """TODO: Docstring for run_hook.

    :event: TODO
    :returns: TODO

    """

    if event == event.RECORD_TRANSFORM_STREAM:
        source = kwargs.get('source')
        # sources = kwargs.get('sources')
        # target = kwargs.get('target')


        # for f in RECORD_TRANSFORM_STREAMS:

        for f in source.record_transforms:
            transform_function = RECORD_TRANSFORMS_STREAM.get(f)

            if not transform_function:
                logger.error('function %s not defined for transform', f)
                raise Exception('function %s not defined for transform', f)

            logger.info('running record transform %s for %s', f, source.source_name)
            transform_function(*args, **kwargs)

        # source._stream = source.stream.rowmap(modify_rows(source), petl.header(source.stream))

        # new_stream = modify_rows(source)
        # print(source.stream)
        # import bpdb; bpdb.set_trace()  # noqa: E702

    elif event == event.RECORD_TRANSFORM:
        # sources = kwargs.get('sources')

        for source_name, source in kwargs.get('sources').items():
            # source = kwargs.get('source')

            for f in source.record_transforms_non_stream:
                transform_function = RECORD_TRANSFORMS.get(f)

                if not transform_function:
                    logger.error('function %s not defined for transform', f)
                    raise Exception('function %s not defined for transform', f)

                logger.info('running record transform %s for %s', f, source.source_name)

                for r in source.records:
                    transform_function(r, *args, **kwargs)
                    # print(r)
                    # print(transform_function)

            # source._stream = source.stream.rowmap(modify_rows(source), petl.header(source.stream))

            # new_stream = modify_rows(source)
            # print(source.stream)

    else:
                                            
        logger.debug('running %s %s hooks',
                     len(HOOK_EVENTS[event]['hooks']),
                     event)


        for hook in HOOK_EVENTS[event]['hooks']:
            # hook(*args, source=source, sources=sources, **kwargs)
            hook(*args, **kwargs)

        return HOOK_EVENTS[event]['results']
