package com.atlassian.stash.internal.i18n.cache;

import com.google.common.base.Predicates;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;

import javax.annotation.Nonnull;
import javax.annotation.concurrent.ThreadSafe;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.function.Supplier;

import static com.atlassian.stash.internal.i18n.cache.CacheUtils.filterKey;
import static com.google.common.collect.ImmutableMap.copyOf;

/**
 * An <i>immutable</i> copy of messages indexed by some key ({@code K}).
 * <p>
 * When the key is deleted all indexed messages are removed as well.
 */
@ThreadSafe
public class MessageMapIndex<K, V> {

    private final Map<K, LocaleMessages<V>> index;
    private final MessageMap<V> messages;
    private final Set<Locale> unresolvedLocales;

    MessageMapIndex() {
        this(ImmutableMap.of(), MessageMap.empty());
    }

    MessageMapIndex(Map<K, LocaleMessages<V>> index, MessageMap<V> messages) {
        this.index = index;
        this.messages = messages;
        this.unresolvedLocales = getLocales(index);
    }

    @Nonnull
    public MessageMapIndex<K, V> add(@Nonnull K key, @Nonnull UnresolvedLocales<V> value) {
        if (value.map.isEmpty()) {
            return this;
        }

        // Can't use ImmutableMap.Builder = doesn't handle duplicates :(
        Map<K, LocaleMessages<V>> newIndex = Maps.newHashMap(index);
        newIndex.put(key, new LocaleMessages<>(ImmutableSet.of(), value));
        return new MessageMapIndex<>(copyOf(newIndex), messages);
    }

    public boolean containsUnresolved(Iterable<Locale> locales) {
        return Iterables.any(locales, Predicates.in(unresolvedLocales));
    }

    @Nonnull
    public MessageLocaleMap<V> getMessages(Iterable<Locale> locales) {
        return new MessageLocaleMap<>(messages, locales);
    }

    @Nonnull
    public MessageMapIndex<K, V> loadAll(Iterable<Locale> locales) {
        MessageMapIndex<K, V> index = this;
        for (Locale locale : locales) {
            index = index.load(locale);
        }
        return index;
    }

    @Nonnull
    public MessageMapIndex<K, V> remove(@Nonnull K key) {
        LocaleMessages<V> keys = index.get(key);
        if (keys == null) {
            return this;
        }
        return new MessageMapIndex<>(filterKey(index, key), messages.remove(keys.keys));
    }

    /**
     * For testing only - normally we only want to load at the top-level, but this is convenient.
     */
    @Nonnull
    MessageLocaleMap<V> loadAndGetMessages(Iterable<Locale> locales) {
        return loadAll(locales).getMessages(locales);
    }

    /**
     * Loads any unloaded messages for a given locale and return a new, updated index.
     */
    @Nonnull
    private MessageMapIndex<K, V> load(@Nonnull Locale locale) {
        // This is the hidden cost of caching per-locale. Is there something even faster?
        if (!unresolvedLocales.contains(locale)) {
            // We haven't got anything to load, just ignore (which is most of the time)
            return this;
        }
        Map<K, LocaleMessages<V>> newIndex = Maps.newHashMap(index);
        MessageMap.Builder<V> newMessages = messages.builder();
        for (Map.Entry<K, LocaleMessages<V>> e : newIndex.entrySet()) {
            Supplier<MessageMap<V>> messageSupplier = e.getValue().locales.map.get(locale);
            // If we have loaded all the locales from this entry then we can skip because we have copied all the results above
            if (messageSupplier != null) {
                MessageMap<V> messages = messageSupplier.get();
                newMessages.putAll(messages);
                ImmutableSet<MessageKey> keys = !messages.keySet().isEmpty()
                        ? ImmutableSet.<MessageKey>builder().addAll(e.getValue().keys).addAll(messages.keySet()).build()
                        : e.getValue().keys;
                e.setValue(new LocaleMessages<>(keys, new UnresolvedLocales<>(filterKey(e.getValue().locales.map, locale))));
            }
        }
        return new MessageMapIndex<>(ImmutableMap.copyOf(newIndex), newMessages.build());
    }

    private static <K, V> ImmutableSet<Locale> getLocales(Map<K, LocaleMessages<V>> index) {
        ImmutableSet.Builder<Locale> newLocales = ImmutableSet.builder();
        for (LocaleMessages<V> value : index.values()) {
            newLocales.addAll(value.locales.map.keySet());
        }
        return newLocales.build();
    }

    private static class LocaleMessages<V> {

        private ImmutableSet<MessageKey> keys;
        private UnresolvedLocales<V> locales;

        public LocaleMessages(ImmutableSet<MessageKey> keys, UnresolvedLocales<V> locales) {
            this.keys = keys;
            this.locales = locales;
        }
    }
}
