package ru.yandex.qe.dispenser.domain.dao.base_resources;

import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import javax.inject.Inject;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.Lists;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;

import ru.yandex.qe.dispenser.domain.base_resources.BaseResource;

@Component
public class BaseResourceCache {

    private final LoadingCache<Long, Optional<BaseResource>> cache;
    private final BaseResourceDao baseResourceDao;

    @Inject
    public BaseResourceCache(BaseResourceDao baseResourceDao) {
        this.baseResourceDao = baseResourceDao;
        this.cache = CacheBuilder.newBuilder()
                .expireAfterWrite(5, TimeUnit.MINUTES)
                .expireAfterAccess(24, TimeUnit.HOURS)
                .maximumSize(10000)
                .build(new Loader(baseResourceDao));
    }

    public Optional<BaseResource> getById(long id) {
        return CacheRefreshScheduler.wrapException(() -> cache.get(id));
    }

    public Set<BaseResource> getByIds(Collection<? extends Long> ids) {
        return CacheRefreshScheduler.wrapException(() -> cache.getAll(ids).values()
                .stream().flatMap(Optional::stream).collect(Collectors.toSet()));
    }

    public void clear() {
        cache.invalidateAll();
    }

    @Transactional(propagation = Propagation.REQUIRED)
    public void refresh() {
        List<Long> keysToRefresh = cache.asMap().keySet().stream().distinct().collect(Collectors.toList());
        Map<Long, BaseResource> refreshed = Lists.partition(keysToRefresh, 500).stream()
                .flatMap(keys -> baseResourceDao.getByIds(keys).stream())
                .collect(Collectors.toMap(BaseResource::getId, Function.identity()));
        Map<Long, Optional<BaseResource>> refreshedById = new HashMap<>();
        keysToRefresh.forEach(key -> refreshedById.put(key, Optional.ofNullable(refreshed.get(key))));
        cache.putAll(refreshedById);
    }

    private static final class Loader extends CacheLoader<Long, Optional<BaseResource>> {

        private final BaseResourceDao baseResourceDao;

        private Loader(BaseResourceDao baseResourceDao) {
            this.baseResourceDao = baseResourceDao;
        }

        @Override
        public Optional<BaseResource> load(Long key) {
            return baseResourceDao.getById(key);
        }

        @Override
        public Map<Long, Optional<BaseResource>> loadAll(Iterable<? extends Long> keys) {
            List<Long> keySet = StreamSupport.stream(keys.spliterator(), false).distinct()
                    .collect(Collectors.toList());
            Map<Long, BaseResource> knownEntitiesByKey = Lists.partition(keySet, 500).stream()
                    .flatMap(page -> baseResourceDao.getByIds(page).stream())
                    .collect(Collectors.toMap(BaseResource::getId, Function.identity()));
            Map<Long, Optional<BaseResource>> result = new HashMap<>();
            keySet.forEach(key -> result.put(key, Optional.ofNullable(knownEntitiesByKey.get(key))));
            return result;
        }

    }

}
