"""Estimates the message sending and receiving times during training rounds."""

import time
from collections import defaultdict
from typing import cast

from flwr.common import Scalar
from flwr.common.message import Message

from mt_dao.callbacks.constants import (
    DAYS,
    HOURS,
    LIST_OF_SUPPORTED_TIME_UNITS,
    MINUTES,
    SECONDS,
)


class FitEvalCallback:
    """Callback to record message sending and receiving times during training rounds.

    This callback tracks the time taken to send and receive messages of different types.

    """

    def __init__(self, time_unit: str = "hours") -> None:
        """Initialize the FitEvalCallback.

        Parameters
        ----------
        time_unit : str, optional
            The unit of time to use for the estimates. Can be one of 'seconds',
            'minutes', 'hours', or 'days'. Defaults to 'hours'.

        Raises
        ------
        ValueError
            If the `time_unit` is not one of 'seconds', 'minutes', 'hours', or 'days'.

        """
        # Message sending and receiving times
        self.record_send_msg_dict: dict[str, dict[int, float]] = defaultdict(
            lambda: defaultdict(float),
        )
        self.record_elapsed_times_dict: dict[str, dict[int, list[float]]] = defaultdict(
            lambda: defaultdict(list),
        )

        # Aggregation task submission times
        self.record_submit_aggregation: float = 0.0
        self.record_submit_aggregation_times: list[float] = []

        # Waiting for aggregation times
        self.record_waiting_for_aggregation: float = 0.0
        self.record_waiting_for_aggregation_times: float = 0.0

        # Server optimization times
        self.record_server_optimization: float = 0.0
        self.record_server_optimization_times: float = 0.0

        self.time_unit = time_unit
        self.divider = 1
        if time_unit == SECONDS:
            self.divider = 1
        elif time_unit == MINUTES:
            self.divider = 60
        elif time_unit == HOURS:
            self.divider = 60 * 60
        elif time_unit == DAYS:
            self.divider = 60 * 60 * 24
        else:
            msg = f"Invalid time_unit: {time_unit}, use {LIST_OF_SUPPORTED_TIME_UNITS}."
            raise ValueError(
                msg,
            )

    def record_send_message(self, message: Message) -> None:
        """Record the time when a message is sent.

        Parameters
        ----------
        message : Message
            The message that is being sent.

        """
        message_type = message.metadata.message_type
        destination_node_id = message.metadata.dst_node_id
        self.record_send_msg_dict[message_type][
            destination_node_id
        ] = time.perf_counter()

    def record_recv_message(self, message: Message) -> None:
        """Record the time when a message is received and calculate elapsed time.

        Parameters
        ----------
        message : Message
            The message that has been received.

        """
        message_type = message.metadata.message_type
        src_node_id = message.metadata.src_node_id
        self.record_elapsed_times_dict[message_type][src_node_id].append(
            time.perf_counter() - self.record_send_msg_dict[message_type][src_node_id],
        )

    def start_record_submit_aggregation(self) -> None:
        """Record the time when an aggregation task is submitted."""
        self.record_submit_aggregation = time.perf_counter()

    def end_record_submit_aggregation(self) -> None:
        """Calculate the elapsed time for the aggregation task submission."""
        elapsed_time = time.perf_counter() - self.record_submit_aggregation
        self.record_submit_aggregation_times.append(elapsed_time)

    def start_record_waiting_for_aggregation(self) -> None:
        """Record the time when waiting for aggregation starts."""
        self.record_waiting_for_aggregation = time.perf_counter()

    def end_record_waiting_for_aggregation(self) -> None:
        """Calculate the elapsed time for waiting for aggregation."""
        elapsed_time = time.perf_counter() - self.record_waiting_for_aggregation
        self.record_waiting_for_aggregation_times = elapsed_time

    def start_record_server_optimization(self) -> None:
        """Record the time when server optimization starts."""
        self.record_server_optimization = time.perf_counter()

    def end_record_server_optimization(self) -> None:
        """Calculate the elapsed time for server optimization."""
        elapsed_time = time.perf_counter() - self.record_server_optimization
        self.record_server_optimization_times = elapsed_time

    def compose_metrics(self) -> dict[str, Scalar]:
        """Compose the metrics based on recorded message times.

        Returns
        -------
        dict[str, Scalar]
            A dictionary containing the average, max, min, and standard deviation of
            message times for each message type.

        """
        metrics: dict[str, Scalar] = {}
        for message_type, node_times in self.record_elapsed_times_dict.items():
            # Get all times for this message type
            times = [
                time / self.divider
                for _node_id, times in node_times.items()
                for time in times
            ]
            # Calculate average time for this message type
            metrics[f"server/time/{message_type}_msg/avg"] = cast(
                "Scalar",
                sum(times) / len(times) if times else 0.0,
            )
            # Calculate max time for this message type
            metrics[f"server/time/{message_type}_msg/max"] = cast(
                "Scalar",
                max(times) if times else 0.0,
            )
            # Calculate min time for this message type
            metrics[f"server/time/{message_type}_msg/min"] = cast(
                "Scalar",
                min(times) if times else 0.0,
            )
            # Calculate the standard deviation for this message type
            metrics[f"server/time/{message_type}_msg/std"] = cast(
                "Scalar",
                (
                    (
                        sum(
                            (x - float(metrics[f"server/time/{message_type}_msg/avg"]))
                            ** 2
                            for x in times
                        )
                        / len(times)
                    )
                    ** 0.5
                    if times
                    else 0.0
                ),
            )
        # Add aggregation task submission time metrics
        metrics["server/time/submit_aggregation/avg"] = cast(
            "Scalar",
            (
                sum(self.record_submit_aggregation_times)
                / len(self.record_submit_aggregation_times)
                if self.record_submit_aggregation_times
                else 0.0
            ),
        )
        metrics["server/time/submit_aggregation/max"] = cast(
            "Scalar",
            (
                max(self.record_submit_aggregation_times)
                if self.record_submit_aggregation_times
                else 0.0
            ),
        )
        metrics["server/time/submit_aggregation/min"] = cast(
            "Scalar",
            (
                min(self.record_submit_aggregation_times)
                if self.record_submit_aggregation_times
                else 0.0
            ),
        )
        metrics["server/time/submit_aggregation/std"] = cast(
            "Scalar",
            (
                (
                    sum(
                        (x - float(metrics["server/time/submit_aggregation/avg"])) ** 2
                        for x in self.record_submit_aggregation_times
                    )
                    / len(self.record_submit_aggregation_times)
                )
                ** 0.5
                if self.record_submit_aggregation_times
                else 0.0
            ),
        )
        # Reset the aggregation task submission times
        self.record_submit_aggregation_times = []
        # Add waiting for aggregation time metrics
        metrics["server/time/waiting_for_aggregation"] = cast(
            "Scalar",
            self.record_waiting_for_aggregation_times / self.divider,
        )
        # Add server optimization time metrics
        metrics["server/time/server_optimization"] = cast(
            "Scalar",
            self.record_server_optimization_times / self.divider,
        )

        return metrics
