from typing import List
from typing import Tuple
from typing import Union
import syft as sy
from syft.generic.pointers.string_pointer import StringPointer
from syft.workers.base import BaseWorker
from syft.generic.object import AbstractObject
from syft.generic.frameworks.overload import overloaded
from syft.generic.frameworks.hook import hook_args
import abc


class String(AbstractObject):
    """
       This is a class that wraps the Python built-in `str` class. In addition to 
       providing access to most of `str`'s method call API, it allows sending
       such wrapped string between workers the same way Syft tensors can be
       moved around among workers.
    """

    # Set of methods from 'str' to hook/override by String
    methods_to_hook = set(
        [
            "__add__",
            "__eq__",
            "__le__",
            "__ge__",
            "__gt__",
            "__lt__",
            "__ne__",
            "__len__",
            "__getitem__",
            "__str__",
            "__repr__",
            "__format__",
            "lower",
            "upper",
            "capitalize",
            "casefold",
            "center",
            "count",
            "encode",
            "endswith",
            "expandtabs",
            "find",
            "format",
            "format_map",
            "index",
            "isalnum",
            "isalpha",
            "isascii",
            "isdecimal",
            "isdigit",
            "isidentifier",
            "islower",
            "isnumeric",
            "isprintable",
            "isspace",
            "istitle",
            "isupper",
            "join",
            "ljust",
            "lstrip",
            "maketrans",
            "partition",
            "replace",
            "rfind",
            "rindex",
            "rjust",
            "rpartition",
            "rsplit",
            "rstrip",
            "split",
            "splitlines",
            "startswith",
            "strip",
            "swapcase",
            "title",
            "translate",
            "zfill",
            "__mod__",
        ]
    )

    def __init__(
        self,
        object: object = None,
        encoding: str = None,
        errors: str = None,
        id: Union[int, str] = None,
        owner: BaseWorker = None,
        tags: List[str] = None,
        description: str = None,
    ):
        """Initialize a String object.

           Args:
               object: This could be any object whose string representation,i.e.,
                   the output of its __str__() method is to be wrapped as a 
                   String object.
               encoding: This should be specified if the above `object` argument is
                   a bytes-like object. It specifies the encoding scheme used to create the
                   bytes-like object b''. for example, encoding could be 'utf-8'.
                   For more details on this argument, please  check the official `str` documentation.
               errors: This should be specified if the above `object` argument is
                   a bytes-like object. Possible values are 'strict', 'ignore' or
                   'replace'. For more details on this argument, please
                   check the official `str` documentation.
               id: An optional string or integer id of the String object
               owner: An optional BaseWorker object to specify the worker on which
                   the String object is located.
               tags: an optional set of hashtags corresponding to this object.
                   They are useful when search for this object.
               description: an optional string describing the purpose of this
                   String object

        """

        # get the specified kwargs for creating the base 'str'
        # class

        self.encoding = encoding
        self.errors = errors

        # String objects have normally a default owner which is the
        # local worker. So prevent 'None' as owner
        if self.owner is None or owner is not None:
            self.owner = owner

        str_kwargs = {}

        if object:
            str_kwargs["object"] = object

        if encoding:
            str_kwargs["encoding"] = encoding

        if errors:
            str_kwargs["errors"] = errors

        # Create a str instance as the 'child' attribute
        self.child = str(**str_kwargs)

        super(String, self).__init__(
            id=id, owner=self.owner, tags=tags, description=description, child=self.child
        )

    def send(self, location: BaseWorker):
        """
           Sends this String object to the worker specified by 'location'. 
           and returns a pointer to that string as a StringPointer object.

           Args:
               location: The BaseWorker object which you want to send this object
                         to. Note that this is never actually the BaseWorker but instead
                         a class which inherits the BaseWorker abstraction.

           Returns:
               A StringPointer objects to self.

        """

        ptr = self.owner.send(self, location)

        return ptr

    def get_class_attributes(self):
        """returns minimal necessary keyword arguments to create a 
           String object
        """
        kwargs = dict(owner=self.owner)

        return kwargs

    def on(self, object: str, wrap=False):
        """Takes and object of type strings and assigns it to
           self.child
        """
        self.child = object

        return self

    def __add__(self, other: Union[str, "String"]):
        """
            [Important] overriding the `__add__` here is not yet
            activated. The real hooking happens in
            syft/generic/frameworks/hook/hook.py.
            Hooking as implemented here (using @overloaded.method)
            is to be activated when hook_args.py is adapted
            to wrapping reponses of `str` types into `String`
            types. This is not yet supported.
        """

        # The following is necessary in order to adapt the
        # below `add_string` method to the args hooking logic in
        # hook_args.py. Please check the doc string of `add_string`
        # to know more.
        if isinstance(other, str):
            other = String(other)

        return self.add_string(other)

    @overloaded.method
    def add_string(self, _self: "String", other: "String"):
        """This method is created in a way adapted to the logic implemented
           in hook_args.py. That is, it can be wrapped with the decorator
           @oerloaded.method.
           
           hook_args.py args hooking logic needs that the data types of 
           argument be unchanged. For instance, 'other' should always
           be of a fixed type 'String' or 'str' but not alternating
           between both. This can cause unexpected behaviou due to caching
           in hook_args.py.

           Args:
               _self: a String object (as received by the decorator).
                      It represents the objects on which we called the add method.
                      It will always be of type `str` inside this method. Since
                      the decorator methods strips the `str` out of the `String`
                      object.
               other: a String object that we wish to concatenate to `_self`.
                      Same as above, it is a String object as received by the 
                      decorator but here it will always be of type `str`.

           Returns:
               The concatentenated `str` object between `_self` and `other`.
               this `str` object will be wrapped by the decorator into a 
               String object
        """

        return _self + other

    @staticmethod
    def create_pointer(
        obj,
        location: BaseWorker = None,
        id_at_location: (str or int) = None,
        register: bool = False,
        owner: BaseWorker = None,
        ptr_id: (str or int) = None,
        garbage_collect_data: bool = True,
    ):
        """
           Creates a StringPointer object that points to a String object 'obj'
           after sending the latter to the worker 'location'.

           Returns:
                  a StringPointer object
        """

        if id_at_location is None:
            id_at_location = obj.id

        if owner is None:
            owner = obj.owner

        string_pointer = StringPointer(
            location=location,
            id_at_location=id_at_location,
            owner=owner,
            id=ptr_id,
            garbage_collect_data=garbage_collect_data,
        )

        return string_pointer

    @staticmethod
    def simplify(worker: BaseWorker, string: "String"):
        """
           Breaks String object into a tuple of simpler objects, its constituting objects that are
           serializable.


           Args:

               worker: a BaseWorker object
               string: the String object to be simplified

           Returns:
               
               A tuple of simpler objects that are sufficient to recreate
               a String object that is a clone of `string`.

        """

        # Encode the string into a bytes object
        simple_child = sy.serde.msgpack.serde._simplify(worker, string.child)
        tags = sy.serde.msgpack.serde._simplify(worker, string.tags)
        description = sy.serde.msgpack.serde._simplify(worker, string.description)

        return (simple_child, string.id, tags, description)

    @staticmethod
    def detail(worker: BaseWorker, simple_obj: Tuple):
        """
           Create an object of type String from the reduced representation in `simple_obj`.


           Args:

               worker: BaseWorker
                       The worker on which the new String object is to be created.
               simple_obj: tuple
                           A tuple resulting from the serialized then deserialized returned tuple
                           from the `simplify` static method above.

           Returns:

               A String object


        """

        # Get the contents of the tuple represening the simplified object
        simple_child, id, tags, description = simple_obj

        # It appears that all strings are converted to bytes objects
        # after deserialization, convert them back to strings
        tags = sy.serde.msgpack.serde._detail(worker, tags)
        description = sy.serde.msgpack.serde._detail(worker, description)

        # Rebuild the str child our of the simplified child (the bytes child)
        child = sy.serde.msgpack.serde._detail(worker, simple_child)

        return String(object=child, id=id, owner=worker, tags=tags, description=description)


### Register the String object with hook_args.py ###
hook_args.default_register_tensor(String)
