/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 accord.coordinate.tracking;

import accord.api.TopologySorter;
import accord.api.Tracing;
import accord.local.Node.Id;
import accord.topology.Shard;
import accord.topology.Topologies;
import accord.topology.Topology;

import accord.utils.Invariants;
import accord.utils.SortedArrays.SortedArrayList;

import com.google.common.annotations.VisibleForTesting;

import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.function.IntFunction;
import java.util.function.Predicate;

import javax.annotation.Nullable;

import static accord.coordinate.tracking.AbstractTracker.ShardOutcomes.NoChange;
import static accord.coordinate.tracking.AbstractTracker.ShardOutcomes.SendMore;
import static accord.coordinate.tracking.RequestStatus.Failed;
import static accord.utils.ArrayBuffers.cachedAny;

public abstract class AbstractTracker<ST extends ShardTracker>
{
    public enum ShardOutcomes implements ShardOutcome<AbstractTracker<?>>
    {
        Fail(Failed),
        Success(RequestStatus.Success),
        SendMore(null),
        NoChange(RequestStatus.NoChange);

        public final RequestStatus result;

        ShardOutcomes(RequestStatus result) {
            this.result = result;
        }

        private static ShardOutcomes min(ShardOutcomes a, ShardOutcomes b)
        {
            return a.compareTo(b) <= 0 ? a : b;
        }

        @Override
        public ShardOutcomes apply(AbstractTracker<?> tracker, int shardIndex)
        {
            if (this == Success)
                return --tracker.waitingOnShards == 0 ? Success : NoChange;
            return this;
        }
    }

    public interface ShardFactory<P, ST extends ShardTracker>
    {
        ST apply(P param, int epochIndex, Shard shard);
    }

    protected final Topologies topologies;
    protected final ST[] trackers;
    protected final int maxShardsPerEpoch;
    protected int waitingOnShards;

    public AbstractTracker(Topologies topologies, IntFunction<ST[]> arrayFactory, Function<Shard, ST> trackerFactory)
    {
        this(topologies, arrayFactory, null, (i1, i2, shard) -> trackerFactory.apply(shard));
    }

    public <P> AbstractTracker(Topologies topologies, IntFunction<ST[]> arrayFactory, P param, ShardFactory<P, ST> trackerFactory)
    {
        Invariants.requireArgument(topologies.totalShards() > 0);
        int topologyCount = topologies.size();
        int maxShardsPerEpoch = topologies.get(0).size();
        int shardCount = maxShardsPerEpoch;
        for (int i = 1 ; i < topologyCount ; ++i)
        {
            int size = topologies.get(i).size();
            maxShardsPerEpoch = Math.max(maxShardsPerEpoch, size);
            shardCount += size;
        }
        this.topologies = topologies;
        this.trackers = arrayFactory.apply(topologyCount * maxShardsPerEpoch);
        for (int i = 0 ; i < topologyCount ; ++i)
        {
            Topology topology = topologies.get(i);
            int size = topology.size();
            for (int j = 0; j < size; ++j)
                trackers[i * maxShardsPerEpoch + j] = trackerFactory.apply(param, i, topology.get(j));
        }
        this.maxShardsPerEpoch = maxShardsPerEpoch;
        this.waitingOnShards = shardCount;
    }

    public abstract RequestStatus recordFailure(Id from);
    // record failure before starting up; can skip any in-flight validation, and perhaps be cheaper
    public RequestStatus prerecordFailure(Id from) { return recordFailure(from); }

    protected int topologyOffset(int topologyIdx)
    {
        return topologyIdx * maxShardsPerEpoch();
    }

    public Topologies topologies()
    {
        return topologies;
    }

    protected <T extends AbstractTracker<ST>, P>
    RequestStatus recordResponse(T self, Id node, BiFunction<? super ST, P, ? extends ShardOutcome<? super T>> function, P param)
    {
        return recordResponse(self, node, function, param, topologies.size());
    }

    protected <T extends AbstractTracker<ST>, P>
    RequestStatus recordResponse(T self, Id node, BiFunction<? super ST, P, ? extends ShardOutcome<? super T>> function, P param, int topologyLimit)
    {
        return toStatusOrTryAgain(recordResponseInternal(self, node, function, param, topologyLimit));
    }

    protected <T extends AbstractTracker<ST>, P>
    ShardOutcomes recordResponseInternal(T self, Id node, BiFunction<? super ST, P, ? extends ShardOutcome<? super T>> function, P param)
    {
        return recordResponseInternal(self, node, function, param, topologies.size());
    }

    protected <T extends AbstractTracker<ST>, P>
    ShardOutcomes recordResponseInternal(T self, Id node, BiFunction<? super ST, P, ? extends ShardOutcome<? super T>> function, P param, int topologyLimit)
    {
        Invariants.require(self == this); // we just accept self as parameter for type safety
        ShardOutcomes status = NoChange;
        int maxShards = maxShardsPerEpoch();
        for (int i = 0; i < topologyLimit ; ++i)
        {
            status = topologies.get(i).mapReduceOn(node, i * maxShards, AbstractTracker::apply, self, function, param, ShardOutcomes::min, status);
        }

        return status;
    }

    protected RequestStatus toStatusOrTryAgain(ShardOutcomes outcomes)
    {
        Invariants.require(outcomes != SendMore);
        return outcomes.result;
    }

    static <ST extends ShardTracker, P, T extends AbstractTracker<ST>>
    ShardOutcomes apply(T tracker, BiFunction<? super ST, P, ? extends ShardOutcome<? super T>> function, P param, int trackerIndex)
    {
        return function.apply(tracker.trackers[trackerIndex], param).apply(tracker, trackerIndex);
    }

    public boolean any(Predicate<ST> test)
    {
        for (ST tracker : trackers)
        {
            if (tracker == null) continue;
            if (test.test(tracker))
                return true;
        }
        return false;
    }

    public boolean all(Predicate<ST> test)
    {
        for (ST tracker : trackers)
        {
            if (tracker == null) continue;
            if (!test.test(tracker))
                return false;
        }
        return true;
    }

    public SortedArrayList<Id> nodes()
    {
        return topologies.nodes();
    }

    public SortedArrayList<Id> filterAndRecordFaulty(@Nullable Tracing tracing)
    {
        return filterAndRecordFaulty(topologies.nodes(), tracing);
    }

    public SortedArrayList<Id> filterAndRecordFaulty(SortedArrayList<Id> ids, @Nullable Tracing tracing)
    {
        return filterAndRecordFaulty(ids, topologies, this, tracing);
    }

    public ST get(int shardIndex)
    {
        int maxShardsPerEpoch = maxShardsPerEpoch();
        return get(shardIndex / maxShardsPerEpoch, shardIndex % maxShardsPerEpoch);
    }

    @VisibleForTesting
    public ST get(int topologyIdx, int shardIdx)
    {
        if (shardIdx >= maxShardsPerEpoch())
            throw new IndexOutOfBoundsException();
        return trackers[topologyOffset(topologyIdx) + shardIdx];
    }

    protected int maxShardsPerEpoch()
    {
        return maxShardsPerEpoch;
    }

    public static SortedArrayList<Id> filterAndRecordFaulty(SortedArrayList<Id> nodes, TopologySorter sorter, AbstractTracker<?> reportTo, @Nullable Tracing tracing)
    {
        Object[] buffer = null;
        int bufferCount = 0;
        int previ = 0;
        for (int i = 0 ; i < nodes.size() ; ++i)
        {
            Id node = nodes.get(i);
            if (sorter.isFaulty(node))
            {
                if (tracing != null)
                    tracing.trace(null, "%s considered faulty; recording failure", node);
                if (Failed == reportTo.prerecordFailure(node))
                    return null;

                if (buffer == null)
                    buffer = cachedAny().get(nodes.size());

                int count = i - previ;
                System.arraycopy(nodes.backingArrayUnsafe(), previ, buffer, bufferCount, count);
                previ = i + 1;
                bufferCount += count;
            }
        }
        if (buffer == null)
            return nodes;

        int count = nodes.size() - previ;
        System.arraycopy(nodes.backingArrayUnsafe(), previ, buffer, bufferCount, count);
        bufferCount += count;

        Id[] newIds = new Id[bufferCount];
        System.arraycopy(buffer, 0, newIds, 0, bufferCount);
        return SortedArrayList.ofSorted(newIds);
    }

    public String summariseTracker()
    {
        StringBuilder sb = new StringBuilder();
        for (int i = 0 ; i < topologies.size() ; ++i)
        {
            Topology topology = topologies.get(i);
            if (sb.length() > 0)
                sb.append(", ");
            sb.append("[epoch=");
            sb.append(topology.epoch());
            for (int j = 0 ; j < topology.size() ; ++j)
            {
                ShardTracker tracker = trackers[topologyOffset(i) + j];
                if (tracker == null)
                    break;
                sb.append(", ");
                sb.append(tracker.summarise());
            }
            sb.append("]");
        }
        return sb.toString();
    }
}
