import torch
import os
import sys
import os.path
import warnings
from sybil.datasets.utils import get_scaled_annotation_mask, IMG_PAD_TOKEN
from sybil.augmentations import ComposeAug
import numpy as np
from abc import ABCMeta, abstractmethod
import hashlib


CACHED_FILES_EXT = ".png"
DEFAULT_CACHE_DIR = "default/"

CORUPTED_FILE_ERR = (
    "WARNING! Error processing file from cache - removed file from cache. Error: {}"
)


def md5(key):
    """
    returns a hashed with md5 string of the key
    """
    return hashlib.md5(key.encode()).hexdigest()


def split_augmentations_by_cache(augmentations):
    """
    Given a list of augmentations, returns a list of tuples. Each tuple
    contains a caching key of the augmentations up to the spiltting point,
    and a list of augmentations that should be applied afterwards.

    split_augmentations will contain all possible splits by cachable augmentations,
    ordered from latest possible one to the former ones.
    The last tuple will have all augmentations.

    Note - splitting will be done for indexes that all augmentations up to them are
    cachable.
    """
    # list of (cache key, post augmentations)
    split_augmentations = []
    split_augmentations.append((DEFAULT_CACHE_DIR, augmentations))
    all_prev_cachable = True
    key = ""
    for ind, trans in enumerate(augmentations):

        # check trans.cachable() first separately to save run time
        if not all_prev_cachable or not trans.cachable():
            all_prev_cachable = False
        else:
            key += trans.caching_keys()
            post_augmentations = (
                augmentations[ind + 1 :] if ind < len(augmentations) else []
            )
            split_augmentations.append((key, post_augmentations))

    return list(reversed(split_augmentations))


def apply_augmentations_and_cache(
    loaded_input, sample, img_path, augmentations, cache, base_key=""
):
    """
    Loads the loaded input by its absolute path and apply the augmentations one
    by one (similar to what the composed one is doing).  All first cachable
    transformer's output is cached (until reaching a non cachable one).
    """
    all_prev_cachable = True
    key = base_key
    for ind, trans in enumerate(augmentations):
        loaded_input = trans(loaded_input, sample)
        if not all_prev_cachable or not trans.cachable():
            all_prev_cachable = False
        else:
            key += trans.caching_keys()
            cache.add(img_path, key, loaded_input["input"])

    return loaded_input


class cache:
    def __init__(self, path, extension=CACHED_FILES_EXT):
        if not os.path.exists(path):
            os.makedirs(path)

        self.cache_dir = path
        self.files_extension = extension
        if ".npy" != extension:
            self.files_extension += ".npy"

    def _file_dir(self, attr_key, par_dir):
        return os.path.join(self.cache_dir, attr_key, par_dir)

    def _file_path(self, attr_key, par_dir, hashed_key):
        return os.path.join(
            self.cache_dir, attr_key, par_dir, hashed_key + self.files_extension
        )

    def _parent_dir(self, path):
        return os.path.basename(os.path.dirname(path))

    def exists(self, image_path, attr_key):
        hashed_key = md5(image_path)
        par_dir = self._parent_dir(image_path)
        return os.path.isfile(self._file_path(attr_key, par_dir, hashed_key))

    def get(self, image_path, attr_key):
        hashed_key = md5(image_path)
        par_dir = self._parent_dir(image_path)
        return np.load(self._file_path(attr_key, par_dir, hashed_key))

    def add(self, image_path, attr_key, image):
        hashed_key = md5(image_path)
        par_dir = self._parent_dir(image_path)
        file_dir = self._file_dir(attr_key, par_dir)
        if not os.path.exists(file_dir):
            os.makedirs(file_dir)
        np.save(self._file_path(attr_key, par_dir, hashed_key), image)

    def rem(self, image_path, attr_key):
        hashed_key = md5(image_path)
        par_dir = self._parent_dir(image_path)
        try:
            os.remove(self._file_path(attr_key, par_dir, hashed_key))
        # Don't raise error if file not exists.
        except OSError:
            pass


class abstract_loader:
    __metaclass__ = ABCMeta

    def __init__(self, cache_path, augmentations, args, apply_augmentations=True):
        self.pad_token = IMG_PAD_TOKEN
        self.augmentations = augmentations
        self.args = args
        self.apply_augmentations = apply_augmentations
        if cache_path is not None:
            self.use_cache = True
            self.cache = cache(cache_path, self.cached_extension)
            self.split_augmentations = split_augmentations_by_cache(augmentations)
        else:
            self.use_cache = False
            self.composed_all_augmentations = ComposeAug(augmentations)

    @abstractmethod
    def load_input(self, path):
        pass

    @property
    @abstractmethod
    def cached_extension(self):
        pass

    def configure_path(self, path, sample=None):
        return path

    def get_image(self, path, sample=None):
        """
        Returns a transformed image by its absolute path.
        If cache is used - transformed image will be loaded if available,
        and saved to cache if not.
        """
        input_dict = {}
        input_path = self.configure_path(path, sample)

        if input_path == self.pad_token:
            return self.load_input(input_path)

        if not self.use_cache:
            input_dict = self.load_input(input_path)
            # hidden loaders typically do not use augmentation
            if self.apply_augmentations:
                input_dict = self.composed_all_augmentations(input_dict, sample)
            return input_dict

        if self.args.use_annotations:
            input_dict["mask"] = get_scaled_annotation_mask(
                input_dict["annotations"], self.args
            )

        for key, post_augmentations in self.split_augmentations:
            base_key = (
                DEFAULT_CACHE_DIR + key
                if key != DEFAULT_CACHE_DIR
                else DEFAULT_CACHE_DIR
            )
            if self.cache.exists(input_path, base_key):
                try:
                    input_dict["input"] = self.cache.get(input_path, base_key)
                    if self.apply_augmentations:
                        input_dict = apply_augmentations_and_cache(
                            input_dict,
                            sample,
                            input_path,
                            post_augmentations,
                            self.cache,
                            base_key=base_key,
                        )
                    return input_dict
                except Exception as e:
                    print(e)
                    hashed_key = md5(input_path)
                    par_dir = self.cache._parent_dir(input_path)
                    corrupted_file = self.cache._file_path(key, par_dir, hashed_key)
                    warnings.warn(CORUPTED_FILE_ERR.format(sys.exc_info()[0]))
                    self.cache.rem(input_path, key)
        all_augmentations = self.split_augmentations[-1][1]
        input_dict = self.load_input(input_path)
        if self.apply_augmentations:
            input_dict = apply_augmentations_and_cache(
                input_dict,
                sample,
                input_path,
                all_augmentations,
                self.cache,
                base_key=key,
            )

        return input_dict