"""
Contracts to enforce loader objects
"""
import asyncio
import functools
import inspect
import os
from collections import namedtuple
from types import ModuleType
from typing import Dict
from typing import Iterable
from typing import List

import pop.exc
import pop.hub
import pop.verify


class ContractedContext(
    namedtuple(
        "ContractedContext",
        ("func", "args", "kwargs", "signature", "ret", "cache", "ref"),
    )
):
    """
    Contracted function calling context
    """

    def __new__(
        cls,
        func: functools.partial,
        args: Iterable,
        kwargs: Dict,
        signature,
        ref,
        ret=None,
        cache=None,
    ):  # pylint: disable=too-many-arguments
        if cache is None:
            cache = {}
        return super().__new__(
            cls, func, list(args), kwargs, signature, ret, cache, ref
        )

    def get_argument(self, name):
        """
        Return the value corresponding to a function argument after binding the contract context
        argument and keyword arguments to the function signature.
        """
        return self.get_arguments()[name]

    def get_arguments(self):
        """
        Return a dictionary of all arguments that will be passed to the function and their
        values, including default arguments.
        """
        if "__bound_signature__" not in self.cache:
            try:
                self.cache["__bound_signature__"] = self.signature.bind(
                    *self.args, **self.kwargs
                )
            except TypeError as e:
                for frame in inspect.trace(0):
                    if frame.function == "bind" and frame.filename.endswith(
                        os.sep + "inspect.py"
                    ):
                        raise pop.exc.BindError(e)
                raise
            # Apply any default values from the signature
            self.cache["__bound_signature__"].apply_defaults()
        return self.cache["__bound_signature__"].arguments


def load_contract(
    contracts: List["Contracted"],
    default_contracts: List[str],
    mod: ModuleType,
    name: str,
) -> List:
    """
    return a Contract object loaded up
    Dynamically create the correct Contracted type
    :param contracts: Contracts functions to add to the sub
    :param default_contracts: The contracts that have been marked as defaults
    :param mod: A loader module
    :param name: The name of the module to get from the loader
    """
    raws = []
    if not contracts:
        return raws
    loaded_contracts = []
    if hasattr(contracts, name):
        loaded_contracts.append(name)
        raws.append(getattr(contracts, name))
    if hasattr(contracts, "init"):
        loaded_contracts.append("init")
        raws.append(getattr(contracts, "init"))
    if default_contracts:
        for contract in default_contracts:
            if contract in loaded_contracts:
                continue
            loaded_contracts.append(contract)
            raws.append(getattr(contracts, contract))
    if hasattr(mod, "__contracts__"):
        cnames = getattr(mod, "__contracts__")
        if not isinstance(cnames, (list, tuple)):
            cnames = cnames.split(",")
        for cname in cnames:
            if cname in contracts:
                if cname in loaded_contracts:
                    continue
                loaded_contracts.append(cname)
                raws.append(getattr(contracts, cname))
    return raws


class Wrapper:
    def __init__(self, func: functools.partial, ref: str, name: str):
        """
        :param func: The contracted function to call
        :param ref: The reference to the function on the hub
        :param name: An alias for the function
        """
        self.__dict__.update(
            getattr(func, "__dict__", {})
        )  # do this first so we later overwrite any conflicts
        self.func = func
        self.ref = ref
        self.__name__ = name
        try:
            self.signature = inspect.signature(self.func)
        except ValueError:
            # inspect.signature raises a ValueError only when a signature could not be provided.
            # Some callables may not be introspectable in certain implementations of Python.
            # For example, in CPython, some built-in functions defined in C provide no metadata about their arguments.
            self.signature = None
        self._sig_errors = []
        self.__wrapped__ = func

    def __call__(self, *args, **kwargs):
        self.func(*args, **kwargs)

    def __repr__(self):
        return "<{} func={}.{}>".format(
            self.__class__.__name__, self.func.__module__, self.__name__
        )


class Contracted(Wrapper):
    """
    This class wraps functions that have a contract associated with them
    and executes the contract routines
    """

    def __init__(
        self,
        hub: "pop.hub.Hub",
        contracts: List["pop.hub.Sub"],
        func: functools.partial,
        ref: str,
        name: str,
        implicit_hub: bool = True,
    ):
        """
        :param hub: The redistributed pop central hub
        :param contracts: Contracts functions to add to the sub
        :param func: The contracted function to call
        :param ref: The reference to the function on the hub
        :param name: An alias for the function
        :param implicit_hub: True if a hub should be implicitly injected into the "call" method
        """
        super().__init__(func, ref, name)
        self.hub = hub
        self.contracts = contracts or []
        self.implicit_hub = implicit_hub
        self._load_contracts()

    def _get_contracts_by_type(self, contract_type: str = "pre") -> List[Wrapper]:
        """
        :param contract_type: One of "call", "pre", "post", or "sig"
        """
        matches = []
        fn_contract_name = f"{contract_type}_{self.__name__}"
        for contract in self.contracts:
            if hasattr(contract, fn_contract_name):
                matches.append(getattr(contract, fn_contract_name))
            if hasattr(contract, contract_type):
                matches.append(getattr(contract, contract_type))

        return matches

    def _load_contracts(self):
        # TODO:
        # if Contracted - only allow regular pre/post
        # if ContractedAsync - allow coroutines and functions
        # if ContractedAsyncGen - allow coroutines and functions

        self.contract_functions = {
            "pre": self._get_contracts_by_type("pre"),
            "call": self._get_contracts_by_type("call")[:1],
            "post": self._get_contracts_by_type("post"),
        }
        # TODO: write test for stack-like behavior (reverse "pre")
        self._has_contracts = sum(len(l) for l in self.contract_functions.values()) > 0

    def __call__(self, *args, **kwargs):
        if self.implicit_hub:
            args = (self.hub,) + args
        if not self._has_contracts:
            return self.func(*args, **kwargs)
        contract_context = ContractedContext(
            self.func, args, kwargs, self.signature, self.ref
        )

        for fn in self.contract_functions["pre"]:
            fn(contract_context)
        if self.contract_functions["call"]:
            ret = self.contract_functions["call"][0](contract_context)
        else:
            ret = self.func(*contract_context.args, **contract_context.kwargs)
        for fn in self.contract_functions["post"]:
            post_ret = fn(contract_context._replace(ret=ret))
            if post_ret is not None:
                ret = post_ret

        return ret

    def __getstate__(self):
        return dict(
            hub=self.hub,
            ref=self.ref,
            name=self.__name__,
            implicit_hub=self.implicit_hub,
            contracts=self.contracts,
        )

    def __setstate__(self, state):
        hub = state["hub"]
        ref = state["ref"]
        name = state["name"]
        func = hub[ref][name].func
        self.__init__(
            hub=hub,
            func=func,
            ref=ref,
            name=name,
            implicit_hub=state["implicit_hub"],
            contracts=state["contracts"],
        )


class ContractedAsyncGen(Contracted):
    async def __call__(self, *args, **kwargs):
        if self.implicit_hub:
            args = (self.hub,) + args
        if not self._has_contracts:
            async for chunk in self.func(*args, **kwargs):
                yield chunk
            return
        contract_context = ContractedContext(
            self.func, args, kwargs, self.signature, self.ref
        )

        for fn in self.contract_functions["pre"]:
            pre_ret = fn(contract_context)
            if asyncio.iscoroutine(pre_ret):
                await pre_ret
        chunk = None
        if self.contract_functions["call"]:
            async for chunk in self.contract_functions["call"][0](contract_context):
                yield chunk
        else:
            async for chunk in self.func(
                *contract_context.args, **contract_context.kwargs
            ):
                yield chunk
        ret = chunk
        for fn in self.contract_functions["post"]:
            if isinstance(fn, ContractedAsync):
                post_ret = await fn(contract_context._replace(ret=ret))
            else:
                post_ret = fn(contract_context._replace(ret=ret))
            if post_ret is not None:
                ret = post_ret


class ContractedAsync(Contracted):
    async def __call__(self, *args, **kwargs):
        if self.implicit_hub:
            args = (self.hub,) + args
        if not self._has_contracts:
            return await self.func(*args, **kwargs)
        contract_context = ContractedContext(
            self.func, args, kwargs, self.signature, self.ref
        )

        for fn in self.contract_functions["pre"]:
            pre_ret = fn(contract_context)
            if asyncio.iscoroutine(pre_ret):
                await pre_ret
        if self.contract_functions["call"]:
            ret = await self.contract_functions["call"][0](contract_context)
        else:
            ret = await self.func(*contract_context.args, **contract_context.kwargs)
        for fn in self.contract_functions["post"]:
            if isinstance(fn, ContractedAsync):
                post_ret = await fn(contract_context._replace(ret=ret))
            else:
                post_ret = fn(contract_context._replace(ret=ret))
            if post_ret is not None:
                ret = post_ret

        return ret


def create_contracted(
    hub: "pop.hub.Hub",
    contracts: List["pop.loader.LoadedMod"],
    func: functools.partial,
    ref: str,
    name: str,
    implicit_hub: bool = True,
) -> Contracted:
    """
    Dynamically create the correct Contracted type
    :param hub: The redistributed pop central hub
    :param contracts: Contracts functions to add to the sub
    :param func: The contracted function to call
    :param ref: The reference to the function on the hub
    :param name: The name of the module to get from the loader
    :param implicit_hub: True if a hub should be implicitly injected into the "call" method
    """
    if asyncio.iscoroutinefunction(func):
        return ContractedAsync(hub, contracts, func, ref, name, implicit_hub)
    elif inspect.isasyncgenfunction(func):
        return ContractedAsyncGen(hub, contracts, func, ref, name, implicit_hub)
    else:
        return Contracted(hub, contracts, func, ref, name, implicit_hub)


class Partial:
    """
    A partial function that can be used for multiprocessing
    """

    def __init__(self, synchronous_function: Contracted, *args, **kwargs):
        self.func = synchronous_function
        self.args = args
        self.kwargs = kwargs

    def __call__(self):
        ret = self.func(*self.args, **self.kwargs)
        # We can't pickle a coroutine, so async functions need to finish here
        # if this is executed in a forked process, we can't guarantee a loop exists
        if asyncio.iscoroutine(ret):
            loop = asyncio.new_event_loop()
            try:
                while asyncio.iscoroutine(ret):
                    ret = loop.run_until_complete(ret)
            finally:
                loop.close()
        return ret

    def __getstate__(self):
        return {
            "func": self.func,
            "args": self.args,
            "kwargs": self.kwargs,
        }

    def __setstate__(self, state):
        self.func = state["func"]
        self.args = state["args"]
        self.kwargs = state["kwargs"]
