# 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.

"""
Pure Python collection serializers for debugging and Python-only execution.

In Cython mode the active collection serializers live in `collection.pxi` and
are imported through `pyfory.serialization`. This module is the pure-Python
fallback only.
"""

from pyfory.serialization import ENABLE_FORY_CYTHON_SERIALIZATION
from pyfory._serializer import Serializer, StringSerializer
from pyfory.resolver import NOT_NULL_VALUE_FLAG, NULL_FLAG
from pyfory.types import TypeId

COLL_DEFAULT_FLAG = 0b0
COLL_TRACKING_REF = 0b1
COLL_HAS_NULL = 0b10
COLL_IS_DECL_ELEMENT_TYPE = 0b100
COLL_IS_SAME_TYPE = 0b1000


def _needs_element_type_info(type_id):
    return type_id in {
        TypeId.STRUCT,
        TypeId.COMPATIBLE_STRUCT,
        TypeId.NAMED_STRUCT,
        TypeId.NAMED_COMPATIBLE_STRUCT,
        TypeId.EXT,
        TypeId.NAMED_EXT,
    }


class CollectionSerializer(Serializer):
    __slots__ = (
        "elem_serializer",
        "elem_tracking_ref",
        "elem_type",
        "elem_type_info",
    )

    def __init__(self, type_resolver, type_, elem_serializer=None, elem_tracking_ref=None):
        super().__init__(type_resolver, type_)
        self.elem_serializer = elem_serializer
        if elem_tracking_ref is not None:
            self.elem_tracking_ref = 1 if elem_tracking_ref else 0
        else:
            self.elem_tracking_ref = -1
        if elem_serializer is None:
            self.elem_type = None
            self.elem_type_info = self.type_resolver.get_type_info(None)
        else:
            self.elem_type = elem_serializer.type_
            self.elem_type_info = self.type_resolver.get_type_info(self.elem_type)
            if elem_tracking_ref is None:
                self.elem_tracking_ref = int(elem_serializer.need_to_write_ref)

    def write_header(self, write_context, value):
        collect_flag = COLL_DEFAULT_FLAG
        elem_type = self.elem_type
        elem_type_info = self.elem_type_info
        has_null = False
        has_same_type = True
        if elem_type is None:
            for item in value:
                if item is None:
                    has_null = True
                    continue
                if elem_type is None:
                    elem_type = type(item)
                elif has_same_type and type(item) is not elem_type:
                    has_same_type = False
            if has_same_type:
                collect_flag |= COLL_IS_SAME_TYPE
                if elem_type is not None:
                    elem_type_info = self.type_resolver.get_type_info(elem_type)
        else:
            collect_flag |= COLL_IS_SAME_TYPE
            if not _needs_element_type_info(elem_type_info.type_id):
                collect_flag |= COLL_IS_DECL_ELEMENT_TYPE
            for item in value:
                if item is None:
                    has_null = True
                    break

        if has_null:
            collect_flag |= COLL_HAS_NULL
        if write_context.track_ref:
            if self.elem_tracking_ref == 1:
                collect_flag |= COLL_TRACKING_REF
            elif self.elem_tracking_ref == -1:
                if not has_same_type or elem_type_info.serializer.need_to_write_ref:
                    collect_flag |= COLL_TRACKING_REF
        write_context.write_var_uint32(len(value))
        write_context.write_int8(collect_flag)
        if has_same_type and (collect_flag & COLL_IS_DECL_ELEMENT_TYPE) == 0:
            self.type_resolver.write_type_info(write_context, elem_type_info)
        return collect_flag, elem_type_info

    def write(self, write_context, value):
        if len(value) == 0:
            write_context.write_var_uint32(0)
            return
        collect_flag, typeinfo = self.write_header(write_context, value)
        serializer = (
            self.elem_serializer if (collect_flag & COLL_IS_DECL_ELEMENT_TYPE) != 0 and self.elem_serializer is not None else typeinfo.serializer
        )
        if (collect_flag & COLL_IS_SAME_TYPE) != 0:
            if (collect_flag & COLL_TRACKING_REF) != 0:
                self._write_same_type_ref(write_context, value, serializer)
            elif (collect_flag & COLL_HAS_NULL) == 0:
                self._write_same_type_no_ref(write_context, value, serializer)
            else:
                self._write_same_type_has_null(write_context, value, serializer)
        else:
            self._write_different_types(write_context, value, collect_flag)

    def _write_same_type_no_ref(self, write_context, value, serializer):
        for item in value:
            serializer.write(write_context, item)

    def _write_same_type_has_null(self, write_context, value, serializer):
        for item in value:
            if item is None:
                write_context.write_int8(NULL_FLAG)
            else:
                write_context.write_int8(NOT_NULL_VALUE_FLAG)
                serializer.write(write_context, item)

    def _write_same_type_ref(self, write_context, value, serializer):
        ref_writer = write_context.ref_writer
        for item in value:
            if not ref_writer.write_ref_or_null(write_context, item):
                serializer.write(write_context, item)

    def _write_different_types(self, write_context, value, collect_flag=0):
        tracking_ref = (collect_flag & COLL_TRACKING_REF) != 0
        has_null = (collect_flag & COLL_HAS_NULL) != 0
        ref_writer = write_context.ref_writer
        if tracking_ref:
            for item in value:
                if not ref_writer.write_ref_or_null(write_context, item):
                    typeinfo = self.type_resolver.get_type_info(type(item))
                    self.type_resolver.write_type_info(write_context, typeinfo)
                    typeinfo.serializer.write(write_context, item)
            return
        if not has_null:
            for item in value:
                typeinfo = self.type_resolver.get_type_info(type(item))
                self.type_resolver.write_type_info(write_context, typeinfo)
                typeinfo.serializer.write(write_context, item)
            return
        for item in value:
            if item is None:
                write_context.write_int8(NULL_FLAG)
            else:
                write_context.write_int8(NOT_NULL_VALUE_FLAG)
                typeinfo = self.type_resolver.get_type_info(type(item))
                self.type_resolver.write_type_info(write_context, typeinfo)
                typeinfo.serializer.write(write_context, item)

    def read(self, read_context):
        length = read_context.read_var_uint32()
        collection_ = self.new_instance(read_context, self.type_)
        if length == 0:
            return collection_
        collect_flag = read_context.read_int8()
        # IMPORTANT: collection readers must obey the ref/null bits written on
        # the wire, not the local Python element annotation or runtime type
        # that may imply a different ref policy. Shared xlang tests
        # intentionally deserialize one ref policy and then serialize another
        # local payload. DO NOT REMOVE this comment.
        if (collect_flag & COLL_IS_SAME_TYPE) != 0:
            if (collect_flag & COLL_IS_DECL_ELEMENT_TYPE) == 0:
                typeinfo = self.type_resolver.read_type_info(read_context)
                serializer = typeinfo.serializer
            else:
                serializer = self.elem_serializer
            if (collect_flag & COLL_TRACKING_REF) != 0:
                self._read_same_type_ref(read_context, length, collection_, serializer)
            elif (collect_flag & COLL_HAS_NULL) == 0:
                self._read_same_type_no_ref(read_context, length, collection_, serializer)
            else:
                self._read_same_type_has_null(read_context, length, collection_, serializer)
        else:
            self._read_different_types(read_context, length, collection_, collect_flag)
        return collection_

    def new_instance(self, read_context, type_):
        raise NotImplementedError

    def _add_element(self, collection_, element):
        raise NotImplementedError

    def _read_same_type_no_ref(self, read_context, length, collection_, serializer):
        read_context.increase_depth()
        for _ in range(length):
            self._add_element(collection_, read_context.read_no_ref(serializer=serializer))
        read_context.decrease_depth()

    def _read_same_type_has_null(self, read_context, length, collection_, serializer):
        read_context.increase_depth()
        for _ in range(length):
            if read_context.read_int8() == NULL_FLAG:
                self._add_element(collection_, None)
            else:
                self._add_element(collection_, read_context.read_no_ref(serializer=serializer))
        read_context.decrease_depth()

    def _read_same_type_ref(self, read_context, length, collection_, serializer):
        read_context.increase_depth()
        ref_reader = read_context.ref_reader
        for _ in range(length):
            ref_id = ref_reader.try_preserve_ref_id(read_context)
            if ref_id < NOT_NULL_VALUE_FLAG:
                obj = ref_reader.get_read_ref()
            else:
                obj = serializer.read(read_context)
                ref_reader.set_read_ref(ref_id, obj)
            self._add_element(collection_, obj)
        read_context.decrease_depth()

    def _read_different_types(self, read_context, length, collection_, collect_flag):
        read_context.increase_depth()
        tracking_ref = (collect_flag & COLL_TRACKING_REF) != 0
        has_null = (collect_flag & COLL_HAS_NULL) != 0
        if tracking_ref:
            for _ in range(length):
                self._add_element(collection_, get_next_element(read_context))
            read_context.decrease_depth()
            return
        if not has_null:
            for _ in range(length):
                typeinfo = self.type_resolver.read_type_info(read_context)
                elem = None if typeinfo is None else read_context.read_no_ref(serializer=typeinfo.serializer)
                self._add_element(collection_, elem)
            read_context.decrease_depth()
            return
        for _ in range(length):
            head_flag = read_context.read_int8()
            if head_flag == NULL_FLAG:
                elem = None
            else:
                typeinfo = self.type_resolver.read_type_info(read_context)
                elem = None if typeinfo is None else read_context.read_no_ref(serializer=typeinfo.serializer)
            self._add_element(collection_, elem)
        read_context.decrease_depth()


class ListSerializer(CollectionSerializer):
    def new_instance(self, read_context, type_):
        instance = []
        read_context.reference(instance)
        return instance

    def _add_element(self, collection_, element):
        collection_.append(element)


class TupleSerializer(CollectionSerializer):
    def new_instance(self, read_context, type_):
        return []

    def _add_element(self, collection_, element):
        collection_.append(element)

    def read(self, read_context):
        return tuple(super().read(read_context))


class StringArraySerializer(ListSerializer):
    def __init__(self, type_resolver, type_):
        super().__init__(type_resolver, type_, StringSerializer(type_resolver, str))


class SetSerializer(CollectionSerializer):
    def new_instance(self, read_context, type_):
        instance = set()
        read_context.reference(instance)
        return instance

    def _add_element(self, collection_, element):
        collection_.add(element)


def get_next_element(read_context):
    ref_reader = read_context.ref_reader
    ref_id = ref_reader.try_preserve_ref_id(read_context)
    if ref_id < NOT_NULL_VALUE_FLAG:
        return ref_reader.get_read_ref()
    typeinfo = read_context.type_resolver.read_type_info(read_context)
    obj = typeinfo.serializer.read(read_context)
    ref_reader.set_read_ref(ref_id, obj)
    return obj


MAX_CHUNK_SIZE = 255
TRACKING_KEY_REF = 0b1
KEY_HAS_NULL = 0b10
KEY_DECL_TYPE = 0b100
TRACKING_VALUE_REF = 0b1000
VALUE_HAS_NULL = 0b10000
VALUE_DECL_TYPE = 0b100000
KV_NULL = KEY_HAS_NULL | VALUE_HAS_NULL
NULL_KEY_VALUE_DECL_TYPE = KEY_HAS_NULL | VALUE_DECL_TYPE
NULL_KEY_VALUE_DECL_TYPE_TRACKING_REF = KEY_HAS_NULL | VALUE_DECL_TYPE | TRACKING_VALUE_REF
NULL_VALUE_KEY_DECL_TYPE = VALUE_HAS_NULL | KEY_DECL_TYPE
NULL_VALUE_KEY_DECL_TYPE_TRACKING_REF = VALUE_HAS_NULL | KEY_DECL_TYPE | TRACKING_KEY_REF


class MapSerializer(Serializer):
    def __init__(
        self,
        type_resolver,
        type_,
        key_serializer=None,
        value_serializer=None,
        key_tracking_ref=None,
        value_tracking_ref=None,
    ):
        super().__init__(type_resolver, type_)
        self.key_serializer = key_serializer
        self.value_serializer = value_serializer
        self.key_tracking_ref = False
        self.value_tracking_ref = False
        if key_serializer is not None:
            self.key_tracking_ref = bool(key_serializer.need_to_write_ref)
            if key_tracking_ref is not None:
                self.key_tracking_ref = bool(key_tracking_ref) and type_resolver.track_ref
        if value_serializer is not None:
            self.value_tracking_ref = bool(value_serializer.need_to_write_ref)
            if value_tracking_ref is not None:
                self.value_tracking_ref = bool(value_tracking_ref) and type_resolver.track_ref

    def write(self, write_context, obj):
        length = len(obj)
        write_context.write_var_uint32(length)
        if length == 0:
            return
        type_resolver = self.type_resolver
        ref_writer = write_context.ref_writer
        key_serializer = self.key_serializer
        value_serializer = self.value_serializer

        items_iter = iter(obj.items())
        key, value = next(items_iter)
        has_next = True
        while has_next:
            while True:
                if key is not None:
                    if value is not None:
                        break
                    if key_serializer is not None:
                        key_write_ref = self.key_tracking_ref
                        if key_write_ref:
                            write_context.write_int8(NULL_VALUE_KEY_DECL_TYPE_TRACKING_REF)
                            if not ref_writer.write_ref_or_null(write_context, key):
                                self._write_obj(key_serializer, write_context, key)
                        else:
                            write_context.write_int8(NULL_VALUE_KEY_DECL_TYPE)
                            self._write_obj(key_serializer, write_context, key)
                    else:
                        write_context.write_int8(VALUE_HAS_NULL | TRACKING_KEY_REF)
                        write_context.write_ref(key)
                else:
                    if value is not None:
                        if value_serializer is not None:
                            value_write_ref = self.value_tracking_ref
                            if value_write_ref:
                                write_context.write_int8(NULL_KEY_VALUE_DECL_TYPE_TRACKING_REF)
                                if not ref_writer.write_ref_or_null(write_context, value):
                                    value_serializer.write(write_context, value)
                            else:
                                write_context.write_int8(NULL_KEY_VALUE_DECL_TYPE)
                                value_serializer.write(write_context, value)
                        else:
                            write_context.write_int8(KEY_HAS_NULL | TRACKING_VALUE_REF)
                            write_context.write_ref(value)
                    else:
                        write_context.write_int8(KV_NULL)
                try:
                    key, value = next(items_iter)
                except StopIteration:
                    has_next = False
                    break

            if not has_next:
                break

            key_cls = type(key)
            value_cls = type(value)
            write_context.enter_flush_barrier()
            write_context.write_int16(-1)
            chunk_size_offset = write_context.get_writer_index() - 1
            chunk_header = 0

            if key_serializer is not None:
                chunk_header |= KEY_DECL_TYPE
            else:
                key_type_info = type_resolver.get_type_info(key_cls)
                type_resolver.write_type_info(write_context, key_type_info)
                key_serializer = key_type_info.serializer

            if value_serializer is not None:
                chunk_header |= VALUE_DECL_TYPE
            else:
                value_type_info = type_resolver.get_type_info(value_cls)
                type_resolver.write_type_info(write_context, value_type_info)
                value_serializer = value_type_info.serializer

            key_write_ref = self.key_tracking_ref if self.key_serializer is not None else bool(key_serializer.need_to_write_ref)
            value_write_ref = self.value_tracking_ref if self.value_serializer is not None else bool(value_serializer.need_to_write_ref)
            if key_write_ref:
                chunk_header |= TRACKING_KEY_REF
            if value_write_ref:
                chunk_header |= TRACKING_VALUE_REF

            write_context.put_uint8(chunk_size_offset - 1, chunk_header)
            chunk_size = 0

            while chunk_size < MAX_CHUNK_SIZE:
                if key is None or value is None or type(key) is not key_cls or type(value) is not value_cls:
                    break
                if not key_write_ref or not ref_writer.write_ref_or_null(write_context, key):
                    self._write_obj(key_serializer, write_context, key)
                if not value_write_ref or not ref_writer.write_ref_or_null(write_context, value):
                    self._write_obj(value_serializer, write_context, value)
                chunk_size += 1
                try:
                    key, value = next(items_iter)
                except StopIteration:
                    has_next = False
                    break

            key_serializer = self.key_serializer
            value_serializer = self.value_serializer
            write_context.put_uint8(chunk_size_offset, chunk_size)
            write_context.exit_flush_barrier()
            write_context.try_flush()

    def read(self, read_context):
        size = read_context.read_var_uint32()
        map_ = {}
        ref_reader = read_context.ref_reader
        read_context.reference(map_)
        chunk_header = read_context.read_uint8() if size != 0 else 0
        key_serializer = self.key_serializer
        value_serializer = self.value_serializer
        read_context.increase_depth()
        while size > 0:
            while True:
                key_has_null = (chunk_header & KEY_HAS_NULL) != 0
                value_has_null = (chunk_header & VALUE_HAS_NULL) != 0
                if not key_has_null and not value_has_null:
                    break
                if not key_has_null:
                    track_key_ref = (chunk_header & TRACKING_KEY_REF) != 0
                    if (chunk_header & KEY_DECL_TYPE) != 0:
                        if track_key_ref:
                            ref_id = ref_reader.try_preserve_ref_id(read_context)
                            if ref_id < NOT_NULL_VALUE_FLAG:
                                key = ref_reader.get_read_ref()
                            else:
                                key = self._read_obj(key_serializer, read_context)
                                ref_reader.set_read_ref(ref_id, key)
                        else:
                            key = self._read_obj_no_ref(key_serializer, read_context)
                    else:
                        key = read_context.read_ref()
                    map_[key] = None
                elif not value_has_null:
                    track_value_ref = (chunk_header & TRACKING_VALUE_REF) != 0
                    if (chunk_header & VALUE_DECL_TYPE) != 0:
                        if track_value_ref:
                            ref_id = ref_reader.try_preserve_ref_id(read_context)
                            if ref_id < NOT_NULL_VALUE_FLAG:
                                value = ref_reader.get_read_ref()
                            else:
                                value = self._read_obj(value_serializer, read_context)
                                ref_reader.set_read_ref(ref_id, value)
                        else:
                            value = self._read_obj_no_ref(value_serializer, read_context)
                    else:
                        value = read_context.read_ref()
                    map_[None] = value
                else:
                    map_[None] = None
                size -= 1
                if size == 0:
                    read_context.decrease_depth()
                    return map_
                chunk_header = read_context.read_uint8()

            # IMPORTANT: map readers must obey the sender-written key/value ref
            # bits in the wire header. Local Python serializer choices must not
            # override that decision while reading. Shared xlang tests
            # intentionally deserialize one ref policy and then serialize
            # another local payload. DO NOT REMOVE this comment.
            track_key_ref = (chunk_header & TRACKING_KEY_REF) != 0
            track_value_ref = (chunk_header & TRACKING_VALUE_REF) != 0
            key_is_declared_type = (chunk_header & KEY_DECL_TYPE) != 0
            value_is_declared_type = (chunk_header & VALUE_DECL_TYPE) != 0
            chunk_size = read_context.read_uint8()
            if not key_is_declared_type:
                key_serializer = self.type_resolver.read_type_info(read_context).serializer
            if not value_is_declared_type:
                value_serializer = self.type_resolver.read_type_info(read_context).serializer
            for _ in range(chunk_size):
                if track_key_ref:
                    ref_id = ref_reader.try_preserve_ref_id(read_context)
                    if ref_id < NOT_NULL_VALUE_FLAG:
                        key = ref_reader.get_read_ref()
                    else:
                        key = self._read_obj(key_serializer, read_context)
                        ref_reader.set_read_ref(ref_id, key)
                else:
                    key = self._read_obj_no_ref(key_serializer, read_context)
                if track_value_ref:
                    ref_id = ref_reader.try_preserve_ref_id(read_context)
                    if ref_id < NOT_NULL_VALUE_FLAG:
                        value = ref_reader.get_read_ref()
                    else:
                        value = self._read_obj(value_serializer, read_context)
                        ref_reader.set_read_ref(ref_id, value)
                else:
                    value = self._read_obj_no_ref(value_serializer, read_context)
                map_[key] = value
                size -= 1
            if size != 0:
                chunk_header = read_context.read_uint8()
        read_context.decrease_depth()
        return map_

    def _write_obj(self, serializer, write_context, obj):
        serializer.write(write_context, obj)

    def _read_obj(self, serializer, read_context):
        return serializer.read(read_context)

    def _read_obj_no_ref(self, serializer, read_context):
        return read_context.read_no_ref(serializer=serializer)


SubMapSerializer = MapSerializer


if ENABLE_FORY_CYTHON_SERIALIZATION:
    from pyfory.serialization import (
        CollectionSerializer as CythonCollectionSerializer,
        ListSerializer as CythonListSerializer,
        TupleSerializer as CythonTupleSerializer,
        StringArraySerializer as CythonStringArraySerializer,
        SetSerializer as CythonSetSerializer,
        MapSerializer as CythonMapSerializer,
    )

    CollectionSerializer = CythonCollectionSerializer
    ListSerializer = CythonListSerializer
    TupleSerializer = CythonTupleSerializer
    StringArraySerializer = CythonStringArraySerializer
    SetSerializer = CythonSetSerializer
    MapSerializer = CythonMapSerializer
    SubMapSerializer = CythonMapSerializer