"""Generic message building functions for mt_dao."""

from typing import cast

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

from mt_dao.constants import (
    CLIENT_ID,
    CODE,
    FAILURE_MSG,
    LOSS,
    MESSAGE,
    METRICS,
    NUM_EXAMPLES,
    STATUS,
    SUCCESS_MSG,
)


def build_success_with_metrics_recorddict(
    msg_str: str,
    metrics: dict[str, Scalar],
    num_examples: int,
) -> RecordDict:
    """Build a RecordDict indicating successful operation with metrics.

    Parameters
    ----------
    msg_str : str
        The message string used as a prefix for the status key in the recorddict.
    metrics : dict[str, Scalar]
        A dictionary containing metric names and their corresponding Scalar values.
    num_examples : int
        The number of examples used for the operation.

    Returns
    -------
    RecordDict
        A RecordDict containing:
        - A ConfigRecord with metrics, num_examples, and status under the key
            "{msg_str}"
        - A status ConfigRecord with code OK and success message under the key
            "{msg_str}.status"

    """
    # Create an empty RecordDict object
    recorddict = RecordDict()
    # Add Status to the recorddict
    recorddict.config_records[f"{msg_str}.{STATUS}"] = ConfigRecord(
        {
            CODE: int(Code.OK.value),
            MESSAGE: SUCCESS_MSG,
        },
    )
    # Add the metrics
    recorddict.config_records[f"{msg_str}.{METRICS}"] = ConfigRecord(
        cast("dict[str, ConfigRecordValues]", metrics),
    )
    # Add the number of examples
    recorddict.metric_records[f"{msg_str}.{NUM_EXAMPLES}"] = MetricRecord(
        {NUM_EXAMPLES: num_examples},
    )
    return recorddict


def get_metrics_from_recorddict(
    recorddict: RecordDict,
    msg_str: str,
) -> tuple[int, dict[str, Scalar]]:
    """Extract metrics from a RecordDict.

    Parameters
    ----------
    recorddict : RecordDict
        The RecordDict containing the metrics.
    msg_str : str
        The message string used as a prefix for the metrics key in the recorddict.

    Returns
    -------
    tuple[int, dict[str, Scalar]]
        A dictionary containing the extracted metrics and the number of examples.

    """
    # Get the metrics config record
    metrics_config_record = recorddict.config_records.get(f"{msg_str}.{METRICS}", None)
    assert (
        metrics_config_record is not None
    ), f"Expected a config record for {msg_str}.{METRICS}, but got None."
    # Get the number of examples
    num_examples_metric_record = recorddict.metric_records.get(
        f"{msg_str}.{NUM_EXAMPLES}", None,
    )
    assert (
        num_examples_metric_record is not None
    ), f"Expected a metric record for {msg_str}.{NUM_EXAMPLES}, but got None."
    num_examples = num_examples_metric_record.get(NUM_EXAMPLES, None)
    assert num_examples is not None, (
        f"Expected a metric for {NUM_EXAMPLES}, but got None.",
    )
    assert not isinstance(num_examples, list), (
        f"Expected a single value for {NUM_EXAMPLES}, but got a list: {num_examples}.",
    )
    return int(num_examples), cast("dict[str, Scalar]", metrics_config_record)


def build_success_recorddict(
    msg_str: str,
) -> RecordDict:
    """Build a RecordDict indicating successful operation.

    Parameters
    ----------
    msg_str : str
        The message string used as a prefix for the status key in the recorddict.

    Returns
    -------
    RecordDict
        A RecordDict containing a status ConfigRecord with code OK and success message.

    Notes
    -----
    This function creates a RecordDict with a status entry indicating a successful
    operation. The status is stored under the key "{msg_str}.status" and contains
    the code and message for success.

    """
    # Create an empty RecordDict object
    recorddict = RecordDict()
    # Add Status to the recorddict
    recorddict.config_records[f"{msg_str}.{STATUS}"] = ConfigRecord(
        {
            CODE: int(Code.OK.value),
            MESSAGE: SUCCESS_MSG,
        },
    )
    return recorddict


def build_failure_recorddict(
    msg_str: str,
) -> RecordDict:
    """Build a RecordDict indicating failed operation.

    Parameters
    ----------
    msg_str : str
        The message string used as a prefix for the status key in the recorddict.

    Returns
    -------
    RecordDict
        A RecordDict containing a status ConfigRecord with code -1 and failure message.

    Notes
    -----
    This function creates a RecordDict with a status entry indicating a failed
    operation. The status is stored under the key "{msg_str}.status" and contains
    the code and message for failure.

    """
    # Create an empty RecordDict object
    recorddict = RecordDict()
    # Add Status to the recorddict
    recorddict.config_records[f"{msg_str}.{STATUS}"] = ConfigRecord(
        {
            CODE: -1,
            MESSAGE: FAILURE_MSG,
        },
    )
    return recorddict


def build_metrics_recorddict(
    msg_str: str,
    metrics: dict[str, float],
    loss: float,
    num_examples: int,
) -> RecordDict:
    """Build a RecordDict containing evaluation metrics, loss, and status.

    Parameters
    ----------
    msg_str : str
        The message string used as a prefix for the metrics and status keys in the
        recorddict.
    metrics : dict[str, float]
        A dictionary containing metric names and their corresponding float values.
    loss : float
        The loss value to include in the recorddict.
    num_examples : int
        The number of examples used for evaluation.

    Returns
    -------
    RecordDict
        A RecordDict containing:
        - A ConfigRecord with metrics, loss, and num_examples under the key "{msg_str}"
        - A status ConfigRecord with code OK and success message under the key
        "{msg_str}.status"

    Notes
    -----
    This function creates a RecordDict with both the evaluation metrics and a status
    entry indicating a successful operation. The metrics and loss are stored under the
    key "{msg_str}", and the status is stored under the key "{msg_str}.status".

    """
    return RecordDict(
        records={
            msg_str: ConfigRecord(
                {
                    **metrics,
                    LOSS: loss,
                    NUM_EXAMPLES: num_examples,
                },
            ),
            f"{msg_str}.{STATUS}": ConfigRecord(
                {
                    CODE: int(Code.OK.value),
                    MESSAGE: SUCCESS_MSG,
                },
            ),
        },
    )


def get_client_id_from_train_eval_ins_recorddict(
    recorddict: RecordDict,
    msg_str: str,
) -> int:
    """Extract the client ID from a train or eval instructions RecordDict.

    Parameters
    ----------
    recorddict : RecordDict
        The RecordDict containing the train or eval instructions.
    msg_str : str
        The message string used as a prefix for the client IDs key in the recorddict.

    Returns
    -------
    int
        The client ID extracted from the RecordDict.

    """
    # Extract the fit or eval instructions config record
    fit_eval_ins_configrecord = recorddict.config_records.get(msg_str, None)
    assert (
        fit_eval_ins_configrecord is not None
    ), f"Expected a config record for {msg_str}, but got None."
    # Extract the list of client IDs
    client_id = fit_eval_ins_configrecord.get(CLIENT_ID, None)
    assert (
        client_id is not None
    ), f"Expected a client ID for {CLIENT_ID}, but got None."
    assert not isinstance(client_id, list), (
        f"Expected a single client ID for {CLIENT_ID}, but got a list: {client_id}."
    )
    # Return the first (and only) client ID
    return int(client_id)
