"""Implementation of the Flower's ServerApp for orchestrating federate learning."""

import contextlib
import os
import time
import warnings
from logging import DEBUG, INFO
from queue import Queue
from typing import cast

import flwr as fl
from flwr.common import (
    Context,
)
from flwr.common.logger import log, update_console_handler
from flwr.server import Grid
from omegaconf import OmegaConf

import ray
import wandb
from mt_dao.callbacks.runtime_estimator import RuntimeEstimator
from mt_dao.conf.base_schema import BaseConfig
from mt_dao.constants import (
    CONFIG_FILE_YAML,
    MAX_BROADCAST_RETRIES,
    mt_dao_SAVE_PATH,
    SERVER_PROCESS_NAME,
)
from mt_dao.server.broadcast_utils import broadcast_parameters_to_nodes
from mt_dao.server.fit_eval import eval_fit_round
from mt_dao.server.init_utils import (
    initialize_server_state,
)
from mt_dao.server.load_balancer import ServerLoadBalancer
from mt_dao.server.s3_utils import (
    cleanup_checkpoints,
    upload_server_checkpoint_model_states,
    upload_server_checkpoint_server_state,
)
from mt_dao.server.server_state import ServerState
from mt_dao.server.utils import (
    wait_for_nodes_to_connect,
)
from mt_dao.strategy_functional.dispatcher import dispatch_strategy_functional
from mt_dao.utils import (
    custom_ray_garbage_collector,
    custom_shm_garbage_collector,
    wandb_init,
)

# Fix the logger
update_console_handler(level=DEBUG, colored=False, timestamps=True)
# Filter user warning from configuration of MPT
warnings.filterwarnings(
    action="ignore",
    category=UserWarning,
    message=("If not using a Prefix Language Model*"),
    append=True,
)
# TODO(<anon>): These don't work -- not sure why
# Filter deprecation warning from pkg_resources
warnings.filterwarnings(
    action="ignore",
    category=DeprecationWarning,
    message=("Deprecated call to *"),
    append=True,
)
warnings.filterwarnings(
    action="ignore",
    category=DeprecationWarning,
    message=("pkg_resources is deprecated*"),
    append=True,
)

# Run via `flower-server-app server:app`
app = fl.server.ServerApp()


def federated_training_loop(  # noqa: C901, PLR0915
    cfg: BaseConfig,
    server_state: ServerState,
    runtime_estimator: RuntimeEstimator,
    grid: Grid,
) -> None:
    """Run the federated learning training loop.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object containing the settings for the server.
    server_state : ServerState
        The current state of the server, including parameters and history.
    runtime_estimator : RuntimeEstimator
        The runtime estimator for tracking the performance of the server.
    grid : Grid
        The driver object for the server, managing the connected nodes.

    Raises
    ------
    ValueError
        If the `mt_dao_SAVE_PATH` environment variable is not set.
    TimeoutError
        If the server app fails to broadcast the parameters to nodes.

    """
    # Run federated learning for number of rounds
    log(INFO, "FL starting from round %s", server_state.start_round)
    start_time = time.perf_counter()

    # Create the LoadBalancer object
    server_load_balancer = ServerLoadBalancer(
        grid=grid,
        total_number_of_clients=cfg.fl.n_total_clients,
        rng=server_state.rng,
        is_production=cfg.mt_dao.is_production,
    )
    # Federated learning loop
    for current_federated_round in range(
        server_state.start_round + 1,
        cfg.fl.n_rounds + 1,
    ):
        runtime_estimator.round_start(server_state)
        server_state.current_round = current_federated_round

        log(DEBUG, f"Commencing server round {server_state.current_round}")

        # Sampling clients for the current round
        server_state.previously_sampled_clients = server_state.sampled_clients
        server_state.sampled_clients = server_state.client_sampler.sample_clients(
            rng=server_state.rng,
        )

        # Broadcast the model parameters to all NodeManagers
        runtime_estimator.broadcast_start(server_state)
        # Check for changes in connected NodeManagers
        server_load_balancer.check_nodes_health()
        retry_count = 0
        broadcast_successful = False
        while retry_count < MAX_BROADCAST_RETRIES and not broadcast_successful:
            try:
                broadcast_parameters_to_nodes(
                    server_load_balancer=server_load_balancer,
                    server_state=server_state,
                    comm_stack=cfg.mt_dao.comm_stack,
                )
                broadcast_successful = True
                log(DEBUG, "Broadcast successful on attempt %s", retry_count + 1)
            except TimeoutError as e:
                retry_count += 1
                if retry_count < MAX_BROADCAST_RETRIES:
                    log(
                        INFO,
                        "Broadcast failed (attempt %s/%s): %s. Retrying...",
                        retry_count,
                        MAX_BROADCAST_RETRIES,
                        str(e),
                    )

                else:
                    log(
                        INFO,
                        "Broadcast failed after %s attempts: %s. Giving up.",
                        MAX_BROADCAST_RETRIES,
                        str(e),
                    )
                    raise
        runtime_estimator.broadcast_end(server_state)

        # Save the checkpoint to S3 Object Store
        # TODO(anon): Move it to where it is not expensive
        if (
            cfg.mt_dao.checkpointing.local_checkpoint
            or cfg.mt_dao.checkpointing.remote_checkpoint
            or cfg.mt_dao.comm_stack.s3
        ):
            if server_state.remote_up_down is None:
                msg = "Cannot checkpoint without a RemoteUploaderDownloader object"
                raise ValueError(msg)
            upload_server_checkpoint_model_states(
                server_state=server_state,
            )

        # Fit the model
        runtime_estimator.eval_fit_round_start(server_state)
        eval_fit_round(
            cfg=cfg,
            server_state=server_state,
            server_load_balancer=server_load_balancer,
            comm_stack=cfg.mt_dao.comm_stack,
            server_optimizer=dispatch_strategy_functional(cfg),
        )
        runtime_estimator.eval_fit_round_end(server_state)
        log(DEBUG, f"Fit and gather round {server_state.current_round} completed")

        # Save the checkpoint to S3 Object Store
        runtime_estimator.checkpoint_start(server_state)
        if (
            cfg.mt_dao.checkpointing.local_checkpoint
            or cfg.mt_dao.checkpointing.remote_checkpoint
            or cfg.mt_dao.comm_stack.s3
        ):
            if server_state.remote_up_down is None:
                log(DEBUG, "Ignoring S3 Upload. Only saving locally.")
            upload_server_checkpoint_server_state(
                server_state=server_state,
            )
            upload_server_checkpoint_model_states(
                server_state=server_state,
            )
        # Clean up checkpoints if asked to
        if (
            cfg.mt_dao.checkpointing.local_cleanup_checkpoints_per_round
            or cfg.mt_dao.checkpointing.remote_cleanup_checkpoints_per_round
        ):
            # NOTE: We want to clean the local directory in the case that we want
            # remote checkpoints without preserving the local checkpoints.
            log(DEBUG, "Cleaning the current round.")
            if not cfg.mt_dao.checkpointing.per_round_cleanup_staleness:
                log(DEBUG, "No staleness factor provided. No deleting.")
                end_idx = 0
            else:
                # NOTE: We enforce this constraint so that we can doo negative
                # indexing.
                assert cfg.mt_dao.checkpointing.per_round_cleanup_staleness > 0, (
                    "The staleness factor needs to be greater than 0."
                )
                end_idx = -(cfg.mt_dao.checkpointing.per_round_cleanup_staleness)
            # NOTE: The run_uuid determines whether it is a remote object or not.
            cleanup_checkpoints(
                run_uuid=cfg.run_uuid,
                end_idx=end_idx,
                server_state=server_state,
                local_cleanup=(
                    cfg.mt_dao.checkpointing.local_cleanup_checkpoints_per_round
                ),
                remote_cleanup=(
                    cfg.mt_dao.checkpointing.remote_cleanup_checkpoints_per_round
                ),
            )
        runtime_estimator.checkpoint_end(server_state)
        runtime_estimator.round_end(server_state)
    log(
        INFO,
        "FL finished in %s",
        time.perf_counter() - start_time + server_state.time_offset,
    )


@app.main()
def main(
    grid: Grid,
    _context: Context,
) -> None:
    """Implement the main function for the Flower ServerApp.

    Parameters
    ----------
    grid : Grid
        The driver object for the server.
    context : Context
        The context object for the server.

    Raises
    ------
    ValueError
        If the environmental variable `mt_dao_SAVE_PATH` is not set.

    """
    start_up_time = time.perf_counter()
    # Get the environmental variable for the dump folder
    mt_dao_save_path = os.environ.get(mt_dao_SAVE_PATH, "")
    run_uuid = os.environ.get("RUN_UUID", "")
    # Raise an error if the environmental variable is not set
    if not mt_dao_save_path:
        msg = f"{mt_dao_SAVE_PATH} is not set."
        raise ValueError(msg)
    # NOTE: This convention is required so we can correctly achieve local
    # restore round functionality.
    if (run_uuid not in mt_dao_save_path) and run_uuid:
        msg = f"We enforce that the run_uuid={run_uuid} is in the {mt_dao_SAVE_PATH}."
        raise ValueError(msg)
    # Load the configuration from the config file
    cfg = cast("BaseConfig", OmegaConf.load(mt_dao_save_path + f"/{CONFIG_FILE_YAML}"))
    log(INFO, "Initializing mt_dao Server")
    custom_gc: (
        contextlib._GeneratorContextManager[None] | contextlib.nullcontext[None]
    ) = contextlib.nullcontext()
    server_state = initialize_server_state(cfg, mt_dao_save_path)
    server_state.local_checkpoint_path = mt_dao_save_path
    # Creating Queue for custom Ray garbage collector
    if cfg.mt_dao.comm_stack.ray:
        log(
            INFO,
            "ServerApp: Ray communication stack is enabled. Initializing Ray."
            " ray.is_initialized()=%s",
            ray.is_initialized(),
        )
        server_state.params_garbage_queue = Queue()
        # Initialize Ray
        ray_context = ray.init("auto")
        log(
            INFO,
            "ServerApp: Ray initialized with context %s\n\t ray.is_initialized()=%s",
            ray_context,
            ray.is_initialized(),
        )
        custom_gc = custom_ray_garbage_collector(
            garbage_queue=server_state.params_garbage_queue,
            list_of_threads=[],
            process_name=SERVER_PROCESS_NAME,
        )
    if cfg.mt_dao.comm_stack.shm:
        server_state.params_garbage_queue = Queue()
        custom_gc = custom_shm_garbage_collector(
            garbage_queue=server_state.params_garbage_queue,
            list_of_threads=[],
            process_name=SERVER_PROCESS_NAME,
        )

    log(
        INFO,
        "Resolved communication stack",
    )
    # Resolve WandB config
    wandb_config = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
    with (
        wandb_init(  # type: ignore[union-attr,misc]
            cfg.use_wandb,
            **cfg.wandb.setup,  # type: ignore[reportCallIssue]
            settings=wandb.Settings(start_method="thread"),  # type: ignore[arg-type]
            config=wandb_config,  # type: ignore[arg-type]
        ),
        custom_gc,
    ):
        log(
            INFO,
            "Wrapping the server execution with context managers",
        )
        # Wait for the minimum number of nodes to connect
        wait_for_nodes_to_connect(grid, cfg.mt_dao.n_nodes)
        # Init the runtime estimator callback
        runtime_estimator = RuntimeEstimator()
        runtime_estimator.init(server_state)

        log(
            INFO,
            "Start-up time for the server is %s",
            time.perf_counter() - start_up_time,
        )
        federated_training_loop(
            cfg=cfg,
            server_state=server_state,
            runtime_estimator=runtime_estimator,
            grid=grid,
        )

        # Bookkeeping
        log(
            DEBUG,
            "app_fit: losses_distributed %s",
            str(server_state.history.losses_distributed),
        )
        log(
            DEBUG,
            "app_fit: metrics_distributed_fit %s",
            str(server_state.history.metrics_distributed_fit),
        )
        log(
            DEBUG,
            "app_fit: metrics_distributed %s",
            str(server_state.history.metrics_distributed),
        )
        log(
            DEBUG,
            "app_fit: losses_centralized %s",
            str(server_state.history.losses_centralized),
        )
        log(
            DEBUG,
            "app_fit: metrics_centralized %s",
            str(server_state.history.metrics_centralized),
        )

        # Clean up checkpoints if asked to
        if (
            cfg.mt_dao.checkpointing.local_cleanup_checkpoints
            or cfg.mt_dao.checkpointing.remote_cleanup_checkpoints
        ):
            # NOTE: The run_uuid determines whether it is a remote object or not.
            log(DEBUG, "Cleaning all checkpoints.")
            cleanup_checkpoints(
                run_uuid=cfg.run_uuid,
                end_idx=None,
                server_state=server_state,
                local_cleanup=(cfg.mt_dao.checkpointing.local_cleanup_checkpoints),
                remote_cleanup=(cfg.mt_dao.checkpointing.remote_cleanup_checkpoints),
            )
