"""Utilities for mt_dao RecordDict handling."""

from __future__ import annotations

import json
from typing import Any, TypeVar, cast

from flwr.common import (
    ArrayRecord,
    RecordDict,
)
from flwr.common.inflatable import (
    InflatableObject,
    add_header_to_object_body,
)

T = TypeVar("T")


class mt_daoRecordDict(RecordDict):
    """A RecordDict that can hold mt_daoArrayRecords."""

    mt_dao_array_records: dict[str, Any]

    def deflate(self) -> bytes:
        """Deflate the RecordDict to bytes.

        This method serializes the RecordDict into a format that can be
        transmitted or stored. It converts the mt_daoArrayRecords into
        ArrayRecords and prepares a dictionary of references to the records.

        Returns
        -------
        bytes
            The deflated object content of the RecordDict.

        """
        # NOTE: These record_refs are supposed to describe with key-values pairs
        # record_name: record_object_id mapping
        record_refs: dict[str, str] = {}
        for record_name, record in self.mt_dao_array_records.items():
            self.array_records[record_name] = ArrayRecord(record)

        for record_name, record in self.items():
            if record_name in self.mt_dao_array_records:
                # Skip already processed records
                continue
            record_refs[record_name] = record.object_id

        # Serialize references dict
        object_body = json.dumps(record_refs).encode("utf-8")
        return add_header_to_object_body(object_body=object_body, obj=self)

    @classmethod
    def inflate(
        cls,
        object_content: bytes,
        children: dict[str, InflatableObject] | None = None,
    ) -> RecordDict:
        """Inflate an RecordDict from bytes.

        TODO(anon): Sayw hwat do.

        Parameters
        ----------
        object_content : bytes
            The deflated object content of the RecordDict.
        children : Optional[dict[str, InflatableObject]]  (default: None)
            Dictionary of children InflatableObjects mapped to their Object IDs.
            These children enable the full inflation of the RecordDict. Default is None.

        Returns
        -------
        RecordDict
            The inflated RecordDict.

        """
        # Instantiate new RecordDict
        state = cast(
            "mt_daoRecordDict",
            super().inflate(
                object_content=object_content,
                children=children,
            ),
        )
        # Inflate mt_daoArrayRecord
        for record_name, record in state.array_records.items():
            state.mt_dao_array_records[record_name] = record.to_numpy_ndarrays()
            del state.array_records[record_name]

        return state
