import logging
import os
import re
import uuid
from pathlib import Path

from ludwig.constants import CHECKSUM, META, TEST, TRAINING, VALIDATION
from ludwig.data.cache.util import calculate_checksum
from ludwig.utils import data_utils
from ludwig.utils.fs_utils import delete, path_exists

logger = logging.getLogger(__name__)


def alphanum(v):
    """Filters a string to only its alphanumeric characters."""
    return re.sub(r"\W+", "", v)


class DatasetCache:
    def __init__(self, config, checksum, cache_map, dataset_manager):
        self.config = config
        self.checksum = checksum
        self.cache_map = cache_map
        self.dataset_manager = dataset_manager

    def get(self):
        training_set_metadata_fp = self.cache_map[META]
        if not path_exists(training_set_metadata_fp):
            return None

        cache_training_set_metadata = data_utils.load_json(training_set_metadata_fp)

        cached_training_set = self.cache_map[TRAINING] if path_exists(self.cache_map[TRAINING]) else None

        cached_test_set = self.cache_map[TEST] if path_exists(self.cache_map[TEST]) else None

        cached_validation_set = self.cache_map[VALIDATION] if path_exists(self.cache_map[VALIDATION]) else None

        valid = self.checksum == cache_training_set_metadata.get(CHECKSUM) and cached_training_set is not None

        return valid, cache_training_set_metadata, cached_training_set, cached_test_set, cached_validation_set

    def put(self, training_set, test_set, validation_set, training_set_metadata):
        logger.info("Writing preprocessed training set cache")
        training_set = self.dataset_manager.save(
            self.cache_map[TRAINING],
            training_set,
            self.config,
            training_set_metadata,
            TRAINING,
        )

        if test_set is not None:
            logger.info("Writing preprocessed test set cache")
            test_set = self.dataset_manager.save(
                self.cache_map[TEST],
                test_set,
                self.config,
                training_set_metadata,
                TEST,
            )

        if validation_set is not None:
            logger.info("Writing preprocessed validation set cache")
            validation_set = self.dataset_manager.save(
                self.cache_map[VALIDATION],
                validation_set,
                self.config,
                training_set_metadata,
                VALIDATION,
            )

        logger.info("Writing train set metadata")
        data_utils.save_json(self.cache_map[META], training_set_metadata)

        return training_set, test_set, validation_set, training_set_metadata

    def delete(self):
        for fname in self.cache_map.values():
            if path_exists(fname):
                delete(fname)


class CacheManager:
    def __init__(self, dataset_manager, cache_dir=None):
        self._dataset_manager = dataset_manager
        self._cache_dir = cache_dir

    def get_dataset_cache(self, config, dataset=None, training_set=None, test_set=None, validation_set=None):
        if dataset is not None:
            key = self.get_cache_key(dataset, config)
            cache_map = {
                META: self.get_cache_path(dataset, key, META, "json"),
                TRAINING: self.get_cache_path(dataset, key, TRAINING),
                TEST: self.get_cache_path(dataset, key, TEST),
                VALIDATION: self.get_cache_path(dataset, key, VALIDATION),
            }
            return DatasetCache(config, key, cache_map, self._dataset_manager)
        else:
            key = self.get_cache_key(training_set, config)
            cache_map = {
                META: self.get_cache_path(training_set, key, META, "json"),
                TRAINING: self.get_cache_path(training_set, key, TRAINING),
                TEST: self.get_cache_path(test_set, key, TEST),
                VALIDATION: self.get_cache_path(validation_set, key, VALIDATION),
            }
            return DatasetCache(config, key, cache_map, self._dataset_manager)

    def get_cache_key(self, dataset, config):
        if not isinstance(dataset, str):
            # TODO(travis): could try hashing the in-memory dataset, but this is tricky for Dask
            return str(uuid.uuid1())
        return calculate_checksum(dataset, config)

    def get_cache_path(self, dataset, key, tag, ext=None):
        if not isinstance(dataset, str):
            dataset = None

        if self._cache_dir is None and dataset is not None:
            # Use the input dataset filename (minus the extension) as the cache path
            stem = Path(dataset).stem
        else:
            # To avoid collisions across different directories, we use the unique checksum
            # as the cache path
            stem = alphanum(key)

        ext = ext or self.data_format
        cache_fname = f"{stem}.{tag}.{ext}"
        return os.path.join(self.get_cache_directory(dataset), cache_fname)

    def get_cache_directory(self, input_fname):
        if self._cache_dir is None:
            if input_fname is not None:
                return os.path.dirname(input_fname)
            return "."
        return self._cache_dir

    def can_cache(self, skip_save_processed_input):
        return self._dataset_manager.can_cache(skip_save_processed_input)

    @property
    def data_format(self):
        return self._dataset_manager.data_format

[(12, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (17, 're.sub', 're.sub', False, 'import re\n'), (32, 'ludwig.utils.data_utils.load_json', 'data_utils.load_json', False, 'from ludwig.utils import data_utils\n'), (75, 'ludwig.utils.data_utils.save_json', 'data_utils.save_json', False, 'from ludwig.utils import data_utils\n'), (114, 'ludwig.data.cache.util.calculate_checksum', 'calculate_checksum', False, 'from ludwig.data.cache.util import calculate_checksum\n'), (29, 'ludwig.utils.fs_utils.path_exists', 'path_exists', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (34, 'ludwig.utils.fs_utils.path_exists', 'path_exists', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (36, 'ludwig.utils.fs_utils.path_exists', 'path_exists', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (38, 'ludwig.utils.fs_utils.path_exists', 'path_exists', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (81, 'ludwig.utils.fs_utils.path_exists', 'path_exists', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (82, 'ludwig.utils.fs_utils.delete', 'delete', False, 'from ludwig.utils.fs_utils import delete, path_exists\n'), (113, 'uuid.uuid1', 'uuid.uuid1', False, 'import uuid\n'), (122, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (135, 'os.path.dirname', 'os.path.dirname', False, 'import os\n')]