import cPickle
import pipeline

from importlib import import_module
from mapreduce import context
from mapreduce.mapper_pipeline import MapperPipeline
from mapreduce.mapreduce_pipeline import MapreducePipeline
from mapreduce import pipeline_base
from mapreduce.model import MapreduceState
from mapreduce.input_readers import RawDatastoreInputReader, GoogleCloudStorageInputReader

from django.utils import six
from djangae.contrib.processing.mapreduce.input_readers import DjangoInputReader

from utils import qualname


class DynamicPipeline(pipeline_base.PipelineBase):
    """
        Horrific class which uses pickle to store pipelines for
        yielding via run(). This wouldn't be necessary if the pipeline
        library wasn't built for Java and had a sensible interface that
        didn't require inheritence.
    """
    def __init__(self, pipelines, *args, **kwargs):
        # This gets reinstantiated somewhere with the already-pickled pipelines argument
        # so we prevent double pickling by checking it's not a string
        if not isinstance(pipelines, six.string_types):
            pipelines = str(cPickle.dumps(pipelines))
        super(DynamicPipeline, self).__init__(pipelines, *args, **kwargs)

    def run(self, pipelines):
        with pipeline.InOrder():
            pipelines = cPickle.loads(str(pipelines))
            for pipe in pipelines:
                yield pipe


def import_callable(dotted_path):
    module_path = dotted_path.rsplit(".", 1)[0]
    while module_path:
        try:
            module = import_module(module_path)
            break
        except ImportError:
            module_path = module_path.rsplit(".", 1)[0]
            continue
    else:
        raise ImportError("Module not found in path: {}".format(dotted_path))

    remainder = dotted_path[len(module_path):].lstrip(".")
    remainder_parts = remainder.split(".")

    func = module
    while remainder_parts:
        next_step = remainder_parts[0]
        if not hasattr(func, next_step):
            raise ImportError("Couldn't find {} in module {}".format(next_step, module))
        func = getattr(func, next_step)
        remainder_parts = remainder_parts[1:]

    if not callable(func):
        raise ImportError("Specified path is not a callable: {}".format(dotted_path))

    return func


class CallbackPipeline(pipeline_base.PipelineBase):
    """
        Simply calls the specified function.
        Takes a dotted-path to the callback
    """
    def run(self, func, *args, **kwargs):
        func = import_callable(func)
        func(*args, **kwargs)


def unpacker(obj):
    params = context.get().mapreduce_spec.mapper.params
    handler = import_callable(params["func"])
    yield handler(obj, *params["args"], **params["kwargs"])


def _do_map(
    input_reader, processor_func, finalize_func, params,
    _shards, _output_writer, _output_writer_kwargs, _job_name, _queue_name,
    *processor_args, **processor_kwargs):

    handler_spec = qualname(unpacker)
    handler_params = {
        "func": qualname(processor_func) if callable(processor_func) else processor_func,
        "args": processor_args,
        "kwargs": processor_kwargs
    }

    handler_params.update(params)

    pipelines = []
    pipelines.append(MapperPipeline(
        _job_name,
        handler_spec=handler_spec,
        input_reader_spec=qualname(input_reader),
        output_writer_spec=qualname(_output_writer) if _output_writer else None,
        params=handler_params,
        shards=_shards
    ))

    if finalize_func:
        pipelines.append(
            CallbackPipeline(
                qualname(finalize_func) if callable(finalize_func) else finalize_func,
                *processor_args,
                **processor_kwargs
            )
        )

    new_pipeline = DynamicPipeline(pipelines)
    new_pipeline.start(queue_name=_queue_name or 'default')
    return new_pipeline


def extract_options(kwargs, additional=None):
    VALID_OPTIONS = {
        "_shards",
        "_output_writer",
        "_output_writer_kwargs",
        "_job_name",
        "_queue_name",
    }