"""
NBT Network Little Endian Serializer for Minecraft Bedrock Edition
"""

from typing import Optional
import struct

from bstream import BinaryStream

# Tag type IDs for NBT format
TAG_END = 0
TAG_BYTE = 1
TAG_SHORT = 2
TAG_INT = 3
TAG_LONG = 4
TAG_FLOAT = 5
TAG_DOUBLE = 6
TAG_BYTE_ARRAY = 7
TAG_STRING = 8
TAG_LIST = 9
TAG_COMPOUND = 10
TAG_INT_ARRAY = 11
TAG_LONG_ARRAY = 12


class NBTSerializer:
    """Serializes NBT tags to network little endian format for Bedrock Edition"""

    def __init__(self, stream):
        """
        Initialize the serializer with a BinaryStream

        Args:
            stream: BinaryStream instance (should be in little endian mode)
        """
        self.stream: BinaryStream = stream

    def write_tag(self, tag, name: Optional[str] = None, write_type: bool = True):
        """
        Write a complete NBT tag with optional name and type

        Args:
            tag: The NBT tag to serialize
            name: Optional tag name (for named tags)
            write_type: Whether to write the tag type ID
        """
        tag_type = self._get_tag_type(tag)

        if write_type:
            self.stream.write_byte(tag_type)

        if name is not None and tag_type != TAG_END:
            self._write_string(name)

        self._write_tag_payload(tag, tag_type)

    def _get_tag_type(self, tag) -> int:
        """Determine the NBT type ID for a tag"""
        type_name = type(tag).__name__

        type_map = {
            'ByteTag': TAG_BYTE,
            'ShortTag': TAG_SHORT,
            'IntTag': TAG_INT,
            'LongTag': TAG_LONG,
            'FloatTag': TAG_FLOAT,
            'DoubleTag': TAG_DOUBLE,
            'ByteArrayTag': TAG_BYTE_ARRAY,
            'StringTag': TAG_STRING,
            'ListTag': TAG_LIST,
            'CompoundTag': TAG_COMPOUND,
            'IntArrayTag': TAG_INT_ARRAY,
        }

        return type_map.get(type_name, TAG_END)

    def _write_string(self, value: str):
        """Write a string with length prefix (unsigned short)"""
        encoded = value.encode('utf-8')
        self.stream.write_unsigned_short(len(encoded))
        self.stream.write_raw_bytes(encoded)

    def _write_tag_payload(self, tag, tag_type: int):
        """Write the payload data for a tag based on its type"""
        if tag_type == TAG_END:
            return
        elif tag_type == TAG_BYTE:
            self.stream.write_byte(tag.value)
        elif tag_type == TAG_SHORT:
            self.stream.write_signed_short(tag.value)
        elif tag_type == TAG_INT:
            self.stream.write_signed_int(tag.value)
        elif tag_type == TAG_LONG:
            self.stream.write_signed_int64(tag.value)
        elif tag_type == TAG_FLOAT:
            self.stream.write_float(tag.value)
        elif tag_type == TAG_DOUBLE:
            self.stream.write_double(tag.value)
        elif tag_type == TAG_BYTE_ARRAY:
            self._write_byte_array(tag)
        elif tag_type == TAG_STRING:
            self._write_string(tag.value)
        elif tag_type == TAG_LIST:
            self._write_list(tag)
        elif tag_type == TAG_COMPOUND:
            self._write_compound(tag)
        elif tag_type == TAG_INT_ARRAY:
            self._write_int_array(tag)

    def _write_byte_array(self, tag):
        """Write a byte array tag"""
        self.stream.write_signed_int(len(tag))
        for byte_val in tag:
            self.stream.write_byte(byte_val)

    def _write_int_array(self, tag):
        """Write an int array tag"""
        self.stream.write_signed_int(len(tag))
        for int_val in tag:
            self.stream.write_signed_int(int_val)

    def _write_list(self, tag):
        """Write a list tag"""
        if len(tag) == 0:
            # Empty list - use TAG_END as element type
            self.stream.write_byte(TAG_END)
            self.stream.write_signed_int(0)
        else:
            # Get type from first element
            element_type = self._get_tag_type(tag[0])
            self.stream.write_byte(element_type)
            self.stream.write_signed_int(len(tag))

            # Write all elements without type/name prefixes
            for element in tag:
                self._write_tag_payload(element, element_type)

    def _write_compound(self, tag):
        """Write a compound tag"""
        for key in tag.keys():
            value = tag[key]
            self.write_tag(value, name=key, write_type=True)

        # Write TAG_END to mark end of compound
        self.stream.write_byte(TAG_END)

    def serialize_root(self, tag, root_name: str = ""):
        """
        Serialize a root NBT tag (typically a CompoundTag)

        Args:
            tag: The root NBT tag
            root_name: Name for the root tag (default: empty string)
        """
        self.write_tag(tag, name=root_name, write_type=True)


# Convenience function
def write_nbt(stream, tag, root_name: str = "") -> bytes:
    """
    Write an NBT tag to a BinaryStream

    Args:
        stream: BinaryStream instance
        tag: NBT tag to serialize
        root_name: Optional root tag name

    Returns:
        The serialized data from the stream
    """
    serializer = NBTSerializer(stream)
    serializer.serialize_root(tag, root_name)
    return stream.data()

#def serialize(tag: CompoundTag)