"""Backend module to transform a python module into a pydantic model.

This module defines the main model in the demessaging framework. It takes a
list of members, or a module, and creates a new Model that can be used to
generate code, connect to the pulsar, and more. See :class:`BackendModule` for
details.
"""
from __future__ import annotations

import atexit
import base64
import datetime as dt
import inspect
import io
from functools import partial
from importlib import import_module
from typing import (
    IO,
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    List,
    Optional,
    Type,
    Union,
    cast,
)

import docstring_parser
from deprogressapi import BaseReport
from pydantic import Field  # pylint: disable=no-name-in-module
from pydantic import BaseModel, ValidationError, create_model
from pydantic.json import (  # pylint: disable=no-name-in-module
    custom_pydantic_encoder,
)

from demessaging.backend import utils
from demessaging.backend.class_ import BackendClass
from demessaging.backend.function import BackendFunction
from demessaging.config import ModuleConfig, append_parameter_docs
from demessaging.messaging.consumer import MessageConsumer
from demessaging.PulsarMessageConstants import PropertyKeys, Status


@append_parameter_docs
class BackendModuleConfig(ModuleConfig):
    """Configuration class for a backend module."""

    # it should be Union[Type[BackendFunction], Type[BackendClass]], but
    # this is not supported by pydantic

    if TYPE_CHECKING:
        models: List[Union[Type[BackendFunction], Type[BackendClass]]]

    models: List[Any] = Field(  # type: ignore
        default_factory=list,
        description=(
            "a list of function or class models for the members of the "
            "backend module"
        ),
    )

    module: Any = Field(
        description="The imported backend module (or none, if there is none)"
    )

    class_name: str = Field(description="Name of the model class")


ModuleMember = Union[
    Type[BackendFunction], Type[BackendClass], Callable, str, Type[object]
]


@append_parameter_docs
class BackendModule(BaseModel):
    """A base class for a backend module.

    Do not directly instantiate from this class, rather use the
    :meth:`create_model` method.
    """

    backend_config: ClassVar[BackendModuleConfig]
    pulsar: ClassVar[MessageConsumer]

    # type that is implemented by subclasses
    __root__: Union[BackendFunction, BackendClass]

    if TYPE_CHECKING:
        # added properties for subclasses generated by create_model
        member: Union[BackendClass, BackendFunction]

    def __call__(self) -> BaseModel:
        """Call the selected member of this backend module."""
        return self.__root__()  # type: ignore

    @classmethod
    def create_model(
        cls,
        module_name: Optional[str] = None,
        members: Optional[List[ModuleMember]] = None,
        config: Optional[ModuleConfig] = None,
        class_name: Optional[str] = None,
        **config_kws,
    ) -> Type[BackendModule]:
        """Generate a module for a backend module.

        Parameters
        ----------
        module_name: str
            The name of the module to import. If none is given, the `members`
            must be specified
        members: list of members
            The list of members that shall be added to this module. It can be
            a list of

            - :class:`~demessaging.backend.function.BackendFunction` classes (
              generated with
              :meth:`~demessaging.backend.function.BackendFunction.create_model`)
            - :class:`~demessaging.backend.class_.BackendClass` classes (
              generated with
              :meth:`~demessaging.backend.class_.BackendClass.create_model`)
            - functions (that will then be transformed using
              :meth:`~demessaging.backend.function.BackendFunction.create_model`)
            - classes (that will then be transformed using
              :meth:`~demessaging.backend.class_.BackendClass.create_model`)
            - strings, in which case they point to the member of the given
              `module_name`
        config: ModuleConfig, optional
            The configuration for the module. If this is not given, you must
            provide ``config_kws`` or define a ``backend_config`` variable
            within the module corresponding to `module_name`
        class_name: str, optional
            The name for the generated subclass of :class:`pydantic.BaseModel`.
            If not given, the name of `Class` is used
        ``**config_kws``
            An alternative way to specify the configuration for the backend
            module.

        Returns
        -------
        Subclass of BackendFunction
            The newly generated class that represents this module.
        """
        if module_name is not None:
            module: Any = import_module(module_name)
        else:
            module = None

        if members is None and module is None:
            raise ValueError("Either members or module need to be provided!")

        if config and config_kws:
            raise ValueError("Either config or config_kws can be used!")
        if config_kws:
            config = ModuleConfig(**config_kws)
        elif module is not None and hasattr(module, "backend_config"):
            config = module.backend_config

        config = cast(ModuleConfig, config)

        # this should not be camelized
        class_name = class_name or module_name or config.messaging_config.topic

        assert config is not None
        config = BackendModuleConfig(
            module=module, class_name=class_name, **config.copy().dict()
        )

        if not members:
            members = list(config.members)
        if not members:
            assert module is not None
            if hasattr(module, "__all__"):
                members = list(module.__all__)
            else:
                functions = inspect.getmembers(
                    module, predicate=inspect.isfunction
                )
                classes = inspect.getmembers(
                    module, predicate=inspect.isfunction
                )
                members = [t[1] for t in functions if not t[0].startswith("_")]
                members += [t[1] for t in classes if not t[0].startswith("_")]

        # finally check if we have any members
        if not members:
            raise ValueError(
                f"Found no members for the given module {module_name}!"
            )
        models: List[Union[Type[BackendFunction], Type[BackendClass]]] = []

        for i, member in enumerate(list(members)):
            member_obj: ModuleMember
            member_model: Union[Type[BackendFunction], Type[BackendClass]]
            if isinstance(member, str):
                member = getattr(module, member)
            if inspect.isclass(member) and issubclass(
                member, (BackendFunction, BackendClass)  # type: ignore
            ):
                member = cast(
                    Union[Type[BackendFunction], Type[BackendClass]], member
                )
                member_model = member
                member_obj = (
                    member.backend_config.Class
                    if issubclass(member, BackendClass)
                    else member.backend_config.function
                )
            elif inspect.isclass(member):
                member_model = BackendClass.create_model(member)
                member_obj = member
            elif callable(member):
                member_model = BackendFunction.create_model(member)
                member_obj = member
            else:
                raise ValueError(
                    f"Cannot transform {member} to a member model!"
                )
            members[i] = member_obj
            models.append(member_model)

        config.members = members
        config.models = models

        if not config.doc and module:
            docstring = docstring_parser.parse(module.__doc__)
            config.doc = utils.get_desc(docstring)

        member_types = models[0]
        for model in models[1:]:
            member_types = Union[member_types, model]  # type: ignore

        kws = {"__module__": module_name} if module_name else {}

        Model: Type[BackendModule] = create_model(  # type: ignore
            class_name,
            __base__=cls,
            __root__=(member_types, Field(description="The member to call.")),
            **kws,  # type: ignore
        )

        Model.__config__.title = config.messaging_config.topic  # type: ignore

        if config.registry.json_encoders:
            # it would be better, to set this via __config__ in create_model,
            # but this is not possible if we use `__base__`
            Model.__config__.json_encoders = config.registry.json_encoders
            Model.__json_encoder__ = partial(
                custom_pydantic_encoder, config.registry.json_encoders
            )

        Model.backend_config = config

        if module is not None:
            config.imports += "\n" + utils.get_module_imports(module)

        Model.__doc__ = config.doc

        return Model

    @classmethod
    def test_connect(cls):
        """Connect to the message pulsar."""
        cls.pulsar = pulsar = MessageConsumer(
            pulsar_config=cls.backend_config.messaging_config,
            handle_request=cls.handle_message,
            module_info=cls.schema(),
        )
        atexit.register(pulsar.disconnect)

        pulsar.connect()

    @classmethod
    def listen(cls):
        """Connect to the message pulsar."""
        cls.pulsar = pulsar = MessageConsumer(
            pulsar_config=cls.backend_config.messaging_config,
            handle_request=cls.handle_message,
            module_info=cls.schema(),
        )
        atexit.register(pulsar.disconnect)

        pulsar.connect()
        pulsar.wait_for_request()

    @classmethod
    async def _send_request(cls, request: Dict[str, Any]) -> Any:
        """Send a request to the backend module.

        Parameters
        ----------
        request: dict
            A request to the backend module.
        """
        from demessaging.messaging.producer import MessageProducer

        producer = MessageProducer(cls.backend_config.messaging_config)

        response = await producer.send_request(request)
        status = response[PropertyKeys.STATUS]
        if status == Status.SUCCESS:
            print("request successful")
            return response["msg"]
        elif status == Status.ERROR:
            print("request failed")
            print(response["msg"])
            raise ValueError(response["error"])

    @classmethod
    def send_request(
        cls: Type[BackendModule],
        request: Union[BackendModule, IO, Dict[str, Any]],
    ) -> BaseModel:
        """Test a request to the backend.

        Parameters
        ----------
        request: dict or file-like object
            A request to the backend module.
        """
        if isinstance(request, io.IOBase):
            model = cls.parse_raw("\n".join(request.readlines()))
        elif hasattr(request, "__root__"):
            request = cast(BackendModule, request)
            model = cls.parse_obj(request.__root__)
        else:
            model = cls.parse_obj(request)
        payload = base64.b64encode(model.json().encode("utf-8")).decode(
            "utf-8"
        )
        request = {
            "properties": {},
            "payload": payload,
        }
        result = utils.run_async(cls._send_request, request)

        return model.__root__.return_model.parse_raw(result)

    def compute(self) -> BaseModel:
        """Send this request to the backend module and compute the result.

        This method updates the model inplace.
        """
        response = self.send_request(self)
        return response

    @classmethod
    def shell(cls):
        """Start a shell with the module defined."""
        from IPython import start_ipython

        start_ipython(argv=[], user_ns=dict(Model=cls))

    @classmethod
    def generate(
        cls,
        line_length: int = 79,
        use_formatters: bool = True,
        use_autoflake: bool = True,
        use_black: bool = True,
        use_isort: bool = True,
    ) -> str:
        """Generate the code for the frontend module."""
        import autoflake
        import black
        import isort

        code = cls.backend_config.render()

        if use_formatters:

            if use_isort:
                code = isort.code(code, float_to_top=True, profile="black")
            if use_black:
                code = black.format_str(
                    code, mode=black.Mode(line_length=line_length)
                )

            # remove unused imports
            if use_autoflake:
                code = autoflake.fix_code(code, remove_all_unused_imports=True)

            if use_isort:
                code = isort.code(code, float_to_top=True, profile="black")

        if cls.backend_config.module:
            # remove __main__, etc.
            name = cls.backend_config.module.__name__
            code = code.replace(name + ".", "")

        return code.strip() + "\n"

    @classmethod
    def handle_message(cls, request_msg):
        print("[{}] processing request".format(dt.datetime.now()))

        payload = base64.b64decode(request_msg["payload"]).decode("utf-8")

        try:
            model = cls.parse_raw(payload)
        except ValidationError as e:
            cls.pulsar.send_error(
                request=request_msg,
                error_message="error validating request: {0}".format(e),
            )
        except Exception as e:
            cls.pulsar.send_error(
                request=request_msg,
                error_message="error processing request: {0}".format(e),
            )
        else:
            try:
                reporter_args = model.__root__.backend_config.reporter_args
                for key, reporter in reporter_args.items():
                    member_reporter = getattr(model.__root__, key)
                    if member_reporter and isinstance(
                        member_reporter, BaseReport
                    ):
                        member_reporter._pulsar = cls.pulsar
                        member_reporter._request = request_msg
                result = model()
            except Exception as e:
                cls.pulsar.send_error(
                    request=request_msg,
                    error_message="error executing request: {0}".format(e),
                )
            else:
                cls.pulsar.send_response(
                    request=request_msg,
                    response_properties={PropertyKeys.STATUS: Status.SUCCESS},
                    response_payload=result.json(),
                )


ModuleConfig.update_forward_refs()
