package ru.yandex.webmaster3.storage.host;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.mutable.MutableLong;
import org.apache.commons.lang3.mutable.MutableObject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import ru.yandex.webmaster3.core.WebmasterException;
import ru.yandex.webmaster3.core.data.WebmasterHostId;
import ru.yandex.webmaster3.core.util.trie.*;
import ru.yandex.webmaster3.storage.user.service.UserHostsService;
import ru.yandex.webmaster3.storage.util.ydb.exception.WebmasterYdbException;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Consumer;

/**
 * @author leonidrom
 */
@Service
@Slf4j
@RequiredArgsConstructor(onConstructor_ = @Autowired)
public class AllVerifiedHostsCacheService {
    private static final long CACHE_REFRESH_TIME_MS = TimeUnit.HOURS.toMillis(1);
    private static final long TRIE_BATCH_SIZE = 200000;

    private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    private final UserHostsService userHostsService;
    private final Lock updateLock = new ReentrantLock();
    private volatile CompactTrie<Byte> trie;

    public void init() {
        executor.scheduleAtFixedRate(
                () -> doResolveTrie(true),
                0,
                CACHE_REFRESH_TIME_MS,
                TimeUnit.MILLISECONDS
        );
    }

    public void destroy() {
        executor.shutdownNow();
    }

    public void foreachHost(Consumer<WebmasterHostId> consumer) {
        resolveTrie();
        trie.foreach(null, (key, keyLen) -> {
            consumer.accept(HOST_CONVERTER.fromRawData(key, keyLen));
        });
    }

    public boolean contains(WebmasterHostId hostId) {
        resolveTrie();
        return trie.contains(HOST_CONVERTER.toRawData(hostId));
    }

    private void resolveTrie() {
        if (trie != null) {
            return;
        }

        doResolveTrie(false);
    }

    private void doResolveTrie(boolean isForceResolve) {
        updateLock.lock();
        try {
            if (trie != null && !isForceResolve) {
                return;
            }

            log.info("Hosts trie is stale, update...");
            long updateStart = System.currentTimeMillis();
            // специально не освобождаем память перед вызовом buildTrie(),
            // что бы после первоначальной инициализации trie всегда был доступен
            trie = buildTrie();
            log.info("Trie loaded: {} bytes, load time {}",
                    trie.getDataSizeBytes(),
                    System.currentTimeMillis() - updateStart);
        } finally {
            updateLock.unlock();
        }
    }

    private CompactTrie<Byte> buildTrie() throws WebmasterException {
        MutableObject<CompactTrie<Byte>> trie = new MutableObject<>();
        MutableObject<CompactTrieBuilder<Byte>> trieBuilder = new MutableObject<>(new CompactTrieBuilder<>());
        MutableLong hosts = new MutableLong();
        MutableLong sumLen = new MutableLong();
        long start = System.currentTimeMillis();
        try {
            userHostsService.forEachUserHostPair(pair -> {
                var hostId = pair.getRight();
                if (hosts.incrementAndGet() % TRIE_BATCH_SIZE == 0) {
                    log.info("Loaded {} hosts", hosts.getValue());
                    trie.setValue(mergeTrie(trie.getValue(), trieBuilder.getValue()));
                    trieBuilder.setValue(new CompactTrieBuilder<>());
                }

                byte[] key = HOST_CONVERTER.toRawData(hostId);
                sumLen.add(key.length);
                trieBuilder.getValue().insert(key, (byte)1);
            });

            trie.setValue(mergeTrie(trie.getValue(), trieBuilder.getValue()));
            long loaded = System.currentTimeMillis();
            log.info("Loaded in {} seconds, hostnames sum length {}", (loaded - start) / 1000, sumLen.getValue());
        } catch (WebmasterYdbException e) {
            throw new RuntimeException();
        }

        return trie.getValue();
    }

    private CompactTrie<Byte> mergeTrie(CompactTrie<Byte> trie, CompactTrieBuilder<Byte> trieBuilder) {
        long startMerge = System.currentTimeMillis();
        CompactTrie<Byte> newTrie;
        if (trie == null) {
            newTrie = CompactTrie.fromNode(trieBuilder, BYTE_DRIVER);
        } else {
            var merged = MergedNodeImpl.createMerged(trie.getRootNode(), trieBuilder.getRoot(), (a, b) -> b);
            newTrie = CompactTrie.fromNode(merged, BYTE_DRIVER);
        }

        newTrie = newTrie.shrink();
        log.info("Merged in {}ms", System.currentTimeMillis() - startMerge);

        return newTrie;
    }

    private final static TrieValueCodec<Byte> BYTE_DRIVER = new TrieValueCodec<Byte>() {
        @Override
        public int write(Byte value, TrieBufferWriter buffer, int offset) {
            return buffer.writeBytes(offset, new byte[] {value}, 0, 1);
        }

        @Override
        public Byte read(TrieBufferReader data) {
            return data.readByte();
        }

        @Override
        public void skip(TrieBufferReader data) {
            data.skip(1);
        }
    };

    private static IDataConverter<WebmasterHostId> HOST_CONVERTER = new IDataConverter<>() {
        @Override
        public byte[] toRawData(WebmasterHostId hostId) {
            return EncodeHostUtil.hostToByteArray(hostId);
        }

        @Override
        public WebmasterHostId fromRawData(byte[] data, int len) {
            return EncodeHostUtil.fromByteArray(data, len);
        }
    };
}
