/*
 * Copyright 2012-2018 Chronicle Map Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package net.openhft.chronicle.hash.impl.stage.iter;

import net.openhft.chronicle.hash.HashEntry;
import net.openhft.chronicle.hash.HashSegmentContext;
import net.openhft.chronicle.hash.impl.CompactOffHeapLinearHashTable;
import net.openhft.chronicle.hash.impl.VanillaChronicleHashHolder;
import net.openhft.chronicle.hash.impl.stage.entry.HashEntryStages;
import net.openhft.chronicle.hash.impl.stage.entry.HashLookupPos;
import net.openhft.chronicle.hash.impl.stage.hash.CheckOnEachPublicOperation;
import net.openhft.sg.StageRef;
import net.openhft.sg.Staged;

import java.util.function.Consumer;
import java.util.function.Predicate;

@Staged
public abstract class HashSegmentIteration<K, E extends HashEntry<K>>
        implements HashEntry<K>, HashSegmentContext<K, E> {

    @StageRef
    public IterationSegmentStages s;
    @StageRef
    public CheckOnEachPublicOperation checkOnEachPublicOperation;
    public boolean entryRemovedOnThisIteration = false;
    public long hashLookupEntry = 0;
    @StageRef
    protected HashLookupPos hlp;
    @StageRef
    HashEntryStages<K> e;
    @StageRef
    VanillaChronicleHashHolder<?> hh;

    public boolean shouldTestEntry() {
        return true;
    }

    public Object entryForIteration() {
        return this;
    }

    public long tierEntriesForIteration() {
        throwExceptionIfClosed();

        return s.tierEntries();
    }

    abstract boolean entryRemovedOnThisIterationInit();

    protected void initEntryRemovedOnThisIteration(boolean entryRemovedOnThisIteration) {
        this.entryRemovedOnThisIteration = entryRemovedOnThisIteration;
    }

    public abstract boolean hashLookupEntryInit();

    public void initHashLookupEntry(long entry) {
        throwExceptionIfClosed();

        hashLookupEntry = entry;
    }

    abstract void closeHashLookupEntry();

    @Override
    public boolean forEachSegmentEntryWhile(Predicate<? super E> predicate) {
        throwExceptionIfClosed();

        checkOnEachPublicOperation.checkOnEachPublicOperation();
        s.innerUpdateLock.lock();
        return innerForEachSegmentEntryWhile(predicate);
    }

    public <T> boolean innerForEachSegmentEntryWhile(Predicate<? super T> predicate) {
        try {
            s.goToLastTier();
            while (true) {
                int currentTier = s.tier;
                long currentTierBaseAddr = s.tierBaseAddr;
                long currentTierIndex = s.tierIndex;
                boolean interrupted = forEachTierEntryWhile(
                        predicate, currentTier, currentTierBaseAddr, currentTierIndex);
                if (interrupted)
                    return false;
                if (currentTier == 0)
                    return true;
                s.prevTier();
            }
        } finally {
            closeHashLookupEntry();
            s.innerReadLock.unlock();
            initEntryRemovedOnThisIteration(false);
        }
    }

    public <T> boolean forEachTierEntryWhile(
            Predicate<? super T> predicate,
            int currentTier, long currentTierBaseAddr, long tierIndex) {
        long leftEntries = tierEntriesForIteration();
        boolean interrupted = false;
        long startPos = 0L;
        CompactOffHeapLinearHashTable hashLookup = hh.h().hashLookup;
        // volatile read not needed because iteration is performed at least under update lock
        while (!hashLookup.empty(hashLookup.readEntry(currentTierBaseAddr, startPos))) {
            startPos = hashLookup.step(startPos);
        }
        hlp.initHashLookupPos(startPos);
        long currentHashLookupPos;
        int steps = 0;
        do {
            // Step from hlp.hashLookupPos, not currentHashLookupPos (with additional initialization
            // of this local variable to startPos outside the loop), because if e.remove() is
            // called in the `predicate`, hlp.hashLookupPos is stepped back in doRemove(), and
            // currentHashLookupPos become invalid
            currentHashLookupPos = hashLookup.step(hlp.hashLookupPos);
            steps++;
            hlp.setHashLookupPos(currentHashLookupPos);
            // volatile read not needed because iteration is performed at least under update lock
            long entry = hashLookup.readEntry(currentTierBaseAddr, currentHashLookupPos);
            initHashLookupEntry(entry);
            if (!hashLookup.empty(entry)) {
                e.readExistingEntry(hashLookup.value(entry));
                if (shouldTestEntry()) {
                    initEntryRemovedOnThisIteration(false);
                    try {
                        if (!predicate.test((T) entryForIteration())) {
                            interrupted = true;
                            break;
                        } else {
                            if (--leftEntries == 0)
                                break;
                        }
                    } finally {
                        hookAfterEachIteration();
                        // if doReplaceValue() -> relocation() -> alloc() -> nextTier()
                        // was called, restore the tier we were iterating over
                        if (s.tier != currentTier) {
                            s.initSegmentTier_WithBaseAddr(
                                    currentTier, currentTierBaseAddr, tierIndex);
                            // To cover shift deleted slot, at the next step forward.
                            // hash lookup entry is relocated to the next chained tier, and the
                            // slot in _current_ tier's hash lookup is shift deleted, see
                            // relocation()
                            currentHashLookupPos = hashLookup.stepBack(currentHashLookupPos);
                            steps--;
                            hlp.initHashLookupPos(currentHashLookupPos);
                        }
                        s.innerWriteLock.unlock();
                        // force entry checksum update (delayedUpdateChecksum depends on keyOffset)
                        e.closeKeyOffset();
                    }
                }
            }
            // the `steps == 0` condition and this variable updates in the loop fix the bug, when
            // shift deletion occurs on the first entry of the tier, and the currentHashLookupPos
            // becomes equal to start pos without making the whole loop, but only visiting a single
            // entry
        } while (currentHashLookupPos != startPos || steps == 0);
        if (!interrupted && leftEntries > 0) {
            throw new IllegalStateException(hh.h().toIdentityString() +
                    ": We went through a tier without interruption, " +
                    "but according to tier counters there should be " + leftEntries +
                    " more entries. Size diverged?");
        }
        return interrupted;
    }

    public void hookAfterEachIteration() {
        throwExceptionIfClosed();

    }

    @Override
    public void forEachSegmentEntry(Consumer<? super E> action) {
        throwExceptionIfClosed();

        forEachSegmentEntryWhile(e -> {
            action.accept(e);
            return true;
        });
    }

    public void checkEntryNotRemovedOnThisIteration() {
        throwExceptionIfClosed();

        if (entryRemovedOnThisIterationInit()) {
            throw new IllegalStateException(
                    hh.h().toIdentityString() + ": Entry was already removed on this iteration");
        }
    }

    @Override
    public void doRemove() {
        throwExceptionIfClosed();

        checkOnEachPublicOperation.checkOnEachPublicOperation();
        s.innerWriteLock.lock();
        try {
            iterationRemove();
        } finally {
            s.innerWriteLock.unlock();
        }
        initEntryRemovedOnThisIteration(true);
    }

    public void iterationRemove() {
        throwExceptionIfClosed();

        // this condition mean -- some other entry taken place of the removed one
        if (hh.h().hashLookup.remove(s.tierBaseAddr, hlp.hashLookupPos) != hlp.hashLookupPos) {
            // if so, should make step back, to compensate step forward on the next iteration,
            // to consume the shifted entry
            hlp.setHashLookupPos(hh.h().hashLookup.stepBack(hlp.hashLookupPos));
        }
        e.innerRemoveEntryExceptHashLookupUpdate();
    }
}
