from dataclasses import dataclass, field
from typing import Optional
from pynumaflow.proto.common import metadata_pb2

"""
 Metadata provides per-message metadata passed between vertices.

 A vertex could create one or more set of key-value pairs per group-name.
 This is required because a vertex could forward a message to a
 Kafka sink with Kafka headers, and a metrics Sink with some key/value points.

 There will be two kinds of metadata,

 system - generated by the system, exposed as read-only to UDFs
 user - user generated with read-write access
	|
	|          +-> [group-m] -> {k1:v1, ... }
	|          |
	+-> [user] +-> [group-n] -> {k1:v1, ... }
	|          |
	|          +-> [group-o] -> {k1:v1, ... }
	|
	|          +-> [group-h] -> {k1:v1, ... }
	|          |
	+-> [sys]  +-> [group-i] -> {k1:v1, ... }
	|
"""


@dataclass
class SystemMetadata:
    """
    System metadata is the mapping of group name to key-value pairs for a given group.
    System metadata wraps the system-generated metadata groups per message. It is read-only to UDFs.
    """

    _data: dict[str, dict[str, bytes]] = field(default_factory=dict)

    def groups(self) -> list[str]:
        """
        Returns the list of group names for the system metadata.
        """
        return list(self._data.keys())

    def keys(self, group: str) -> list[str]:
        """
        Returns the list of keys for a given group.
        """
        return list(self._data.get(group, {}).keys())

    def value(self, group: str, key: str) -> Optional[bytes]:
        """
        Returns the value for a given group and key.
        """
        return self._data.get(group, {}).get(key)


@dataclass
class UserMetadata:
    """
    UserMetadata wraps the user-generated metadata groups per message. It is read-write to UDFs.
    """

    _data: dict[str, dict[str, bytes]] = field(default_factory=dict)

    def groups(self) -> list[str]:
        """
        Returns the list of group names for the user metadata.
        """
        return list(self._data.keys())

    def keys(self, group: str) -> list[str]:
        """
        Returns the list of keys for a given group.
        """
        keys = self._data.get(group) or {}
        return list(keys.keys())

    def __contains__(self, group: str) -> bool:
        """
        Returns True if the group exists.
        """
        return group in self._data

    def __getitem__(self, group: str) -> dict[str, bytes]:
        """
        Returns the data for a given group.
        Raises KeyError if the group does not exist.
        """
        return self._data[group]

    def __setitem__(self, group: str, data: dict[str, bytes]):
        """
        Sets the data for a given group.
        """
        self._data[group] = data

    def __delitem__(self, group: str):
        """
        Removes the group and all its keys and values.
        Raises KeyError if the group does not exist.
        """
        del self._data[group]

    def __len__(self) -> int:
        """
        Returns the number of groups.
        """
        return len(self._data)

    def value(self, group: str, key: str) -> Optional[bytes]:
        """
        Returns the value for a given group and key.
        If the group or key does not exist, returns None.
        """
        value = self._data.get(group)
        if value is None:
            return None
        return value.get(key)

    def add_key(self, group: str, key: str, value: bytes):
        """
        Adds the value for a given group and key.
        """
        self._data.setdefault(group, {})[key] = value

    def remove_key(self, group: str, key: str) -> Optional[bytes]:
        """
        Removes the key and its value for a given group and returns the value.
        If this key is the only key in the group, the group will be removed.
        Returns None if the group or key does not exist.
        """
        group_data = self._data.pop(group, None)
        if group_data is None:
            return None
        value = group_data.pop(key, None)
        if group_data:
            self._data[group] = group_data
        return value

    def remove_group(self, group: str) -> Optional[dict[str, bytes]]:
        """
        Removes the group and all its keys and values and returns the data.
        Returns None if the group does not exist.
        """
        return self._data.pop(group, None)

    def clear(self):
        """
        Clears all the groups and all their keys and values.
        """
        self._data.clear()

    def _to_proto(self) -> metadata_pb2.Metadata:
        return metadata_pb2.Metadata(
            user_metadata={
                group: metadata_pb2.KeyValueGroup(key_value=value)
                for group, value in self._data.items()
            },
        )


def _user_and_system_metadata_from_proto(
    proto: metadata_pb2.Metadata,
) -> tuple[UserMetadata, SystemMetadata]:
    """
    Converts the protobuf metadata to the UserMetadata and SystemMetadata objects.
    """
    user_metadata = {group: dict(kv.key_value) for group, kv in proto.user_metadata.items()}
    system_metadata = {group: dict(kv.key_value) for group, kv in proto.sys_metadata.items()}
    return UserMetadata(user_metadata), SystemMetadata(system_metadata)
