/*
 * 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 org.apache.cassandra.service.accord;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import com.google.common.collect.ImmutableMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import accord.api.Data;
import accord.api.DataStore;
import accord.api.Query;
import accord.api.Read;
import accord.api.Update;
import accord.coordinate.tracking.AbstractTracker;
import accord.impl.AbstractFetchCoordinator;
import accord.local.CommandStore;
import accord.local.Node;
import accord.local.SafeCommandStore;
import accord.primitives.PartialDeps;
import accord.primitives.PartialTxn;
import accord.primitives.Participants;
import accord.primitives.Range;
import accord.primitives.Ranges;
import accord.primitives.Routable;
import accord.primitives.Seekable;
import accord.primitives.Seekables;
import accord.primitives.SyncPoint;
import accord.primitives.Timestamp;
import accord.primitives.Txn;
import accord.primitives.TxnId;
import accord.topology.TopologyException;
import accord.utils.Invariants;
import accord.utils.async.AsyncChain;
import accord.utils.async.AsyncChains;

import org.apache.cassandra.db.TypeSizes;
import org.apache.cassandra.io.ParameterisedVersionedSerializer;
import org.apache.cassandra.io.UnversionedSerializer;
import org.apache.cassandra.io.util.DataInputPlus;
import org.apache.cassandra.io.util.DataOutputPlus;
import org.apache.cassandra.locator.InetAddressAndPort;
import org.apache.cassandra.locator.RangesAtEndpoint;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.accord.serializers.CommandSerializers;
import org.apache.cassandra.service.accord.serializers.IVersionedSerializer;
import org.apache.cassandra.service.accord.serializers.KeySerializers;
import org.apache.cassandra.service.accord.serializers.TableMetadatasAndKeys;
import org.apache.cassandra.service.accord.serializers.Version;
import org.apache.cassandra.streaming.PreviewKind;
import org.apache.cassandra.streaming.StreamCoordinator;
import org.apache.cassandra.streaming.StreamManager;
import org.apache.cassandra.streaming.StreamOperation;
import org.apache.cassandra.streaming.StreamPlan;
import org.apache.cassandra.streaming.StreamResultFuture;
import org.apache.cassandra.streaming.StreamSession;
import org.apache.cassandra.tcm.ClusterMetadata;
import org.apache.cassandra.utils.FBUtilities;
import org.apache.cassandra.utils.TimeUUID;

import static accord.primitives.Routables.Slice.Minimal;
import static org.apache.cassandra.utils.CollectionSerializers.deserializeMap;
import static org.apache.cassandra.utils.CollectionSerializers.serializeMap;
import static org.apache.cassandra.utils.CollectionSerializers.serializedMapSize;

public class AccordFetchCoordinator extends AbstractFetchCoordinator implements StreamManager.StreamListener
{
    private static final Logger logger = LoggerFactory.getLogger(AccordFetchCoordinator.class);
    private static final Query noopQuery = (txnId, executeAt, keys, data, read, update) -> null;

    public static class StreamData implements Data
    {
        public static class SessionInfo
        {
            final TimeUUID planId;
            final boolean hasData;

            public SessionInfo(TimeUUID planId, boolean hasData)
            {
                this.planId = planId;
                this.hasData = hasData;
            }

            static final UnversionedSerializer<SessionInfo> serializer = new UnversionedSerializer<>()
            {
                public void serialize(SessionInfo info, DataOutputPlus out) throws IOException
                {
                    TimeUUID.Serializer.instance.serialize(info.planId, out);
                    out.writeBoolean(info.hasData);
                }

                public SessionInfo deserialize(DataInputPlus in) throws IOException
                {
                    return new SessionInfo(TimeUUID.Serializer.instance.deserialize(in), in.readBoolean());
                }

                public long serializedSize(SessionInfo info)
                {
                    return TimeUUID.Serializer.instance.serializedSize(info.planId) + TypeSizes.BOOL_SIZE;
                }
            };
        }
        public static final UnversionedSerializer<StreamData> serializer = new UnversionedSerializer<>()
        {
            @Override
            public void serialize(StreamData data, DataOutputPlus out) throws IOException
            {
                serializeMap(data.streams, out, TokenRange.serializer, SessionInfo.serializer);
            }

            @Override
            public StreamData deserialize(DataInputPlus in) throws IOException
            {

                return new StreamData(ImmutableMap.copyOf(deserializeMap(in,
                                                                         TokenRange.serializer,
                                                                         SessionInfo.serializer)));
            }

            @Override
            public long serializedSize(StreamData data)
            {
                return serializedMapSize(data.streams, TokenRange.serializer, SessionInfo.serializer);
            }
        };

        private final ImmutableMap<TokenRange, SessionInfo> streams;

        public StreamData(ImmutableMap<TokenRange, SessionInfo> streams)
        {
            this.streams = streams;
        }

        public static StreamData of(TokenRange range, TimeUUID streamId, boolean hasData)
        {
            return new StreamData(ImmutableMap.of(range, new SessionInfo(streamId, hasData)));
        }

        @Override
        public StreamData merge(Data data)
        {
            StreamData that = (StreamData) data;
            if (that.streams.keySet().stream().anyMatch(this.streams::containsKey))
                throw new IllegalStateException(String.format("Unable to merge: key found in multiple StreamData %s %s",
                                                              this.streams.keySet(), that.streams.keySet()));
            Invariants.require(!that.streams.keySet().stream().anyMatch(this.streams::containsKey));
            ImmutableMap.Builder<TokenRange, SessionInfo> builder = ImmutableMap.builder();
            builder.putAll(this.streams);
            builder.putAll(that.streams);
            return new StreamData(builder.build());
        }

        @Override
        public Data without(Ranges ranges)
        {
            ImmutableMap.Builder<TokenRange, SessionInfo> copy = ImmutableMap.builder();
            for (Map.Entry<TokenRange, SessionInfo> e : streams.entrySet())
            {
                if (!ranges.intersects(e.getKey()))
                {
                    copy.put(e.getKey(), e.getValue());
                }
                else
                {
                    Ranges remainder = Ranges.of(e.getKey()).without(ranges);
                    for (Range range : remainder)
                        copy.put((TokenRange) range, e.getValue());
                }
            }
            return new StreamData(copy.build());
        }
    }

    // needs to be externally synchronized
    private class IncomingStream
    {
        private final TimeUUID planId;
        private Range range;
        private Node.Id from;
        private StreamResultFuture future;

        public IncomingStream(TimeUUID planId)
        {
            this.planId = planId;
        }

        private void rangeReceived(Range range, Node.Id from)
        {
            Invariants.nonNull(range);
            Invariants.nonNull(from);
            Invariants.require(this.range == null, "range was not null: %s", this.range);
            Invariants.require(this.from == null, "from was not null: %s", this.from);
            this.range = range;
            this.from = from;
            maybeListen();
        }

        private void futureReceived(StreamResultFuture future)
        {
            Invariants.nonNull(future);
            Invariants.require(this.future == null, "future was not null: %s", this.future);
            this.future = future;
            maybeListen();
        }

        private void maybeListen()
        {
            // TODO (required): if for some reason the stream isn't initiated this hangs
            if (range == null || future == null)
                return;

            logger.info("StreamFuture for plan {} received for bootstrap of {} from {}", planId, range, from);

            Invariants.nonNull(from);
            future.addCallback((state, fail) -> {
                if (fail == null)
                {
                    logger.info("Reporting success of plan {} for bootstrap of {} from {}", planId, range, from);
                    success(from, Ranges.of(range));
                }
                else
                {
                    logger.info("Reporting failure of plan {} for bootstrap of {} from {}", planId, range, from, fail);
                    fail(from, Ranges.of(range), fail);
                }
            }, ((AccordCommandStore) commandStore()).taskExecutor());
        }
    }

    public static class StreamingRead implements Read
    {
        public static final ParameterisedVersionedSerializer<StreamingRead, TableMetadatasAndKeys, Version> serializer = new ParameterisedVersionedSerializer<StreamingRead, TableMetadatasAndKeys, Version>()
        {
            @Override
            public void serialize(StreamingRead read, TableMetadatasAndKeys seekables, DataOutputPlus out, Version version) throws IOException
            {
                InetAddressAndPort.Serializer.inetAddressAndPortSerializer.serialize(read.to, out, version.messageVersion());
                KeySerializers.ranges.serialize(read.ranges, out);
            }

            @Override
            public StreamingRead deserialize(TableMetadatasAndKeys seekables, DataInputPlus in, Version version) throws IOException
            {
                return new StreamingRead(InetAddressAndPort.Serializer.inetAddressAndPortSerializer.deserialize(in, version.messageVersion()),
                                         KeySerializers.ranges.deserialize(in));
            }

            @Override
            public long serializedSize(StreamingRead read, TableMetadatasAndKeys seekables, Version version)
            {
                return InetAddressAndPort.Serializer.inetAddressAndPortSerializer.serializedSize(read.to, version.messageVersion())
                       + KeySerializers.ranges.serializedSize(read.ranges);
            }
        };

        private final InetAddressAndPort to;
        private final Ranges ranges;

        public StreamingRead(InetAddressAndPort to, Ranges ranges)
        {
            this.to = to;
            this.ranges = ranges;
        }

        @Override
        public Seekables<?, ?> keys() { return ranges; }

        private static boolean hasDataToStream(StreamCoordinator coordinator, InetAddressAndPort to)
        {
            for (StreamSession session : coordinator.getAllStreamSessions())
            {
                if (!session.peer.equals(to))
                    continue;

                Invariants.require(session.getNumRequests() == 0, "Requested to send data: %s", session);
                if (session.getNumTransfers() > 0)
                    return true;
            }
            return false;
        }

        @Override
        public AsyncChain<Data> read(SafeCommandStore commandStore, Seekable key, Timestamp executeAt)
        {
            try
            {
                Invariants.requireArgument(key.domain() == Routable.Domain.Range, "Required Range but saw %s: %s", key.domain(), key);
                TokenRange range = (TokenRange) key;

                // TODO (required): check epoch
                // TODO (required): handle dropped tables
                TableId tableId = range.table();
                TableMetadata table = ClusterMetadata.current().schema.getKeyspaces().getTableOrViewNullable(tableId);
                Invariants.require(table != null, "Table with id %s not found", tableId);

                // TODO (required): may also be relocation
                StreamPlan plan = new StreamPlan(StreamOperation.BOOTSTRAP, 1, false,
                                                 null, PreviewKind.NONE).flushBeforeTransfer(true);

                logger.info("StreamPlan {} created for bootstrap of {} to {}", plan.planId(), range, to);
                RangesAtEndpoint ranges = RangesAtEndpoint.toDummyList(Collections.singleton(range.toKeyspaceRange()));
                plan.transferRanges(to, table.keyspace, ranges, table.name);
                StreamResultFuture future = plan.execute();
                return AsyncChains.success(StreamData.of(range, future.planId, hasDataToStream(future.getCoordinator(), to)));
            }
            catch (Throwable t)
            {
                return AsyncChains.failure(t);
            }
        }

        @Override
        public Read slice(Ranges ranges) { return new StreamingRead(to, this.ranges.slice(ranges, Minimal)); }

        @Override
        public Read intersecting(Participants<?> participants) { return new StreamingRead(to, this.ranges.slice(ranges, Minimal)); }

        @Override
        public Read merge(Read other) { throw new UnsupportedOperationException(); }
    }

    public static class StreamingTxn
    {
        private static final ParameterisedVersionedSerializer<Read, TableMetadatasAndKeys, Version> read = (ParameterisedVersionedSerializer)StreamingRead.serializer;

        private static final UnversionedSerializer<Query> query = new UnversionedSerializer<>()
        {
            @Override
            public void serialize(Query t, DataOutputPlus out)
            {
                Invariants.requireArgument(t == noopQuery);
            }

            @Override
            public Query deserialize(DataInputPlus in)
            {
                return noopQuery;
            }

            @Override
            public long serializedSize(Query t)
            {
                Invariants.requireArgument(t == noopQuery);
                return 0;
            }
        };

        private static final ParameterisedVersionedSerializer<Update, TableMetadatasAndKeys, Version> update = new ParameterisedVersionedSerializer<>()
        {
            @Override
            public void serialize(Update t, TableMetadatasAndKeys seekables, DataOutputPlus out, Version version)
            {
                Invariants.requireArgument(t == null);
            }

            @Override
            public Update deserialize(TableMetadatasAndKeys seekables, DataInputPlus in, Version version)
            {
                return null;
            }

            @Override
            public long serializedSize(Update t, TableMetadatasAndKeys seekables, Version version)
            {
                Invariants.requireArgument(t == null);
                return 0;
            }
        };

        // TODO (desired): this could be serialized as an InetAddressAndPort and Ranges if we had a special case PartialTxn implementation
        public static final IVersionedSerializer<PartialTxn> serializer = new CommandSerializers.PartialTxnSerializer(read, query, update, TableMetadatasAndKeys.serializer);
    }

    private final Map<TimeUUID, IncomingStream> streams = new HashMap<>();

    public AccordFetchCoordinator(Node node, Ranges ranges, SyncPoint syncPoint, DataStore.FetchRanges fetchRanges, CommandStore commandStore) throws TopologyException
    {
        super(node, node.someSequentialExecutor(), ranges, syncPoint, fetchRanges, commandStore);
    }

    @Override
    public AbstractTracker<?> tracker()
    {
        return null;
    }

    @Override
    public void start()
    {
        StreamManager.instance.addListener(this);
        super.start();
    }

    private IncomingStream stream(TimeUUID id)
    {
        return streams.computeIfAbsent(id, IncomingStream::new);
    }

    // called from stream thread
    @Override
    public synchronized void onRegister(StreamResultFuture result)
    {
        stream(result.planId).futureReceived(result);
    }

    protected void onDone(Ranges success, Throwable failure)
    {
        StreamManager.instance.removeListener(this);
        super.onDone(success, failure);
    }

    @Override
    protected PartialTxn rangeReadTxn(Ranges ranges)
    {
        StreamingRead read = new StreamingRead(FBUtilities.getBroadcastAddressAndPort(), ranges);
        return new PartialTxn.InMemory(Txn.Kind.Read, ranges, read, noopQuery, null, TableMetadatasAndKeys.none(Routable.Domain.Range));
    }

    @Override
    protected synchronized void onReadOk(Node.Id from, CommandStore commandStore, Data data, Ranges received)
    {
        if (data == null)
            return;

        StreamData streamData = (StreamData) data;
        streamData.streams.forEach((range, streamInfo) -> {
            if (streamInfo.hasData)
            {
                logger.info("StreamPlan {} created for bootstrap of {} from {}", streamInfo.planId, range, from);
                stream(streamInfo.planId).rangeReceived(range, from);
            }
            else
            {
                logger.info("StreamPlan {} created for bootstrap of {} from {} with no data to stream; succeeding immediately.", streamInfo.planId, range, from);
                // if there was no data to stream, no connection is initiated, and we aren't notified via the stream
                // listener, so the stream initiator notifies us and we mark it complete here
                success(from, Ranges.of(range));
            }
        });
    }

    public static class AccordFetchRequest extends FetchRequest
    {
        public AccordFetchRequest(long sourceEpoch, TxnId syncId, Ranges ranges, PartialDeps partialDeps, PartialTxn partialTxn)
        {
            super(sourceEpoch, syncId, ranges, partialDeps, partialTxn);
        }

        @Override
        protected AsyncChain<Data> beginRead(SafeCommandStore safeStore, Timestamp executeAt, PartialTxn txn, Participants<?> execute)
        {
            AsyncChain<Data> result = super.beginRead(safeStore, executeAt, txn, execute);
            // TODO (required): verify that streaming snapshots have all been created by now, so we won't stream any data that arrives after this
            readStarted(safeStore);
            return result;
        }
    }

    @Override
    protected FetchRequest newFetchRequest(long sourceEpoch, TxnId syncId, Ranges ranges, PartialDeps partialDeps, PartialTxn partialTxn)
    {
        return new AccordFetchRequest(sourceEpoch, syncId, ranges, partialDeps, partialTxn);
    }
}
