from typing import (
    Callable,
    Union,
    Dict,
)

from messagebus import (
    Message,
    AckMessage,
    MessageBusException,
)
from ._SendMessageMock import SendMessageMock
from ._SendMessageRequestValidator import SendMessageRequestValidator


class SendMessageMockBuilder(object):
    def __init__(
        self,
    ):
        self.component = "InvalidComponentName"
        self.token = None
        self.data_type = None
        self.data_validator = None
        self.headers_validator = None
        self.callback = None
        self.response = AckMessage()

    def when_sent_to(self, component: str):
        self.component = component
        return self

    def when_sent_to_any_component(self):
        self.component = None
        return self

    def with_token(self, token: int):
        self.token = token
        return self

    def with_data_type(self, data_type: int):
        self.data_type = data_type
        return self

    def with_data(self, data_validator: Callable[[bytes], bool]):
        self.data_validator = data_validator
        return self

    def with_headers(self, headers_validator: Callable[[Dict[str, str]], bool]):
        self.headers_validator = headers_validator
        return self

    def will_reply_with(self, response: Union[Message, MessageBusException, Callable[[Message], Message]]):
        self.response = response
        return self

    def with_callback(self, callback: Callable[[Message], bool]):
        self.callback = callback
        return self

    def build(self):
        return SendMessageMock(
            request_validator=SendMessageRequestValidator(
                component_name=self.component,
                token=self.token,
                data_type=self.data_type,
                data_validator=self.data_validator,
                headers_validator=self.headers_validator,
                callback=self.callback,
            ),
            response=self.response
        )
