"""Helper functions defined to help with s3_utils.py."""

from concurrent.futures import ThreadPoolExecutor
import re
import shutil
import time
from logging import DEBUG, INFO
from pathlib import Path
from typing import Any

from botocore.client import BaseClient
from botocore.exceptions import ClientError
from composer.loggers import RemoteUploaderDownloader
from composer.utils import S3ObjectStore, maybe_create_object_store_from_uri, parse_uri
from composer.utils.file_helpers import list_remote_objects
from flwr.common import ConfigRecord, NDArrays, log
from omegaconf import OmegaConf

from mt_dao.conf.base_schema import S3CommConfig
from mt_dao.constants import (
    ENDPOINT_ID,
    FIRST_MOMENTUM,
    MODEL_PARAMETERS,
    SECOND_MOMENTUM,
)
from mt_dao.file_utils import (
    _walk_load_directory,
    chunk_and_save_file,
    create_remote_up_down,
    download_file_from_s3,
    upload_file_to_s3,
)
from mt_dao.utils import ModelStateNames


def key_exists(client: BaseClient, bucket: str, key: str) -> bool:
    """Determine whether a key exists at S3 bucket.

    A helper method to determine whether a key exists at a remote
    S3 bucket.
    Parameters.
    ----------
    client : _type_
        The boto3 client used to search the remote location.
    bucket : str
        The name of the bucket
    key : str
        The name of the remote file.

    Returns
    -------
    bool
        Whether the key exists at the remote bucket.

    Raises
    ------
    ClientError
        Error in accessing the file in the remote s3 bucket.

    """
    try:
        client.head_object(
            Bucket=bucket,
            Key=key,
        )
    except ClientError as e:
        if e.response["Error"]["Code"] == "404":
            return False
        raise
    return True


def list_objects(
    run_uuid_path: str,
) -> tuple[bool, list[str]]:
    """List objects in a local/(remote s3) given path.

    This function attempts to list objects in the specified path. If the path is remote
    S3 path, it uses the `list_remote_objects` function to list the objects. If the path
    is a local path, it lists all files recursively within the directory.

    Parameters
    ----------
    run_uuid_path : str
        The path to list objects from. This can be either a local path or a remote S3
        path.

    Returns
    -------
    tuple[bool, list[str]]
        A tuple where the first element is a boolean indicating whether the path is
        remote (True for remote, False for local), and the second element is a list of
        object paths.

    Example
    -------
    >>> is_remote, objects = list_objects("s3://my-bucket/my-folder")
    >>> print(is_remote)
    True
    >>> print(objects)
    ['s3://my-bucket/my-folder/file1.txt', 's3://my-bucket/my-folder/file2.txt']

    >>> is_remote, objects = list_objects("/local/path/to/folder")
    >>> print(is_remote)
    False
    >>> print(objects)
    ['/local/path/to/folder/file1.txt', '/local/path/to/folder/file2.txt']

    """
    try:
        return True, list_remote_objects(run_uuid_path)
    except ValueError:
        # Local path
        return False, [str(p) for p in Path(run_uuid_path).rglob("*") if p.is_file()]


def delete_remote_object(object_name: str) -> None:
    """Delete an object from an S3 bucket.

    This function deletes an object from an S3 bucket using the provided object name.
    It creates an S3 object store from the object name, parses the URI to extract the
    prefix, and then deletes the object from the object store.

    Parameters
    ----------
    object_name : str
        The name of the object to delete.

    Raises
    ------
        ValueError: If the object name is not a valid URI or if the object cannot be
            deleted.

    """
    # Create an object store from the object name
    object_store: S3ObjectStore | None = maybe_create_object_store_from_uri(object_name)  # type: ignore[assignment,reportAssignmentType]
    if object_store is None:
        msg = f"Invalid object name: {object_name}"
        raise ValueError(msg)
    # Parse the URI to extract the prefix to use as the key to delete the file
    _, _, prefix = parse_uri(object_name)
    # Delete the object from the object store
    object_store.client.delete_object(
        Bucket=object_store.bucket,
        Key=object_store.get_key(prefix),
    )


def delete_object(
    object_path: str,
) -> None:
    """Delete a local/remote object from a specified path.

    This function attempts to delete the specified object from the provided path. If the
    path is a remote S3 path, it uses the `delete_remote_object` function to delete the
    object. If the path is a local path, it deletes the local file using the `unlink`
    method.

    Parameters
    ----------
    object_path : str
        The path to the object to be deleted. This can be either a local path or a
        remote S3 path.

    Example
    -------
    >>> delete_object("s3://my-bucket/my-folder/myfile.txt")
    >>> delete_object("/local/path/to/myfile.txt")

    Notes
    -----
    This function uses the `delete_remote_object` function to delete objects from a
    remote S3 path. For local paths, it uses the `Path.unlink` method to delete local
    files.

    """
    try:
        delete_remote_object(object_path)
    except ValueError:
        # Local path
        Path(object_path).unlink()


def generic_vector_upload_to_s3(  # noqa: PLR0917, PLR0913
    vector: NDArrays,
    layer_names: list[str],
    current_round: int,
    remote_up_down: RemoteUploaderDownloader | None,
    local_checkpoint_path: str,
    model_type: str,
    *,
    parameter_upload: bool = True,
    is_second_momentum: bool = False,
    max_shard_size: str = "5GB",
) -> None:
    """Encapsulate general vector uploadings to remote S3 buckets.

    Parameters
    ----------
    vector : NDArrays
        The vector to upload
    layer_names : list[str]
        The corresponding layer names to the vector.
    current_round : int
        The current round at which the upload is happening.
    remote_up_down : RemoteUploaderDownloader
        RemoteUploaderDownloader object that manages I/O with remote S3 bucket.
    local_checkpoint_path : str
        The local checkpoint directory.
    model_type : str
        The model type being uploaded.
    parameter_upload : bool, optional
        Whether the upload is a parameter upload, by default True
    is_second_momentum : bool, optional
        Whether the momentum term being uploaded is the second moment, by default False
    max_shard_size : str, optional
        The maxmimum shard size of the upload file., by default "5GB"

    """

    if parameter_upload:
        upload_type = MODEL_PARAMETERS
    elif is_second_momentum:
        upload_type = SECOND_MOMENTUM
    else:
        upload_type = FIRST_MOMENTUM
    local_dir = (
        Path(local_checkpoint_path)
        / "server"
        / f"{current_round}"
        / f"{upload_type}"
        / f"{model_type}"
    )
    dump_time = time.perf_counter()
    if not local_dir.exists():
        chunk_and_save_file(
            layer_names=layer_names,
            model_arrays=vector,
            save_directory=local_dir,
            max_shard_size=max_shard_size,
        )
    log(DEBUG, f"Time to dump to disk: {time.perf_counter() - dump_time}")
    if remote_up_down:
        log(
            DEBUG,
            f"Push {model_type} of {upload_type} to S3 Object Store.",
        )
        for file_name in _walk_load_directory(local_dir):
            relative_path = Path(file_name).relative_to(local_dir)
            upload_file_to_s3(
                remote_up_down,
                f"{current_round}/{upload_type}/{model_type}/{relative_path}",
                file_name,
            )


def generic_model_state_upload_to_s3(  # noqa: PLR0913
    model_states: NDArrays,
    layer_names_and_types: tuple[tuple[str, ModelStateNames], ...],
    current_round: int,
    remote_up_down: RemoteUploaderDownloader | None,
    local_checkpoint_path: str,
    *,
    parameter_upload: bool = True,
    is_second_momentum: bool = False,
    max_shard_size: str = "5GB",
) -> None:
    """Generate a generic model state uploader to remote S3 bucket.

    Parameters
    ----------
    model_states : NDArrays
        The full model state to be uploaded.
    layer_names_and_types : tuple[tuple[str, ModelStateNames], ...]
        The corresponding layer names and types to the model state.
    current_round : int
        The current round at which the upload is happening.
    remote_up_down : RemoteUploaderDownloader
        RemoteUploaderDownloader object that manages I/O with remote S3 bucket.
    local_checkpoint_path : str
        The local checkpoint directory.
    parameter_upload : bool, optional
        Whether the upload is a parameter upload, by default True
    is_second_momentum : bool, optional
        Whether the momentum term being uploaded is the second moment, by default False
    max_shard_size : str, optional
        The maxmimum shard size of the upload file., by default "5GB"

    """

    def generate_model_type_state_dict(
        model_states: NDArrays,
        layer_names_and_types: tuple[tuple[str, ModelStateNames], ...],
    ) -> dict[str, tuple[NDArrays, list[str]]]:
        """Generate the state dict for the model per model type.

        Parameters
        ----------
        model_states : NDArrays
            The model states.
        layer_names_and_types : tuple[tuple[str, ModelStateNames], ...]
            The corresponding layer names and types to the model state.

        Returns
        -------
        dict[str, tuple[NDArrays, list[str]]]
            A dictionary containing the model parameters and the corresponding
            names of the layers per model type.

        """
        model_state_dict = {}
        for layer_ndarray, (layer_name, layer_type) in zip(
            model_states,
            layer_names_and_types,
            strict=False,
        ):
            if layer_type.value not in model_state_dict:
                model_state_dict[layer_type.value] = ([layer_ndarray], [layer_name])
            else:
                model_state_dict[layer_type.value][0].append(layer_ndarray)
                model_state_dict[layer_type.value][1].append(layer_name)
        return model_state_dict

    model_state = generate_model_type_state_dict(
        model_states=model_states,
        layer_names_and_types=layer_names_and_types,
    )
    with ThreadPoolExecutor(max_workers=8) as executor:
        upload_futures = []

        log(DEBUG, "Submitting model states to thread pool executor.")
        for state_type, (
            state_layer_ndarrays,
            state_layer_names,
        ) in model_state.items():
            # 2. Submit the function to the executor. This is non-blocking and
            # immediately returns a Future object representing the background job.
            future = executor.submit(
                generic_vector_upload_to_s3,
                vector=state_layer_ndarrays,
                layer_names=state_layer_names,
                current_round=current_round,
                remote_up_down=remote_up_down,
                local_checkpoint_path=local_checkpoint_path,
                model_type=state_type,
                parameter_upload=parameter_upload,
                is_second_momentum=is_second_momentum,
                max_shard_size=max_shard_size,
            )
            upload_futures.append(future)

        log(
            DEBUG,
            f"All {len(upload_futures)} upload tasks have been submitted. The main thread can continue.",
        )


def get_file_from_path(
    input_file_path: str,
    run_uuid: str,
    s3_comm_config: S3CommConfig,
    output_parent_path: str,
    relative_header: str = "",
) -> Path:
    """Retrieve a file from a given path, which can be either a local path or an S3 URI.

    This function interprets the input file path to determine whether it's a local path
    or an S3 URI. If it is an S3 URI, the function downloads the file from the specified
    S3 bucket to a temporary directory. If it is a local path, the function verifies the
    existence of the file. The function returns the local path to the file.

    Parameters
    ----------
    input_file_path : str
        The path to the file to be retrieved. Can be a local file path or an S3 URI.
    run_uuid : str
        The unique identifier for the run, used for S3 operations.
    s3_comm_config : S3CommConfig
        The S3 communication configuration.
    output_parent_path : str
        The output parent path for storing the downloaded file.
    relative_header: str
        The relative header the file should be saved against such as the remote
        saving location structure is preserved.

    Returns
    -------
    Path
        The local file path to the retrieved file.

    Raises
    ------
    ValueError
        If the backend specified in the URI is unknown.

    """
    # Interpret the URI
    backend, bucket_name, remote_file_name = parse_uri(input_file_path)
    local_file_path: Path | None = None
    if backend == "s3":
        log(
            INFO,
            "Downloading model %s from S3 bucket %s",
            remote_file_name,
            bucket_name,
        )
        # Create RemoteUploaderDownloader object
        remote_up_down = create_remote_up_down(
            bucket_name=bucket_name,
            prefix="",
            run_uuid=run_uuid,
            num_attempts=5,
            client_config=OmegaConf.to_container(
                s3_comm_config.backend_kwargs.client_config,
            ),  # type: ignore[reportArgumentType, arg-type]
        )
        if relative_header:
            file_suffix = Path(remote_file_name).relative_to(relative_header)
            local_file_path = Path(output_parent_path) / file_suffix
        else:
            local_file_path = Path(output_parent_path) / (
                "checkpoint" + Path(remote_file_name).suffix
            )
        download_file_from_s3(remote_up_down, remote_file_name, local_file_path)
    elif not backend:
        log(
            INFO,
            "File path %s is local.",
            Path(input_file_path),
        )
        local_file_path = Path(input_file_path)
    else:
        msg = f"Unknown backend: {backend}"
        raise ValueError(msg)
    assert local_file_path is not None, "Local file path is None"
    assert local_file_path.exists(), f"Local file path {local_file_path} does not exist"
    return local_file_path


def get_num_batches_from_checkpoint_name(checkpoint_name: str) -> int:
    """Extract the number of batches from the checkpoint name.

    The checkpoint name is expected to be in the format:
    ep{n_epochs}-ba{n_batches}-rank{rank}.pt

    Parameters
    ----------
    checkpoint_name : str
        The name of the checkpoint file.

    Returns
    -------
        int: The number of batches extracted from the checkpoint name.

    Raises
    ------
        ValueError: If the checkpoint name does not match the expected format.

    """
    match = re.search(r"-ba(\d+)-", checkpoint_name)
    if match:
        return int(match.group(1))
    msg = f"Invalid checkpoint name format: {checkpoint_name}"
    raise ValueError(msg)


def extract_s3_comm_config_from_configrecord(
    s3_comm_config: ConfigRecord,
) -> tuple[str, str, str]:
    """Extract S3 communication configuration details from a ConfigsRecord object.

    This function parses a ConfigsRecord object containing S3 communication
    configuration and extracts essential information required for S3 operations.
    Specifically, it retrieves the `endpoint_id`, `file_name`, and `current_round`
    from the ConfigsRecord. These values are crucial for identifying the correct S3
    bucket and path, and for versioning or round-specific operations.

    Parameters
    ----------
    s3_comm_config : ConfigRecord
        A ConfigRecord object containing the S3 communication configuration. Expected
        to have keys for `endpoint_id`, `file_name`, and `current_round`.

    Returns
    -------
    tuple[str, str, str]
        A tuple containing `endpoint_id`, `file_name`, and `current_round` as strings.

    Raises
    ------
    ValueError
        If any of the required keys (`endpoint_id`, `file_name`, or `current_round`) are
        missing from the ConfigsRecord.

    Notes
    -----
    The function ensures that all returned values are strings, even if they are provided
    as different types in the ConfigsRecord. This standardization facilitates their use
    in S3 operations without further type checking or conversion.

    """
    # Extract endpoint id from the content of the message
    endpoint_id: Any
    if ENDPOINT_ID in s3_comm_config:
        endpoint_id = str(s3_comm_config[ENDPOINT_ID])
    else:
        msg = f"{ENDPOINT_ID} is not present in the message"
        raise ValueError(msg)
    file_name: Any
    if "file_name" in s3_comm_config:
        file_name = str(s3_comm_config["file_name"])
    else:
        msg = "file_name is not present in the message"
        raise ValueError(msg)
    folder_name: Any
    if "folder_name" in s3_comm_config:
        folder_name = str(s3_comm_config["folder_name"])
    else:
        msg = "folder_name is not present in the message"
        raise ValueError(msg)
    return endpoint_id, file_name, folder_name


def obtain_sorted_runs(
    run_uuid_path: str,
    state_keys: tuple[str, ...],
    *,
    catch_all: bool = True,
) -> list[int]:
    """Obtain the sorted runs from the server path.

    This function lists the objects in the specified run UUID path, extracts unique run
    numbers from the paths under `{run_uuid_path}/server/`, and filters them based on
    the provided state keys. It returns the sorted list of valid run numbers.

    Parameters
    ----------
    run_uuid_path : str
        The path to the run UUID root.
    state_keys : tuple[str, ...]
        The state keys to check in the paths. Keys are intended to be the prefixes of
        any file name. For example, if the keys are ("state", "model"), then the
        function will return any path that starts with "state" and "model".
    catch_all : bool
        Whether to force that all state_keys should match to give a valid run number.

    Returns
    -------
    list[int]
        The sorted list of valid run numbers.

    Example
    -------
    >>> run_uuid_path = "s3://my-bucket/my-folder"
    >>> state_keys = ("state", "model")
    >>> sorted_runs = obtain_sorted_runs(run_uuid_path, state_keys)
    >>> print(sorted_runs)
    [1, 2, 3]

    Notes
    -----
    This function uses the `list_objects` function to list objects in the given path.
    It extracts unique run numbers from the paths under `{run_uuid_path}/server/` and
    filters them based on the provided state keys.

    """
    _, remote_objects = list_objects(run_uuid_path)
    # Extract unique run numbers
    run_numbers = {
        int(reg.group(1))
        for path in remote_objects
        if (reg := re.search(r"server/(\d+)/.*$", path)) is not None
    }

    valid_runs = set()
    catch_all_fn = all if catch_all else any
    for run in run_numbers:
        # Filter paths for the current run
        run_paths = [path for path in remote_objects if f"server/{run}/" in path]

        # Check if all state_keys are present in the paths for this run
        if catch_all_fn(
            any(state_key in path for path in run_paths) for state_key in state_keys
        ):
            valid_runs.add(run)

    return sorted(valid_runs)


def delete_clients_checkpoints(run_uuid_path: str, end_idx: int | None = -1) -> None:
    """Delete client checkpoints from a specified path. Can be either a local or remote.

    This function deletes the specified client checkpoints from the provided run UUID
    path. It lists all the objects in the path, extracts unique client IDs, and removes
    the corresponding checkpoints for each client based on the `end_idx` parameter. The
    function supports both local and remote S3 paths.

    Parameters
    ----------
    run_uuid_path : str
        The path to the run UUID, which can be either a local path or a remote S3 path.
    end_idx : int, optional
        The index up to which checkpoints should be deleted. Defaults to -1, which means
        all checkpoints except for the last.

    Example
    -------
    >>> delete_clients_checkpoints("s3://my-bucket/my-folder", end_idx=5)
    >>> delete_clients_checkpoints("/local/path/to/folder", end_idx=5)

    Notes
    -----
    This function uses the `list_objects` function to list objects in the given path.
    For remote S3 paths, it uses the `delete_object` function to delete objects from the
    S3 bucket. For local paths, it uses the `delete_object` function to delete local
    files.

    """
    # List all the remote objects in the run UUID path
    _, remote_objects = list_objects(run_uuid_path)
    # Extract unique client IDs from the remote objects
    unique_client_ids = {
        int(reg.group(1))
        for path in remote_objects
        if (reg := re.search(r"client_(\d+)/.*$", path)) is not None
    }
    for client_id in unique_client_ids:
        # List all the remote objects for the client
        is_remote, client_remote_objects = list_objects(
            f"{run_uuid_path}/client_{client_id}/",
        )
        # Remove symlinks from the list of files
        client_remote_objects = [
            cro for cro in client_remote_objects if not cro.endswith(".symlink")
        ]
        # Sort by number of batches the client trained on
        sorted_client_objects = sorted(
            client_remote_objects,
            key=get_num_batches_from_checkpoint_name,
        )
        # Delete only the last `end_idx` checkpoints
        objects_to_remove = sorted_client_objects[:end_idx]
        for object_to_remove in objects_to_remove:
            if is_remote:
                # Parse the URI to extract the backend and bucket name
                backend, bucket_name, _prefix = parse_uri(run_uuid_path)
                delete_object(f"{backend}://{bucket_name}/{object_to_remove}")
            else:
                delete_object(object_to_remove)


def delete_rounds(
    run_uuid_path: str,
    state_keys: tuple[str, ...],
    end_idx: int | None = -1,
    *,
    catch_all: bool = True,
) -> None:
    """List objects in a local/(remote s3) given path.

    This function attempts to list objects in the specified path. If the path is remote
    S3 path, it uses the `list_remote_objects` function to list the objects. If the path
    is a local path, it lists all files recursively within the directory.

    Parameters
    ----------
    run_uuid_path : str
        The path to list objects from. This can be either a local path or a remote S3
        path.
    state_keys : tuple[str, ...]
        The state keys to check in the paths. Keys are intended to be the prefixes of
        any file name. For example, if the keys are ("state", "model"), then the
        function will return any path that starts with "state" and "model".
    catch_all : bool,
        Whether to ensure that all the state_keys that are passed in to delete
        the run or not.
    end_idx : int, optional
        The index up to which rounds should be deleted. Defaults to -1, which means all

    Example
    -------
    >>> is_remote, objects = list_objects("s3://my-bucket/my-folder")
    >>> print(is_remote)
    True
    >>> print(objects)
    ['s3://my-bucket/my-folder/file1.txt', 's3://my-bucket/my-folder/file2.txt']

    >>> is_remote, objects = list_objects("/local/path/to/folder")
    >>> print(is_remote)
    False
    >>> print(objects)
    ['/local/path/to/folder/file1.txt', '/local/path/to/folder/file2.txt']

    Notes
    -----
    This function uses the `list_remote_objects` function to list objects in a remote
    S3 path. For local paths, it uses the `Path.rglob` method to recursively list all
    files in the directory.

    """
    # List all the federated rounds in the run UUID path
    sorted_rounds = obtain_sorted_runs(
        run_uuid_path=run_uuid_path,
        state_keys=state_keys,
        catch_all=catch_all,
    )
    # Delete only the last `end_idx` rounds
    rounds_to_delete = sorted_rounds[:end_idx]
    for round_to_delete in rounds_to_delete:
        # List all the remote objects for the server at the round specified
        directory = f"{run_uuid_path}/server/{round_to_delete}/"
        is_remote, objects_to_remove = list_objects(directory)
        if is_remote:
            # Remove the objects found one by one since it is s3.
            for object_to_remove in objects_to_remove:
                # Parse the URI to extract the backend and bucket name
                backend, bucket_name, _ = parse_uri(run_uuid_path)
                delete_object(f"{backend}://{bucket_name}/{object_to_remove}")
        else:
            # Deleting locally can remove the entire tree for us.
            shutil.rmtree(directory)
