import logging
import random
from typing import MutableMapping, Optional, Union

import random_name
from fsspec.spec import AbstractFileSystem

from scaffold.data.artifact_manager.base import ArtifactManager

FILESYSTEM = Union[MutableMapping, AbstractFileSystem]

logger = logging.getLogger(__name__)


def validate_or_generate_target_afid(
    artifact_manager: ArtifactManager, afid: Optional[str] = None, prefix: Optional[str] = None
) -> str:
    """Validates a given target afid or generates a new one.

    Takes the given afid and checks its validity, raises a ValueError if not. If the given afid is None, a new afid
    is returned with the optional prefix. Can be used when validating any afid parameters that are meant for
    targets. Target afids refer to the ids for artifacts that will be generated by your entrypoint and therefore
    should not exist yet, in contrast to "input afids" which are ids for artifacts an entrypoint needs as input to
    do something, and therefore _should_ exist in the artifact store.

    Args:
        afid: An artifact id that is supposed to be not used yet on the server, default None.
        prefix: Only used when afid is None. Will be the prefix of the new afid.

    Raises:
        ValueError if an afid was given, but it already exists in the artifact store.

    Returns:
        Either the given afid if its valid, or a new one.
    """
    if afid is None:
        afid = get_new_afid(artifact_manager=artifact_manager, prefix=prefix)
        logger.info(f"Artifact id was not set in config. Got new unique id: {afid}")
    else:
        if artifact_manager.exists(afid):
            raise ValueError(
                f"The given artifact id ({afid}) to store the result would not be unique in this artifact store"
                f". Please specify an artifact id that is unique across ALL runs of the artifact store,"
                f" or leave it empty to get a new id assigned."
            )
    return afid


def get_new_afid(artifact_manager: ArtifactManager, prefix: str = None) -> str:
    """Get a new artifact id.

    The identifiers are triplets of random words, to make them more distinguishable for the human eye, plus the
    timestamp to ensure uniqueness. Optionally, a prefix can be added to the id.

    Args:
        prefix: String to be added to the uid, empty by default

    Returns:
        String in form of prefix__new_id__timestamp
    """
    import time

    random.seed()  # Overwrite the seed set by torch_lightning.utilities.seed.seed_everything during training.
    timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
    new_id = (prefix + "__" if prefix is not None else "") + random_name.generate_name() + "__" + timestamp

    return new_id


def validate_input_afid(artifact_manager: ArtifactManager, afid: str) -> None:
    """Validates input afid, i.e. an afid that is meant to be existing in the artifact store, as it is needed as
    input to a entrypoint.

    Args:
            afid: An artifact id that is supposed to be existing in the artifact store.

    Raises:
            ValueError if an afid was given, but it does not exist in the artifact store.
    """
    if not artifact_manager.exists(afid):
        raise ValueError(f"The given artifact id ({afid}) does not exist in the artifact store.")
