"""Gather message handling for mt_dao."""

import uuid
from typing import cast

from flwr.common import (
    Code,
    ConfigRecord,
    MetricRecord,
    RecordDict,
)

from mt_dao.constants import (
    BATCH_ID,
    CODE,
    COMM_STACK,
    ENDPOINT_ID,
    FILE_NAME,
    FOLDER_NAME,
    GATHER_MSG,
    GATHER_P,
    GATHER_RES,
    MASK,
    MESSAGE,
    NUM_EXAMPLES,
    PARAMETERS,
    QUERY,
    STATUS,
    TYPE,
)
from mt_dao.messages.generic import build_success_recorddict


def build_gather_base_recorddict(
    gather_msg: str,
    aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]],
    batch_id: int = 0,
) -> RecordDict:
    """Build a RecordDict for gathering model updates and aggregation mask from clients.

    Parameters
    ----------
    gather_msg : str
        The message string used as a prefix for the gather operation in the recorddict.
    aggregation_mask : tuple[tuple[bool, ...], list[str], list[str]]
        A tuple containing:
        - A tuple of booleans indicating which parameters to aggregate.
        - A list of layer names.
        - A list of layer types.
    batch_id : int, optional
        The batch identifier for the gather operation (default is 0).

    Returns
    -------
    RecordDict
        A RecordDict containing:
        - A gather query entry.
        - A status ConfigRecord with code OK and gather message.
        - Communication stack instructions including endpoint ID, folder and file names,
          mask, and batch ID.

    Notes
    -----
    This function creates a RecordDict for gathering model updates and the associated
    aggregation mask from clients. The communication stack instructions are stored
    under the key "{gather_msg}.{COMM_STACK}" and include metadata required for the
    gather operation.

    """
    # Create an empty RecordDict object
    recorddict = build_success_recorddict(gather_msg)
    # Add the broadcast message to the recorddict
    recorddict.config_records[QUERY] = ConfigRecord({TYPE: GATHER_P})
    # Add Status to the recorddict
    recorddict.config_records[f"{gather_msg}.{STATUS}"] = ConfigRecord(
        {
            CODE: int(Code.OK.value),
            MESSAGE: GATHER_MSG,
        },
    )
    # Add communication stack instructions to the recorddict
    recorddict.config_records[f"{gather_msg}.{COMM_STACK}"] = ConfigRecord(
        {
            ENDPOINT_ID: str(uuid.uuid4()),
            FOLDER_NAME: COMM_STACK,
            FILE_NAME: PARAMETERS,
            MASK: repr(aggregation_mask),
            BATCH_ID: batch_id,
        },
    )
    return recorddict


def build_gather_res_recorddict(
    aggregation_mask: tuple[tuple[bool, ...], list[str], list[str]],
    num_examples: int,
    batch_id: int = 0,
) -> RecordDict:
    """Build a RecordDict for gathering results, with metrics and number of examples.

    Parameters
    ----------
    aggregation_mask : tuple[tuple[bool, ...], list[str], list[str]]
        A tuple containing:
        - A tuple of booleans indicating which parameters to aggregate.
        - A list of layer names.
        - A list of layer types.
    num_examples : int
        The number of examples used for aggregation.
    batch_id : int, optional
        The batch identifier for the gather operation (default is 0).

    Returns
    -------
    RecordDict
        A RecordDict containing:
        - The base gather configuration and communication stack.
        - A MetricRecord with the number of examples under the key
        "{GATHER_RES}.{NUM_EXAMPLES}".

    Notes
    -----
    This function creates a RecordDict for the gather response, including the
    aggregation mask and number of examples. The number of examples are stored
    under the keys "{GATHER_RES}.{NUM_EXAMPLES}".

    """
    # Create an empty RecordDict object
    recorddict = build_gather_base_recorddict(
        gather_msg=GATHER_RES,
        aggregation_mask=aggregation_mask,
        batch_id=batch_id,
    )
    # Add the number of examples
    recorddict.metric_records[f"{GATHER_RES}.{NUM_EXAMPLES}"] = MetricRecord(
        {NUM_EXAMPLES: num_examples},
    )
    return recorddict


def gather_res_recorddict_to_num_samples(
    gather_res_recorddict: RecordDict,
) -> int:
    """Extract the number of samples from a gather response RecordDict.

    Parameters
    ----------
    gather_res_recorddict : RecordDict
        The RecordDict containing the gather response, expected to have the number of
        examples under the key "{GATHER_RES}.{NUM_EXAMPLES}".

    Returns
    -------
    int
        The number of examples used for aggregation.

    """
    # Obtain the number of examples from the recorddict
    num_examples_dict = gather_res_recorddict.metric_records.get(
        f"{GATHER_RES}.{NUM_EXAMPLES}",
        None,
    )
    assert (
        num_examples_dict is not None
    ), "The number of examples is missing in the gather response recorddict."
    num_examples = num_examples_dict.get(NUM_EXAMPLES, None)
    assert (
        num_examples is not None
    ), "The number of examples is missing in the gather response recorddict."
    return cast("int", num_examples)
