# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# Licensed 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 types import MappingProxyType

from libcpp.utility cimport move
from thrift.py3.types cimport CompiledEnum, Struct
from thrift.py3.exceptions cimport GeneratedError
from thrift.py3.server cimport ServiceInterface
from thrift.py3.client cimport Client
from thrift.py3.common cimport MetadataBox
from apache.thrift.metadata.types cimport (
    ThriftMetadata,
    ThriftStruct,
    ThriftException,
    ThriftService,
    ThriftEnum,
    ThriftType,
    ThriftField,
    cThriftMetadata,
    cThriftStruct,
    cThriftException,
    cThriftService,
    cThriftEnum,
    ThriftStructType,
    ThriftUnionType,
    ThriftEnumType,
    ThriftPrimitiveType,
    ThriftListType,
    ThriftSetType,
    ThriftMapType,
    ThriftFunction,
    ThriftTypedefType,
    ThriftSinkType,
    ThriftStreamType,
    ThriftConstStruct,
    ThriftConstValue,
)


cpdef enum ThriftKind:
    PRIMITIVE = 0
    LIST = 1
    SET = 2
    MAP = 3
    ENUM = 4
    STRUCT = 5
    UNION = 6
    TYPEDEF = 7
    STREAM = 8
    SINK = 9


cpdef enum ThriftConstKind:
    CV_BOOL = 0
    CV_INT = 1
    CV_FLOAT = 2
    CV_STRING = 3
    CV_MAP = 4
    CV_LIST = 5
    CV_STRUCT = 6


cdef class ThriftTypeProxy:
    # A union of a bunch of thrift metadata types
    cdef readonly object thriftType
    cdef readonly ThriftMetadata thriftMeta
    cdef readonly ThriftKind kind

    def __init__(self, object thriftType not None, ThriftMetadata thriftMeta not None):
        if not isinstance(thriftType, (
            ThriftStruct,
            ThriftEnum,
            ThriftPrimitiveType,
            ThriftListType,
            ThriftSetType,
            ThriftMapType,
            ThriftTypedefType,
            ThriftSinkType,
            ThriftStreamType,
        )):
            raise TypeError(f"{thriftType!r} is not a known thrift type.")
        self.thriftType = thriftType
        self.thriftMeta = thriftMeta

    @staticmethod
    cdef _fbthrift_create(ThriftType thriftType, ThriftMetadata thriftMeta):
        # Determine value and kind
        if thriftType.type is ThriftType.Type.t_list:
            return ThriftListProxy(thriftType.value, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_set:
            return ThriftSetProxy(thriftType.value, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_map:
            return ThriftMapProxy(thriftType.value, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_enum:
            specialType = ThriftTypeProxy(thriftMeta.enums[thriftType.value.name], thriftMeta)
            specialType.kind = ThriftKind.ENUM
            return specialType
        elif thriftType.type is ThriftType.Type.t_struct:
            return ThriftStructProxy(thriftType.value.name, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_union:
            return ThriftStructProxy(thriftType.value.name, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_typedef:
            return ThriftTypedefProxy(thriftType.value, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_stream:
            return ThriftStreamProxy(thriftType.value, thriftMeta)
        elif thriftType.type is ThriftType.Type.t_sink:
            return ThriftSinkProxy(thriftType.value, thriftMeta)
        specialType = ThriftTypeProxy(thriftType.value, thriftMeta)
        specialType.kind = ThriftKind.PRIMITIVE
        return specialType

    def as_primitive(self):
        if self.kind == ThriftKind.PRIMITIVE:
            return self.thriftType
        raise TypeError('Type is not primitive')

    def as_struct(self):
        if self.kind == ThriftKind.STRUCT or self.kind == ThriftKind.UNION:
            return self
        raise TypeError('Type is not a struct')

    def as_union(self):
        if self.kind == ThriftKind.UNION:
            return self
        raise TypeError('Type is not a union')

    def as_enum(self):
        if self.kind == ThriftKind.ENUM:
            return self.thriftType
        raise TypeError('Type is not an enum')

    def as_list(self):
        if self.kind == ThriftKind.LIST:
            return self
        raise TypeError('Type is not a list')

    def as_set(self):
        if self.kind == ThriftKind.SET:
            return self
        raise TypeError('Type is not a set')

    def as_map(self):
        if self.kind == ThriftKind.MAP:
            return self
        raise TypeError('Type is not a map')

    def as_typedef(self):
        if self.kind == ThriftKind.TYPEDEF:
            return self
        raise TypeError('Type is not a typedef')

    def as_stream(self):
        if self.kind == ThriftKind.STREAM:
            return self
        raise TypeError('Type is not a stream')

    def as_sink(self):
        if self.kind == ThriftKind.SINK:
            return self
        raise TypeError('Type is not a sink')


cdef class ThriftSetProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy valueType

    def __init__(self, ThriftSetType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.SET
        self.valueType = ThriftTypeProxy._fbthrift_create(self.thriftType.valueType, self.thriftMeta)


cdef class ThriftListProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy valueType

    def __init__(self, ThriftListType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.LIST
        self.valueType = ThriftTypeProxy._fbthrift_create(self.thriftType.valueType, self.thriftMeta)


cdef class ThriftMapProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy valueType
    cdef readonly ThriftTypeProxy keyType

    def __init__(self, ThriftMapType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.MAP
        self.valueType = ThriftTypeProxy._fbthrift_create(self.thriftType.valueType, self.thriftMeta)
        self.keyType = ThriftTypeProxy._fbthrift_create(self.thriftType.keyType, self.thriftMeta)


cdef class ThriftTypedefProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy underlyingType
    cdef readonly str name

    def __init__(self, ThriftTypedefType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.TYPEDEF
        self.name = self.thriftType.name
        self.underlyingType = ThriftTypeProxy._fbthrift_create(self.thriftType.underlyingType, self.thriftMeta)


cdef class ThriftSinkProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy elemType
    cdef readonly ThriftTypeProxy initialResponseType
    cdef readonly ThriftTypeProxy finalResponseType

    def __init__(self, ThriftSinkType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.SINK
        self.elemType = ThriftTypeProxy._fbthrift_create(self.thriftType.elemType, self.thriftMeta)
        if self.thriftType.initialResponseType is not None:
            self.initialResponseType = ThriftTypeProxy._fbthrift_create(self.thriftType.initialResponseType, self.thriftMeta)
        if self.thriftType.finalResponseType is not None:
            self.finalResponseType = ThriftTypeProxy._fbthrift_create(self.thriftType.finalResponseType, self.thriftMeta)


cdef class ThriftStreamProxy(ThriftTypeProxy):
    cdef readonly ThriftTypeProxy elemType
    cdef readonly ThriftTypeProxy initialResponseType

    def __init__(self, ThriftStreamType thriftType not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftType, thriftMeta)
        self.kind = ThriftKind.STREAM
        self.elemType = ThriftTypeProxy._fbthrift_create(self.thriftType.elemType, self.thriftMeta)
        if self.thriftType.initialResponseType is not None:
            self.initialResponseType = ThriftTypeProxy._fbthrift_create(self.thriftType.initialResponseType, self.thriftMeta)


cdef class ThriftFieldProxy:
    cdef readonly ThriftTypeProxy type
    cdef readonly ThriftField thriftType
    cdef readonly ThriftMetadata thriftMeta
    cdef readonly int id
    cdef readonly str name
    cdef readonly int is_optional
    cdef readonly tuple structuredAnnotations

    def __init__(self, ThriftField thriftType not None, ThriftMetadata thriftMeta not None):
        self.type = ThriftTypeProxy._fbthrift_create(thriftType.type, thriftMeta)
        self.thriftType = thriftType
        self.thriftMeta = thriftMeta
        self.id = self.thriftType.id
        self.name = self.thriftType.name
        self.is_optional = self.thriftType.is_optional
        self.structuredAnnotations = tuple(ThriftConstStructProxy(annotation) for annotation in self.thriftType.structured_annotations)

    @property
    def pyname(self):
        if self.thriftType.unstructured_annotations is None:
            raise TypeError('The pyname field requires the thrift option `thrift_cpp2_options = ["deprecated_unstructured_annotations_in_metadata"]` to be enabled')
        return self.thriftType.unstructured_annotations.get("py3.name", self.name)



cdef class ThriftStructProxy(ThriftTypeProxy):
    cdef readonly str name
    cdef readonly int is_union
    cdef readonly tuple structuredAnnotations

    def __init__(self, str name not None, ThriftMetadata thriftMeta not None):
        super().__init__(thriftMeta.structs[name], thriftMeta)
        self.name = self.thriftType.name
        self.is_union = self.thriftType.is_union
        self.structuredAnnotations = tuple(ThriftConstStructProxy(annotation) for annotation in self.thriftType.structured_annotations)

        if self.is_union:
            self.kind = ThriftKind.UNION
        else:
            self.kind = ThriftKind.STRUCT

    @property
    def fields(self):
        for field in self.thriftType.fields:
            yield ThriftFieldProxy(field, self.thriftMeta)


cdef class ThriftConstValueProxy:
    cdef readonly ThriftConstValue thriftType
    cdef readonly ThriftConstKind kind
    cdef readonly object type
    def __init__(self, ThriftConstValue value not None):
        self.thriftType = value
        if self.thriftType.type in (ThriftConstValue.Type.cv_bool, ThriftConstValue.Type.cv_integer, ThriftConstValue.Type.cv_double, ThriftConstValue.Type.cv_string):
            self.type = self.thriftType.value
            if self.thriftType.type is ThriftConstValue.Type.cv_bool:
                self.kind = CV_BOOL
            elif self.thriftType.type is ThriftConstValue.Type.cv_integer:
                self.kind = CV_INT
            elif self.thriftType.type is ThriftConstValue.Type.cv_double:
                self.kind = CV_FLOAT
            else:
                self.kind = CV_STRING
        if self.thriftType.type is ThriftConstValue.Type.cv_struct:
            self.type = ThriftConstStructProxy(self.thriftType.value)
            self.kind = CV_STRUCT
        if self.thriftType.type is ThriftConstValue.Type.cv_list:
            self.type = tuple(ThriftConstValueProxy(ele) for ele in self.thriftType.value)
            self.kind = CV_LIST
        if self.thriftType.type is ThriftConstValue.Type.cv_map:
            self.type = MappingProxyType({ThriftConstValueProxy(ele.key).type: ThriftConstValueProxy(ele.value) for ele in self.thriftType.value})
            self.kind = CV_MAP

    def as_bool(self):
        if self.kind == ThriftConstKind.CV_BOOL:
            return self.type
        raise TypeError('Type is not a boolean')

    def as_int(self):
        if self.kind == ThriftConstKind.CV_INT:
            return self.type
        raise TypeError('Type is not an integer')

    def as_float(self):
        if self.kind == ThriftConstKind.CV_FLOAT:
            return self.type
        raise TypeError('Type is not a float')

    def as_string(self):
        if self.kind == ThriftConstKind.CV_STRING:
            return self.type
        raise TypeError('Type is not a string')

    def as_struct(self):
        if self.kind == ThriftConstKind.CV_STRUCT:
            return self.type
        raise TypeError('Type is not a struct')

    def as_list(self):
        if self.kind == ThriftConstKind.CV_LIST:
            return self.type
        raise TypeError('Type is not a list')
    def as_map(self):
        if self.kind == ThriftConstKind.CV_MAP:
            return self.type
        raise TypeError('Type is not a map')

cdef class ThriftConstStructProxy:
    cdef readonly ThriftConstStruct thriftType
    cdef readonly str name
    cdef readonly ThriftKind kind
    def __init__(self, ThriftConstStruct struct not None):
        self.name = struct.type.name
        self.kind = ThriftKind.STRUCT
        self.thriftType = struct

    @property
    def fields(self):
        return MappingProxyType({key: ThriftConstValueProxy(self.thriftType.fields[key]) for key in self.thriftType.fields})


cdef class ThriftExceptionProxy:
    cdef readonly ThriftException thriftType
    cdef readonly ThriftMetadata thriftMeta
    cdef readonly str name
    cdef readonly tuple structuredAnnotations

    def __init__(self, str name not None, ThriftMetadata thriftMeta not None):
        self.thriftType = thriftMeta.exceptions[name]
        self.thriftMeta = thriftMeta
        self.name = self.thriftType.name
        self.structuredAnnotations = tuple(ThriftConstStructProxy(annotation) for annotation in self.thriftType.structured_annotations)

    @property
    def fields(self):
        for field in self.thriftType.fields:
            yield ThriftFieldProxy(field, self.thriftMeta)


cdef class ThriftFunctionProxy:
    cdef readonly str name
    cdef readonly ThriftFunction thriftType
    cdef readonly ThriftMetadata thriftMeta
    cdef readonly ThriftTypeProxy return_type
    cdef readonly int is_oneway
    cdef readonly tuple structuredAnnotations

    def __init__(self, ThriftFunction thriftType not None, ThriftMetadata thriftMeta not None):
        self.name = thriftType.name
        self.thriftType = thriftType
        self.thriftMeta = thriftMeta
        self.return_type = ThriftTypeProxy._fbthrift_create(self.thriftType.return_type, self.thriftMeta)
        self.is_oneway = self.thriftType.is_oneway
        self.structuredAnnotations = tuple(ThriftConstStructProxy(annotation) for annotation in self.thriftType.structured_annotations)

    @property
    def arguments(self):
        for argument in self.thriftType.arguments:
            yield ThriftFieldProxy(argument, self.thriftMeta)

    @property
    def exceptions(self):
        for exception in self.thriftType.exceptions:
            yield ThriftFieldProxy(exception, self.thriftMeta)


cdef class ThriftServiceProxy:
    cdef readonly ThriftService thriftType
    cdef readonly str name
    cdef readonly ThriftMetadata thriftMeta
    cdef readonly ThriftServiceProxy parent
    cdef readonly tuple structuredAnnotations

    def __init__(self, str name not None, ThriftMetadata thriftMeta not None):
        self.thriftType = thriftMeta.services[name]
        self.name = self.thriftType.name
        self.thriftMeta = thriftMeta
        self.parent = None if self.thriftType.parent is None else ThriftServiceProxy(
            self.thriftType.parent,
            self.thriftMeta
        )
        self.structuredAnnotations = tuple(ThriftConstStructProxy(annotation) for annotation in self.thriftType.structured_annotations)

    @property
    def functions(self):
        for function in self.thriftType.functions:
            yield ThriftFunctionProxy(function, self.thriftMeta)


def gen_metadata(obj_or_cls):
    if hasattr(obj_or_cls, "getThriftModuleMetadata"):
        return obj_or_cls.getThriftModuleMetadata()

    cls = obj_or_cls if isinstance(obj_or_cls, type) else type(obj_or_cls)

    if not issubclass(cls, (Struct, GeneratedError, ServiceInterface, Client, CompiledEnum)):
        raise TypeError(f'{cls!r} is not a thrift-py3 type.')

    # get the box
    cdef MetadataBox box = cls.__get_metadata__()
    # unbox the box
    cdef ThriftMetadata meta = ThriftMetadata._fbthrift_create(move(box._cpp_obj))
    cdef str name = cls.__get_thrift_name__()

    if issubclass(cls, Struct):
        return ThriftStructProxy(name, meta)
    elif issubclass(cls, GeneratedError):
        return ThriftExceptionProxy(name, meta)
    elif issubclass(cls, ServiceInterface):
        return ThriftServiceProxy(name, meta)
    elif issubclass(cls, Client):
        return ThriftServiceProxy(name, meta)
    elif issubclass(cls, CompiledEnum):
        return meta.enums[name]
    else:
        raise TypeError(f'unsupported thrift-py3 type: {cls!r}.')