"""Client Application for a Flower-based LLM training node.

This module defines a client-side application using the Flower framework. It handles
parameter broadcasting, training, evaluation, and communication with a node manager.
"""

import ast
import copy
import os
import time
import warnings
from collections.abc import Generator
from logging import DEBUG
from typing import Any, cast

from flwr.common import (
    Context,
    Message,
    log,
)
from flwr.common.logger import update_console_handler

from mt_dao.constants import (
    BATCH_ID,
    BROADCAST,
    BROADCAST_INS,
    BROADCAST_RES,
    CLIENT_STATE_ACCUMULATOR,
    COMM_STACK,
    EVAL_INS,
    EVALUATE_RESULTS,
    EXPECTED_LATENCY,
    FIT_INS,
    FIT_RES,
    FREE_RESOURCES,
    GATHER,
    GATHER_INS,
    GATHER_RES,
    MASK,
)
from mt_dao.masks_utils import (
    combine_masks,
    generate_empty_mask,
    mask_to_batches,
)
from mt_dao.messages.gather import build_gather_res_recorddict
from mt_dao.messages.generic import (
    build_metrics_recorddict,
    build_success_recorddict,
    build_success_with_metrics_recorddict,
)
from mt_dao.node_manager.messages_utils import (
    GatherParametersState,
    extract_fit_eval_ins,
    extract_gather_parameters_state,
    extract_set_parameters_state,
    get_and_set_parameters_and_transmission_mask,
    insert_gather_parameters_state,
    insert_set_parameters_state,
    transform_record_dict_to_mt_dao,
)
from mt_dao.node_manager.node_manager_app import NodeManagerApp
from mt_dao.server.comms import (
    load_recorddict_parameters_from_remote,
    offload_recorddict_parameters_to_remote,
)
from mt_dao.shm.utils import (
    compress_with_strict,
)
from mt_dao.utils import set_parameters_in_place

# 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,
)

# Flower ClientApp
app = NodeManagerApp()


@app.lifespan()
def lifespan(ctx: Context) -> Generator[None, Any, None]:  # noqa: ARG001
    """Wrap ClientApp function with a context manager.

    The current functionality allows to initializing the NodeManagerApp shared memory
    parameters at the time where the communication happens so as to gain some wall time.

    Parameters
    ----------
    ctx : Context
        Context passed by Flower, which contains the stateful records for the
        application, including parameters, masks, and configuration.

    Yields
    ------
    Generator[None, Any, None]
        Standard return time for context managers.

    """
    # Perform initialization tasks before the app starts
    log(DEBUG, "Initializing ClientApp")

    yield  # ClientApp is running

    # Perform cleanup tasks after the app stops
    log(DEBUG, "Cleaning up ClientApp")
    start_time = time.perf_counter()

    # Close the workers after training if the environment variable
    # CLIENTAPP_NO_RELOAD is not set.
    if "CLIENTAPP_NO_RELOAD" not in os.environ:
        log(DEBUG, "Environmental variable CLIENTAPP_NO_RELOAD hasn't been set.")
        # NOTE: We need to deep copy the worker parameters to avoid segmentation faults
        # when the workers are closed.
        app.worker_parameters = copy.deepcopy(app.worker_parameters)
        app.close_workers()

    log(
        DEBUG,
        "ClientApp lifespan cleanup took %s seconds",
        time.perf_counter() - start_time,
    )


@app.query(BROADCAST)
def set_parameters(msg: Message, ctx: Context) -> Message:
    """Set model parameters in shared memory from a broadcast message.

    This function handles the reception of model parameters from the server via a
    broadcast message. It manages mask handling for sparse or batched parameter updates,
    loads the parameters from the communication stack (S3, shared memory, or Ray), and
    writes the received parameters into local shared memory. The function also updates
    the context with the current transmission mask and parameter configuration for
    subsequent training or evaluation steps.

    Parameters
    ----------
    msg : Message
        The incoming broadcast message containing parameter data and configuration.
    ctx : Context
        The execution context, including stateful records for parameters, masks, and
        configuration.

    Returns
    -------
    Message
        A reply message indicating successful parameter update, to be sent back to the
        server.

    Notes
    -----
    - Supports both full and batched parameter updates using masks.
    - Ensures only the relevant subset of parameters is written to shared memory.
    - Updates context state for downstream training/evaluation steps.

    """
    ctx.state = transform_record_dict_to_mt_dao(ctx.state)
    # Get relevant objects from the context
    set_parameters_state = extract_set_parameters_state(
        state=ctx.state,
        node_manager_app=app,
    )
    # The payload we receive may be sparse based what the server sent thus, when we
    # write to the shm we only need to write a subset of the parameters. The mask is a
    # tuple of (mask, layer_names, layer_types) where mask is a tuple of: (a) booleans
    # indicating which parameters to keep; (b) layer_names is a list of strings
    # indicating the names of the layers; (c) layer_types is a list of strings
    # indicating the types of the layers.
    mask: tuple[tuple[bool, ...], list[str], list[str]] = ast.literal_eval(
        str(msg.content.config_records[f"{BROADCAST_INS}.{COMM_STACK}"][MASK]),
    )
    # For batched communication, only set what has been received currently
    if app.cfg.mt_dao.comm_stack.n_batches:
        # Get the batch ID from the message
        batch_id = msg.content.config_records[f"{BROADCAST_INS}.{COMM_STACK}"][BATCH_ID]
        # If this is the first batch, we need to initialize the whole transmission mask
        # that will be used by the client with the empty mask so we can combine it with
        # all the partial masks
        if batch_id == 0:
            set_parameters_state.transmission_mask = generate_empty_mask(
                layer_names_and_types=app.layer_names_and_types,
            )
        # Combine the currently received mask with the previous one
        set_parameters_state.transmission_mask = combine_masks(
            mask_a=set_parameters_state.transmission_mask,
            mask_b=mask,
            layer_names_and_types=app.layer_names_and_types,
        )
    else:
        # Handle null case, we need to set the mask to the one received
        # from the server
        set_parameters_state.transmission_mask = mask

    # Obtain parameters from the communication stack
    msg, parameters = load_recorddict_parameters_from_remote(
        remote_uploader_downloader=app.remote_up_down,
        incoming_message=msg,
        comm_stack=app.cfg.mt_dao.comm_stack,
        msg_str=BROADCAST_INS,
    )

    # Set parameters to the local shared memory
    set_parameters_in_place(
        list(
            compress_with_strict(
                data=set_parameters_state.parameters,
                selectors=mask[0],
                strict=True,
            ),
        ),
        parameters,
        parameter_pos=range(len(parameters)),
    )

    # Insert the set parameters context into the context
    insert_set_parameters_state(
        state=ctx.state,
        set_parameters_state=set_parameters_state,
    )

    # Create the reply of success
    return Message(
        content=build_success_recorddict(msg_str=BROADCAST_RES),
        reply_to=msg,
    )


@app.query(GATHER)
def gather_parameters(msg: Message, ctx: Context) -> Message:
    """Gather model parameters and metrics for aggregation.

    This function prepares and returns the client's model parameters and training
    metrics for aggregation by the server. It handles mask batching for partial
    aggregation, compresses the parameters according to the aggregation mask, and
    offloads the result using the configured communication stack (S3, shared memory,
    or Ray). Only the first batch includes metrics and client state accumulator.

    Parameters
    ----------
    msg : Message
        The incoming gather message containing aggregation instructions and batch info.
    ctx : Context
        The execution context, including stateful records for parameters, masks,
        metrics, and training results.

    Returns
    -------
    Message
        A reply message containing the compressed parameters, metrics (if first batch),
        and aggregation mask, ready for server-side aggregation.

    Notes
    -----
    - Supports both full and batched aggregation using masks.
    - Only the first batch includes metrics and client state accumulator.
    - Ensures parameters are compressed and offloaded efficiently according to the
      communication stack configuration.

    """
    ctx.state = transform_record_dict_to_mt_dao(ctx.state)
    # Get relevant objects from the context
    gather_parameters_state = extract_gather_parameters_state(ctx.state)

    batch_id = cast(
        "int",
        msg.content.config_records[f"{GATHER_INS}.{COMM_STACK}"][BATCH_ID],
    )

    # Handle batching of masks
    if app.cfg.mt_dao.comm_stack.n_batches:
        gather_parameters_state.aggregation_mask = mask_to_batches(
            full_mask=gather_parameters_state.aggregation_mask,
            layer_names_and_types=app.layer_names_and_types,
            n_batches=app.cfg.mt_dao.comm_stack.n_batches,
        )[batch_id]

    gather_recorddict, _ = offload_recorddict_parameters_to_remote(
        parameters=list(
            compress_with_strict(
                gather_parameters_state.parameters,
                gather_parameters_state.aggregation_mask[0],
                strict=True,
            ),
        ),
        remote_uploader_downloader=app.remote_up_down,
        outgoing_recorddict=build_gather_res_recorddict(
            aggregation_mask=gather_parameters_state.aggregation_mask,
            num_examples=gather_parameters_state.num_samples,
            batch_id=batch_id,
        ),
        comm_stack=app.cfg.mt_dao.comm_stack,
        msg_str=GATHER_RES,
    )
    return Message(content=gather_recorddict, reply_to=msg)


@app.query(FREE_RESOURCES)
def free_resources(msg: Message, ctx: Context) -> Message:
    """Free resources for this client node.

    Parameters
    ----------
    msg : Message
        Incoming message requesting to free resources.
    ctx : Context
        Execution context passed by Flower.

    Returns
    -------
    Message
        Reply message confirming resource deallocation.

    """
    ctx.state = transform_record_dict_to_mt_dao(ctx.state)
    return Message(content=msg.content, reply_to=msg)


@app.train()
def train(msg: Message, ctx: Context) -> Message:
    """Execute training on the client with parameters and configuration from server.

    This function handles the training phase of the federated learning workflow. It
    retrieves the transmission mask and model parameters from the context, applies the
    received parameters to the local model (using shared memory and masking for partial
    updates), and executes the training process. After training, it updates the context
    with the new parameters, training metrics, and aggregation mask, and returns a
    success message to the server.

    Parameters
    ----------
    msg : Message
        The incoming training message containing configuration and instructions from the
        server.
    ctx : Context
        The execution context, including stateful records for parameters, masks, and
        configuration.

    Returns
    -------
    Message
        A reply message indicating successful completion of the training step, including
        updated parameters, metrics, and aggregation mask in the context.

    Notes
    -----
    - Applies only the relevant subset of parameters using the transmission mask.
    - Updates shared memory with the new parameters for downstream steps.
    - Aggregation mask is updated in the context for use in the aggregation phase.
    - Ensures proper cleanup of worker resources after training.

    """
    ctx.state = transform_record_dict_to_mt_dao(ctx.state)
    start_time = time.perf_counter()

    # Get the update parameters and transmission masks from the context and set the
    # NodeManagerApp parameters and transmission mask
    get_and_set_parameters_and_transmission_mask(
        state=ctx.state,
        node_manager_app=app,
    )

    # Extract the fit instructions from the message
    fit_ins = extract_fit_eval_ins(msg.content, ctx.state, FIT_INS)

    # Make sure workers are up and running
    app.create_and_start_workers()

    # Execute the training process using the training instructions
    parameters, n_samples, metrics = app.fit(
        fit_ins=fit_ins,
    )

    # Add relevant parameters back to the context
    gather_parameters_state = GatherParametersState(
        parameters=parameters,
        aggregation_mask=app.aggregation_mask,
        num_samples=n_samples,
    )
    insert_gather_parameters_state(ctx.state, gather_parameters_state)

    # Only return metrics for the first batch
    metrics_to_return = {
        CLIENT_STATE_ACCUMULATOR: repr(metrics.pop(CLIENT_STATE_ACCUMULATOR)),
        **metrics,
    }
    # Return the success message
    return Message(
        content=build_success_with_metrics_recorddict(
            msg_str=FIT_RES,
            metrics=metrics_to_return,
            num_examples=n_samples,
        ),
        reply_to=msg,
        ttl=EXPECTED_LATENCY + int(time.perf_counter() - start_time),
    )


@app.evaluate()
def evaluate(msg: Message, ctx: Context) -> Message:
    """Execute evaluation on the client with parameters and configuration from server.

    This function handles the evaluation phase of the federated learning workflow.
    It retrieves the transmission mask and model parameters from the context, applies
    the received parameters to the local model (using shared memory and masking for
    partial updates), and executes the evaluation process. After evaluation, it
    returns a message to the server containing the evaluation loss, number of samples,
    and evaluation metrics.

    Parameters
    ----------
    msg : Message
        The incoming evaluation message containing configuration and instructions from
        the server.
    ctx : Context
        The execution context, including stateful records for parameters, masks, and
        configuration.

    Returns
    -------
    Message
        A reply message containing the evaluation results, including loss, metrics,
        and number of samples.

    Notes
    -----
    - Applies only the relevant subset of parameters using the transmission mask.
    - Updates shared memory with the parameters for evaluation.
    - Ensures proper cleanup of worker resources after evaluation.

    """
    ctx.state = transform_record_dict_to_mt_dao(ctx.state)
    start_time = time.perf_counter()

    # Get the update parameters and transmission masks from the context and set the
    # NodeManagerApp parameters and transmission mask
    get_and_set_parameters_and_transmission_mask(
        state=ctx.state,
        node_manager_app=app,
    )

    # Extract the fit instructions from the message
    eval_ins = extract_fit_eval_ins(msg.content, ctx.state, EVAL_INS)

    # Make sure workers are up and running
    app.create_and_start_workers()

    # Execute the evaluation process
    loss, n_samples, metrics = app.eval(eval_ins=eval_ins)

    # Compose and return the success message with the evaluation results
    return Message(
        content=build_metrics_recorddict(
            msg_str=EVALUATE_RESULTS,
            metrics=metrics,
            loss=loss,
            num_examples=n_samples,
        ),
        reply_to=msg,
        ttl=EXPECTED_LATENCY + int(time.perf_counter() - start_time),
    )
