"""Module handling expert assignments for clients."""

import ast
import os
import secrets
from logging import INFO

from flwr.common import log

from mt_dao.conf.base_schema import BaseConfig, ExpertsAssignmentPolicyName

FIRST_EXECUTION = True
mt_dao_CLIENT_TO_EXPERTS_IDS: dict[int, list[int]] = ast.literal_eval(
    os.getenv("mt_dao_CLIENT_TO_EXPERTS_IDS", "{}"),
)


class ExpertsAssignmentPolicyNotImplementedError(Exception):
    """Error raised when an experts assignment policy is not implemented."""

    def __init__(self, message: str) -> None:
        """Create the exception with the message.

        Parameters
        ----------
        message: str
            The error message to display.

        """
        super().__init__(message)


def random_assignment(
    expert_ids_to_assign: set[int],
    clients_per_round: int,
    client_id: str | int,
    experts_overlapping_factor: int,
    rng: secrets.SystemRandom,
) -> list[int]:
    """Assign experts to clients randomly with consistent overlapping factor.

    Parameters
    ----------
    expert_ids_to_assign : set[int]
        The set of expert IDs to assign to clients.
    clients_per_round : int
        The number of clients participating in each round.
    client_id : str | int
        The ID of the client requesting expert assignments.
    experts_overlapping_factor : int
        The number of clients each expert should be assigned to.
    rng : secrets.SystemRandom
        A random number generator for shuffling expert IDs.

    Returns
    -------
    list[int]
        A list of expert IDs assigned to the client.

    Raises
    ------
    ValueError
        If the total number of experts is not divisible by the product of clients per
        round and the experts overlapping factor.

    """
    # Sort and shuffle the list of expert IDs for reproducibility
    expert_ids = sorted(expert_ids_to_assign)
    rng.shuffle(expert_ids)
    n_experts_to_assign = len(expert_ids)

    # Calculate how many experts each client gets
    # Each expert should be assigned to `experts_overlapping_factor` clients
    # NOTE: We assume the practitioners know what they are dealing with, but, just in
    # case, we raise an error if the values mismatch
    if n_experts_to_assign % (clients_per_round * experts_overlapping_factor) != 0:
        msg = (
            "The total number of experts must be divisible by "
            f"{clients_per_round} * {experts_overlapping_factor}. "
            f"Got {n_experts_to_assign}, {clients_per_round}, and"
            f" {experts_overlapping_factor}."
        )
        raise ValueError(msg)
    experts_per_client = (
        n_experts_to_assign * experts_overlapping_factor
    ) // clients_per_round

    # Convert client_id to int for consistent indexing
    client_idx = int(client_id)

    # Calculate start and end indices for this client
    start_idx = (client_idx * experts_per_client) % n_experts_to_assign
    end_idx = start_idx + experts_per_client

    # Handle wraparound if we go beyond the total number of experts
    if end_idx <= n_experts_to_assign:
        client_expert_ids = expert_ids[start_idx:end_idx]
    else:
        client_expert_ids = (
            expert_ids[start_idx:] + expert_ids[: end_idx - n_experts_to_assign]
        )

    return client_expert_ids


def fixed_assignment(expert_ids_to_assign: set[int], client_id: str | int) -> list[int]:
    """Assign experts to clients based on a fixed mapping from env var.

    This function assigns experts to clients based on a predefined mapping
    stored in the environment variable `mt_dao_CLIENT_TO_EXPERTS_IDS`.

    Parameters
    ----------
    expert_ids_to_assign : set[int]
        The set of expert IDs to assign to clients.
    client_id : str | int
        The ID of the client requesting expert assignments.

    Returns
    -------
    list[int]
        A list of expert IDs assigned to the client.

    Raises
    ------
    ValueError
        If the client ID is not found in the expert assignments or
        if the expert assignments are inconsistent.

    """
    # Read the env var
    if int(client_id) not in mt_dao_CLIENT_TO_EXPERTS_IDS:
        msg = (
            f"Client ID {client_id} not found in expert assignments"
            f" ({mt_dao_CLIENT_TO_EXPERTS_IDS})."
        )
        raise ValueError(msg)
    client_experts_to_assign = sorted(mt_dao_CLIENT_TO_EXPERTS_IDS[int(client_id)])
    # Check if they are consistent
    if any(
        expert_id not in expert_ids_to_assign for expert_id in client_experts_to_assign
    ):
        msg = (
            f"Expert assignments for client ID {client_id} are inconsistent: we are"
            f" expecting to assign expert IDs in {expert_ids_to_assign}, but trying"
            f" to assign {client_experts_to_assign}."
        )
        raise ValueError(msg)
    return client_experts_to_assign


def experts_assignment(
    cfg: BaseConfig,
    client_id: str | int,
    server_round: int,
) -> str | None:
    """Determine the expert IDs assigned to a client for a given server round.

    This function uses the specified experts assignment policy to determine
    which expert IDs are assigned to the client.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object.
    client_id : str | int
        The ID of the client.
    server_round : int
        The current server round.

    Returns
    -------
    str | None
        The expert IDs assigned to the client for the given server round.

    Raises
    ------
    ExpertsAssignmentPolicyNotImplementedError
        If the experts assignment policy is not implemented.

    """
    if not (
        hasattr(cfg.llm_config.model, "ffn_config")
        and hasattr(cfg.llm_config.model.ffn_config, "ff_n_experts")
        and hasattr(cfg.llm_config.model.ffn_config, "ff_n_active_experts")
    ):
        global FIRST_EXECUTION  # noqa: PLW0603
        if FIRST_EXECUTION:
            log(INFO, "Experts assignment not applicable.")
            FIRST_EXECUTION = False
        return None

    # Get configuration values
    total_experts = cfg.llm_config.model.ffn_config.ff_n_experts
    clients_per_round = cfg.fl.n_clients_per_round
    experts_assignment_policy_name = (
        cfg.fl.experts_assignment_policy.experts_assignment_policy_name
    )
    experts_overlapping_factor = (
        cfg.fl.experts_assignment_policy.experts_overlapping_factor
    )
    shared_expert_ids = cfg.fl.experts_assignment_policy.shared_expert_ids

    # Shuffle with seeded RNG for reproducibility
    # TODO(anon): Ensure this doesn't impact other PRNGs and it is completely
    # reproducible irrespective of what other PRNG do.
    rng = secrets.SystemRandom(cfg.seed + server_round)

    # Dispatch the assignment strategy
    expert_ids_to_assign = set(range(total_experts)) - set(shared_expert_ids)
    if experts_assignment_policy_name == ExpertsAssignmentPolicyName.RANDOM:
        return str(
            shared_expert_ids
            + random_assignment(
                expert_ids_to_assign=expert_ids_to_assign,
                clients_per_round=clients_per_round,
                client_id=client_id,
                rng=rng,
                experts_overlapping_factor=experts_overlapping_factor,
            ),
        )
    if experts_assignment_policy_name == ExpertsAssignmentPolicyName.FIXED:
        return str(
            shared_expert_ids
            + fixed_assignment(
                expert_ids_to_assign=expert_ids_to_assign,
                client_id=client_id,
            ),
        )
    msg = f"Unknown experts assignment policy: {experts_assignment_policy_name}"
    raise ExpertsAssignmentPolicyNotImplementedError(msg)
