#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
from collections import namedtuple, OrderedDict
from future.utils import iteritems

from thrift.Thrift import TType

__all__ = ['create_client', 'Serializer', 'struct_to_dict', 'parse_struct_spec']
StructField = namedtuple('StructField',
                         'id type name type_args default req_type')


def create_client(client_klass, host=None, port=None, client_type=None):
    """
    Given a thrift client class, and a host/port
    return a client using HeaderTransport
    """
    from thrift.transport.TSocket import TSocket
    from thrift.protocol.THeaderProtocol import THeaderProtocol

    sock = TSocket(host=host, port=port)
    protocol = THeaderProtocol(
        sock,
        client_types=client_type,  # We accept the same as our inital send_
        client_type=client_type  # Used for the inital send_
    )
    sock.open()
    return client_klass(protocol)


def parse_struct_spec(struct):
    """
    Given a thrift struct return a generator of parsed field information

    StructField fields:
        id - the field number
        type - a Thrift.TType
        name - the field name
        type_args - type arguments (ex: Key type Value type for maps)
        default - the default value
        req_type - the field required setting
            (0: Required, 1: Optional, 2: Optional IN, Required OUT)

    :param struct: a thrift struct
    :return: a generator of StructField tuples
    """
    for field in struct.thrift_spec:
        if not field:
            continue
        yield StructField._make(field)


def struct_to_dict(struct, defaults=False):
    """
    Given a Thrift Struct convert it into a dict
    :param struct: a thrift struct
    :param defaults: return default values
    :return: OrderedDict
    """
    adict = OrderedDict()
    union = struct.isUnion()
    if union and struct.field == 0:
        # if struct.field is 0 then it is unset escape
        return adict
    for field in parse_struct_spec(struct):
        if union:
            if field.id == struct.field:
                value = struct.value
            else:
                continue
        else:
            value = getattr(struct, field.name, field.default)
        if value != field.default or defaults:
            if field.type == TType.STRUCT:
                if value is not None:
                    sub_dict = struct_to_dict(value, defaults=defaults)
                    if sub_dict or defaults:  # Do not include empty sub structs
                        adict[field.name] = sub_dict
            elif field.type == TType.LIST or field.type == TType.SET:
                sub_list = __list_to_dict(value,
                                          field.type_args,
                                          defaults=defaults)
                if sub_list or defaults:
                    adict[field.name] = sub_list
            elif field.type == TType.MAP:
                sub_map = __map_to_dict(value,
                                        field.type_args,
                                        defaults=defaults)
                if sub_map or defaults:
                    adict[field.name] = sub_map
            else:
                adict[field.name] = value
        if union:  # If we got this far then we have the union value
            break
    return adict


def __list_to_dict(alist, type_args, defaults=False):
    """
    Given a python list-like collection, potentially containing Thrift Structs,
    convert it into a dict
    :param alist: a list or set
    :param defaults: return default values
    :return: OrderedDict
    """
    if not alist:
        return alist

    element_type = type_args[0]
    if element_type == TType.STRUCT:
        return [struct_to_dict(element, defaults=defaults) for element in alist]
    if element_type == TType.LIST or element_type == TType.SET:
        sub_list = [__list_to_dict(element, type_args[1], defaults=defaults)
                    for element in alist]
        return set(sub_list) if element_type == TType.SET else sub_list
    else:
        return alist


def __map_to_dict(amap, type_args, defaults=False):
    """
    Given a python dictionary, potentially containing Thrift Structs, convert it
    into a dict
    :param alist: a list or set
    :param defaults: return default values
    :return: OrderedDict
    """
    if not amap:
        return amap

    keys, values = zip(*iteritems(amap))
    keys = __list_to_dict(keys, type_args[:2], defaults=defaults)
    values = __list_to_dict(values, type_args[2:4], defaults=defaults)

    return dict(zip(keys, values))
