"""Utility functions related to file objects."""

import pickle  # noqa: S403
from pathlib import Path
from typing import Any

import numpy as np
import torch
from composer.loggers import RemoteUploaderDownloader
from flwr.common import (
    NDArrays,
    parameters_to_ndarrays,
)
from huggingface_hub import split_torch_state_dict_into_shards
from safetensors import safe_open
from safetensors.torch import save_file

from mt_dao.constants import (
    BIN_SUFFIX,
    CURRENT_CODEBASE_NAME,
    DEFAULT_MAX_SHARD_SIZE,
    NUMPY_ARRAY_COMPRESSED_SUFFIX,
    NUMPY_ARRAY_SUFFIX,
    OLD_CODEBASE_NAME,
    SAFETENSORS_SUFFIX,
)


def download_file_from_s3(
    remote_up_down: RemoteUploaderDownloader,
    remote_file_name: str,
    local_file_name: Path | str,
) -> None:
    """Download a file from S3.

    Parameters
    ----------
    remote_up_down : RemoteUploaderDownloader
        The remote uploader/downloader object.
    remote_file_name : str
        The name of the remote file to download.
    local_file_name : Path | str
        The local file path where the file will be saved.

    """
    remote_up_down._check_workers()  # noqa: SLF001
    remote_up_down.download_file(
        remote_file_name=remote_file_name,
        destination=str(local_file_name),
        overwrite=True,
    )


def upload_file_to_s3(
    remote_up_down: RemoteUploaderDownloader,
    remote_file_name: str,
    local_file_name: Path,
) -> None:
    """Upload a file to S3.

    Parameters
    ----------
    remote_up_down : RemoteUploaderDownloader
        The remote uploader/downloader object.
    remote_file_name : str
        The name of the remote file to upload.
    local_file_name : Path
        The local file path of the file to upload.

    """
    remote_up_down._check_workers()  # noqa: SLF001
    remote_up_down.upload_file(
        state=None,
        remote_file_name=remote_file_name,
        file_path=local_file_name,
        overwrite=True,
    )


def load_model_states_from_file(
    file_path: Path,
) -> tuple[list[str], NDArrays]:
    """Load model states from a file.

    This function loads model states from a file and returns the layer names and
    corresponding tensors. It supports multiple file formats including NumPy arrays,
    binary files, and safetensors. The function handles both cases where layer names
    are provided or not.

    Parameters
    ----------
    file_path : Path
        The path to the file containing the model states.

    Returns
    -------
    tuple[list[str], NDArrays]
        A tuple containing:
         - A list of layer names.
         - A list of tensors corresponding to the loaded model states.

    Raises
    ------
    ValueError
        If the file format is not supported or if the file does not contain valid model
        states.

    """
    # NumPy format, no layer names provided
    if file_path.suffix in {NUMPY_ARRAY_SUFFIX, NUMPY_ARRAY_COMPRESSED_SUFFIX}:
        with file_path.open("rb") as file:
            data = np.load(file)
            return [], [data[key] for key in data.files]
    # .bin format, no layer names provided
    elif file_path.suffix == BIN_SUFFIX:
        with file_path.open("rb") as file:
            return [], parameters_to_ndarrays(pickle.load(file))  # noqa: S301
    # Safetensors format, layer names are provided
    elif file_path.suffix == SAFETENSORS_SUFFIX:
        all_layer_names: list[str] = []
        all_tensors: NDArrays = []
        # NOTE: We assume numpy because of underlying framework.
        # This is used in conjunction with loading sharded models.
        with safe_open(file_path, framework="numpy", device="cpu") as f:
            for key in f.keys():  # noqa: SIM118
                all_layer_names.append(key)
                all_tensors.append(f.get_tensor(key))
        return all_layer_names, all_tensors
    else:
        msg = f"Unsupported file format: {file_path.suffix}"
        raise ValueError(msg)


def dump_model_parameters_to_file(
    file_path: Path,
    model_parameters: NDArrays | dict[str, torch.Tensor],
) -> None:
    """Load model parameters from a file.

    Parameters
    ----------
    file_path : Path
        The file path.
    model_parameters : NDArrays
        The model parameters.

    Raises
    ------
    ValueError
        If the file format is not supported.

    """
    # NOTE: Very slow for big models b/c compression. Good benchmark available here: https://stackoverflow.com/questions/30329726/fastest-save-and-load-options-for-a-numpy-array
    if file_path.suffix == NUMPY_ARRAY_COMPRESSED_SUFFIX:
        assert isinstance(
            model_parameters, list,
        ), "The model parameters should be a list for compressed NumPy arrays."
        with file_path.open("wb") as file:
            np.savez_compressed(file, *model_parameters)
    elif file_path.suffix == BIN_SUFFIX:
        assert isinstance(
            model_parameters, list,
        ), "The model parameters should be a list for compressed NumPy arrays."
        with file_path.open("wb") as file:
            pickle.dump(model_parameters, file)
    elif file_path.suffix == NUMPY_ARRAY_SUFFIX:
        assert isinstance(
            model_parameters, list,
        ), "The model parameters should be a list for compressed NumPy arrays."
        with file_path.open("wb") as file:
            np.savez(file, *model_parameters)
    elif file_path.suffix == SAFETENSORS_SUFFIX:
        assert isinstance(
            model_parameters, dict,
        ), "The model parameters should be a dictionary for safetensors."
        save_file(model_parameters, file_path)
    else:
        msg = f"Unsupported file format: {file_path.suffix}"
        raise ValueError(msg)


def chunk_and_save_file(
    layer_names: list[str],
    model_arrays: NDArrays,  # type: ignore[PGH003]
    save_directory: Path,
    max_shard_size: str = DEFAULT_MAX_SHARD_SIZE,
) -> None:
    """Save models into chunks of a maximally defined size.

    Re-implements a similar sharding procedure as to HuggingFace transformers
    to ensure that the saving of a single checkpoint is not "too" big.

    Parameters
    ----------
    layer_names : list[str]
        List of layer names
    model_arrays : NDArrays
        List of model or optimizer state parameters.
    save_directory : Path
        List of model or optimizer state parameters.
    max_shard_size : str, optional
        the maximum size of a single shard, by default "5GB"

    Raises
    ------
    NotADirectoryError
        If the save directory is not a directory.

    """
    if save_directory.is_file():
        message = "The save_directory has to be a directory and not a file."
        raise NotADirectoryError(message)
    if not save_directory.exists():
        save_directory.mkdir(parents=True, exist_ok=True)
    assert len(model_arrays) == len(
        layer_names,
    ), "The tensors and model names should be of the same length."
    # NOTE: Explicit cast to Torch tensors, and ensure the device is on CPU
    # to avoid VRAM overflow.
    model_tensors: list[torch.Tensor] = [
        torch.Tensor(x, device="cpu") for x in model_arrays
    ]
    state_dict_construction = dict(zip(layer_names, model_tensors, strict=False))

    sharded_state_dict = split_torch_state_dict_into_shards(
        state_dict=state_dict_construction,
        max_shard_size=max_shard_size,
    )

    for shard_file, tensors in sharded_state_dict.filename_to_tensors.items():
        shard = {}
        for tensor in tensors:
            shard[tensor] = state_dict_construction[tensor].contiguous()
            # NOTE: This functionality is similary to that used in HF transformers
            # For the delete reference, see https://github.com/huggingface/transformers/pull/34890
            del state_dict_construction[tensor]
        dump_model_parameters_to_file(
            file_path=save_directory / shard_file,
            model_parameters=shard,
        )


def _walk_load_directory(directory: Path) -> list[Path]:
    """Retrieve all `.safetensors` files within a directory.

    Parameters
    ----------
    directory : Path
       The directory path to search.

    Returns
    -------
    list[Path]
        A list of paths to the found `.safetensors` files.

    """
    return list(directory.rglob(f"*{SAFETENSORS_SUFFIX}"))


def load_chunked_model(load_directory: Path) -> tuple[NDArrays, list[str]]:
    """Load a chunked model containing safe tensors files.

    This function loads a model from a directory that potentially contains
    the model in a chunked format.

    Parameters
    ----------
    load_directory : Path
        The path containing the directory to load files from.

    Returns
    -------
    tuple[NDArrays, list[str]]
        A tuple containing:
         - the tensors corresponding to the loaded model.
         - the corresponding layer names.

    """
    safetensor_files = _walk_load_directory(load_directory)
    all_tensors: NDArrays = []
    all_layer_names: list[str] = []
    for safetensor_file in safetensor_files:
        layer_names, tensors = load_model_states_from_file(Path(safetensor_file))
        all_layer_names.extend(layer_names)
        all_tensors.extend(tensors)
    all_layer_names, all_tensors = map(list, zip(
        *sorted(zip(all_layer_names, all_tensors, strict=False)),
        strict=False,
    ))
    return all_tensors, all_layer_names


class RenameUnpickler(pickle.Unpickler):  # noqa: S301
    """Custom unpickler to rename modules when loading pickled objects."""

    def find_class(self, module: str, name: str) -> Any:  # noqa: ANN401
        """Find and return the class specified by the module and name.

        Parameters
        ----------
        module : str
            The module name.
        name : str
            The class name.

        Returns
        -------
        Any
            The class object.

        """
        renamed_module = module
        if OLD_CODEBASE_NAME in module:
            renamed_module = module.replace(OLD_CODEBASE_NAME, CURRENT_CODEBASE_NAME)

        return super().find_class(renamed_module, name)


def custom_pickle_load(
    file_obj: Any,  # noqa: ANN401
) -> Any:  # noqa: ANN401
    """Read a pickled object from the file object.

    Parameters
    ----------
    file_obj : Any
        The file object to read from.

    Returns
    -------
    Any
        The unpickled object.

    """
    return RenameUnpickler(file_obj).load()


def create_remote_up_down(  # noqa: PLR0913
    bucket_name: str,
    prefix: str,
    run_uuid: str | None,
    num_attempts: int,
    client_config: dict[str, Any],
    *,
    num_concurrent_uploads: int = 1,
    upload_staging_folder: str | None = None,  # Don't touch, it's /tmp by default
    use_procs: bool = True,
) -> RemoteUploaderDownloader:
    """Create the remote uploader/downloader.

    Parameters
    ----------
    bucket_name : str
        The name of the bucket.
    prefix : str
        The prefix of the bucket.
    run_uuid : str | None
        The UUID of the run.
    num_attempts : int
        The number of attempts.
    client_config : dict[str, Any]
        The configuration of the client.
    num_concurrent_uploads : int, optional
        The number of concurrent uploads, by default 1.
    upload_staging_folder : str | None, optional
        The upload staging folder, dont't touch, by default None.
    use_procs : bool, optional
        Whether to use processes, by default True. Don't touch.

    Returns
    -------
    RemoteUploaderDownloader
        The remote uploader/downloader.

    """
    bucket_uri = f"s3://{bucket_name}"
    remote_up_down = RemoteUploaderDownloader(
        bucket_uri=bucket_uri,
        backend_kwargs={
            "bucket": bucket_name,
            "prefix": prefix,  # Don't touch
            "region_name": None,  # Not necessary
            "endpoint_url": None,  # Will be read from env var
            "aws_access_key_id": None,  # Will be read from config file
            "aws_secret_access_key": None,  # Will be read from config file
            "aws_session_token": None,  # Will be automatically generated
            "client_config": client_config,  # And using defaults
            "transfer_config": None,  # Using defaults
        },
        file_path_format_string="{remote_file_name}",  # Don't touch
        num_concurrent_uploads=num_concurrent_uploads,
        upload_staging_folder=upload_staging_folder,  # Don't touch, default: /tmp
        use_procs=use_procs,  # Don't touch
        num_attempts=num_attempts,
    )
    remote_up_down.init(run_name=run_uuid)  # Don't touch
    return remote_up_down
