"""Estimate total time of training."""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

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

if TYPE_CHECKING:
    from mt_dao.server.server_state import ServerState


__all__ = ["RuntimeEstimator"]


class RuntimeEstimator:
    """Estimates total training time.

    This callback estimates the remaining time for training based on the time taken
    for previous rounds. It can be configured to use different time units such as
    seconds, minutes, hours, or days. The estimated time is added to the server's
    history after each round.

    """

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

        Parameters
        ----------
        time_unit : str, optional
            The unit of time to use for the estimate. 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'.

        """
        self._enabled = True
        self.start_time: float | None = None
        self.total_rounds: int | None = None

        # Timer attributes for each phase
        self._broadcast_start_time: float | None = None
        self._eval_fit_round_start_time: float | None = None
        self._checkpoint_start_time: float | None = None
        self._previously_elapsed_time = 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 init(self, server_state: ServerState) -> None:
        """Initialize the callback.

        This method is called at the start of training, before any rounds have started.

        Parameters
        ----------
        server_state : ServerState
            The current state of the server.

        """
        if self._enabled:
            server_state.history.add_metrics_centralized(
                server_round=server_state.current_round,
                metrics={"server/time/remaining_estimate_unit": self.time_unit},
            )
            # NOTE: This is the number of total rounds left to train from the latest
            # resumption, so we need to subtract the ServerState's start_round
            if self.total_rounds is None:
                self.total_rounds = server_state.total_rounds - server_state.start_round

    def round_start(self, server_state: ServerState) -> None:  # noqa: ARG002
        """Start the timer for the current round.

        Parameters
        ----------
        server_state : ServerState
            The current state of the server.

        """
        if not self._enabled:
            return
        if self.start_time is None:
            self.start_time = time.perf_counter()
        return

    def round_end(self, server_state: ServerState) -> None:
        """Estimate the remaining time after the current round.

        Parameters
        ----------
        server_state : ServerState
            The current state of the server.

        """
        if not self._enabled:
            return

        assert self.start_time is not None, (
            "start_time should be set on round_start if enabled"
        )
        assert self.total_rounds is not None, (
            "total_rounds should be set on init if enabled"
        )
        elapsed_time = time.perf_counter() - self.start_time
        last_round_completion_time = elapsed_time - self._previously_elapsed_time
        elapsed_rounds = server_state.current_round - server_state.start_round
        rounds_left = self.total_rounds - elapsed_rounds
        avg_round_completion_time = elapsed_time / elapsed_rounds
        remaining_time = avg_round_completion_time * rounds_left
        self._previously_elapsed_time = elapsed_time
        server_state.history.add_metrics_centralized(
            server_round=server_state.current_round,
            metrics={
                "server/time/total_time_elapsed": elapsed_time / self.divider,
                "server/time/remaining_estimate": remaining_time / self.divider,
                "server/time/round_completion_time": (
                    last_round_completion_time / self.divider
                ),
                "server/time/round_elapsed": elapsed_rounds,
                "server/time/rounds_left": rounds_left,
                "server/time/avg_round_completion_time": (
                    avg_round_completion_time / self.divider
                ),
            },
        )

    def broadcast_start(self, server_state: ServerState) -> None:  # noqa: ARG002
        """Mark the start of the broadcast phase."""
        if not self._enabled:
            return
        self._broadcast_start_time = time.perf_counter()

    def broadcast_end(self, server_state: ServerState) -> None:
        """Log the elapsed time for the broadcast phase."""
        if not self._enabled:
            return
        assert self._broadcast_start_time is not None, (
            "broadcast_start must be called before broadcast_end"
        )
        elapsed = time.perf_counter() - self._broadcast_start_time
        server_state.history.add_metrics_centralized(
            server_round=server_state.current_round,
            metrics={"server/time/broadcast_elapsed": elapsed / self.divider},
        )

    def eval_fit_round_start(self, server_state: ServerState) -> None:  # noqa: ARG002
        """Mark the start of the fit and gather phase."""
        if not self._enabled:
            return
        self._eval_fit_round_start_time = time.perf_counter()

    def eval_fit_round_end(self, server_state: ServerState) -> None:
        """Log the elapsed time for the fit and gather phase."""
        if not self._enabled:
            return
        assert self._eval_fit_round_start_time is not None, (
            "eval_fit_round_start must be called before eval_fit_round_end"
        )
        elapsed = time.perf_counter() - self._eval_fit_round_start_time
        server_state.history.add_metrics_centralized(
            server_round=server_state.current_round,
            metrics={"server/time/eval_fit_round_elapsed": elapsed / self.divider},
        )

    def checkpoint_start(self, server_state: ServerState) -> None:  # noqa: ARG002
        """Mark the start of the checkpoint phase."""
        if not self._enabled:
            return
        self._checkpoint_start_time = time.perf_counter()

    def checkpoint_end(self, server_state: ServerState) -> None:
        """Log the elapsed time for the checkpoint phase."""
        if not self._enabled:
            return
        assert self._checkpoint_start_time is not None, (
            "checkpoint_start must be called before checkpoint_end"
        )
        elapsed = time.perf_counter() - self._checkpoint_start_time
        server_state.history.add_metrics_centralized(
            server_round=server_state.current_round,
            metrics={"server/time/checkpoint_elapsed": elapsed / self.divider},
        )
