/*
 * 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.txn;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;

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

import accord.api.AsyncExecutor;
import accord.api.Write;
import accord.local.CommandStore;
import accord.local.SafeCommandStore;
import accord.primitives.PartialTxn;
import accord.primitives.Routable.Domain;
import accord.primitives.RoutableKey;
import accord.primitives.Seekable;
import accord.primitives.Seekables;
import accord.primitives.Timestamp;
import accord.primitives.TxnId;
import accord.utils.SimpleBitSet;
import accord.utils.SimpleBitSets;
import accord.utils.async.AsyncChain;
import accord.utils.async.AsyncChains;

import org.apache.cassandra.cql3.RowUpdateBuilder;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.Columns;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.Mutation;
import org.apache.cassandra.db.ReadCommand.PotentialTxnConflicts;
import org.apache.cassandra.db.RegularAndStaticColumns;
import org.apache.cassandra.db.TypeSizes;
import org.apache.cassandra.db.partitions.PartitionUpdate;
import org.apache.cassandra.db.rows.Row;
import org.apache.cassandra.io.ParameterisedVersionedSerializer;
import org.apache.cassandra.io.util.DataInputBuffer;
import org.apache.cassandra.io.util.DataInputPlus;
import org.apache.cassandra.io.util.DataOutputBuffer;
import org.apache.cassandra.io.util.DataOutputPlus;
import org.apache.cassandra.schema.ColumnMetadata;
import org.apache.cassandra.service.accord.AccordCommandStore;
import org.apache.cassandra.service.accord.AccordExecutor;
import org.apache.cassandra.service.accord.api.PartitionKey;
import org.apache.cassandra.service.accord.serializers.TableMetadatas;
import org.apache.cassandra.service.accord.serializers.TableMetadatasAndKeys;
import org.apache.cassandra.service.accord.serializers.Version;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.ObjectSizes;
import org.apache.cassandra.utils.SimpleBitSetSerializers;

import static com.google.common.base.Preconditions.checkState;
import static org.apache.cassandra.db.rows.DeserializationHelper.Flag.FROM_REMOTE;
import static org.apache.cassandra.utils.ArraySerializers.deserializeArray;
import static org.apache.cassandra.utils.ArraySerializers.serializeArray;
import static org.apache.cassandra.utils.ArraySerializers.serializedArraySize;
import static org.apache.cassandra.utils.ArraySerializers.skipArray;

public class TxnWrite extends AbstractKeySorted<TxnWrite.Update> implements Write
{
    public static final long NO_TIMESTAMP = 0;

    @SuppressWarnings("unused")
    private static final Logger logger = LoggerFactory.getLogger(TxnWrite.class);

    private static final long EMPTY_SIZE = ObjectSizes.measure(new TxnWrite(TableMetadatas.none(), Collections.emptyList(), SimpleBitSets.allUnset(1)));

    public static class Update extends AbstractParameterisedVersionedSerialized<PartitionUpdate, TableMetadatas>
    {
        private static final long EMPTY_SIZE = ObjectSizes.measure(new Update(null, 0, ByteBufferUtil.EMPTY_BYTE_BUFFER));
        public final PartitionKey key;
        public final int index;

        public Update(PartitionKey key, int index, PartitionUpdate update, TableMetadatas tables)
        {
            this(key, index, serializeInternal(update, tables, Version.LATEST));
        }

        private Update(PartitionKey key, int index, ByteBuffer latestVersionBytes)
        {
            super(latestVersionBytes);
            this.key = key;
            this.index = index;
        }

        @Override
        public long estimatedSizeOnHeap()
        {
            // we don't measure the key, as this is shared
            if (latestVersionBytes == null)
                return EMPTY_SIZE;
            return EMPTY_SIZE + ByteBufferUtil.estimatedSizeOnHeap(latestVersionBytes);
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            if (!super.equals(o)) return false;
            Update update = (Update) o;
            return index == update.index && key.equals(update.key);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(super.hashCode(), key, index);
        }

        @Override
        public String toString()
        {
            return "Complete{" +
                   "key=" + key +
                   ", index=" + index +
                   '}';
        }

        public AsyncChain<Void> write(AsyncExecutor executor, TableMetadatas tables, boolean preserveTimestamps, long timestamp)
        {
            PartitionUpdate update = deserialize(tables);
            if (!preserveTimestamps)
                update = new PartitionUpdate.Builder(update, 0).updateAllTimestamp(timestamp).build();
            Mutation mutation = new Mutation(update, PotentialTxnConflicts.ALLOW);
            return executor.chain(() -> mutation.apply(false, false));
        }

        @Override
        protected ByteBuffer serialize(PartitionUpdate value, TableMetadatas tables, Version version)
        {
            return serializeInternal(value, tables, version);
        }

        @Override
        protected ByteBuffer reserialize(ByteBuffer bytes, TableMetadatas param, Version srcVersion, Version trgVersion)
        {
            return bytes;
        }

        @Override
        protected PartitionUpdate deserialize(TableMetadatas tables, ByteBuffer bytes, Version version)
        {
            return deserialize(key, tables, bytes, version);
        }

        private static ByteBuffer serializeInternal(PartitionUpdate value, TableMetadatas tables, Version version)
        {
            try (DataOutputBuffer out = DataOutputBuffer.scratchBuffer.get())
            {
                PartitionUpdate.serializer.serializeWithoutKey(value, tables, out, version.messageVersion());
                return out.asNewBuffer();
            }
            catch (IOException e)
            {
                throw new UncheckedIOException(e);
            }
        }

        private static PartitionUpdate deserialize(PartitionKey key, TableMetadatas tables, ByteBuffer bytes, Version version)
        {
            try (DataInputBuffer in = new DataInputBuffer(bytes, true))
            {
                return PartitionUpdate.serializer.deserialize(key, tables, in, version.messageVersion(), FROM_REMOTE);
            }
            catch (IOException e)
            {
                throw new UncheckedIOException(e);
            }
        }

        public static final ParameterisedVersionedSerializer<Update, TableMetadatasAndKeys, Version> serializer = new ParameterisedVersionedSerializer<Update, TableMetadatasAndKeys, Version>()
        {
            @Override
            public void serialize(Update write, TableMetadatasAndKeys tablesAndKeys, DataOutputPlus out, Version version) throws IOException
            {
                tablesAndKeys.serializeKey(write.key, out);
                out.writeInt(write.index);
                ByteBufferUtil.writeWithVIntLength(write.bytes(tablesAndKeys.tables, version), out);
            }

            ByteBuffer reserialize(ByteBuffer buffer, TableMetadatasAndKeys tablesAndKeys, Version srcVersion, Version trgVersion)
            {
                return buffer;
            }

            @Override
            public Update deserialize(TableMetadatasAndKeys tablesAndKeys, DataInputPlus in, Version version) throws IOException
            {
                PartitionKey key = tablesAndKeys.deserializeKey(in);
                int index = in.readInt();
                ByteBuffer bytes = ByteBufferUtil.readWithVIntLength(in);
                if (version != Version.LATEST)
                    bytes = reserialize(bytes, tablesAndKeys, version, Version.LATEST);
                return new Update(key, index, bytes);
            }

            @Override
            public void skip(TableMetadatasAndKeys tablesAndKeys, DataInputPlus in, Version version) throws IOException
            {
                PartitionKey key = tablesAndKeys.deserializeKey(in);
                int index = in.readInt();
                ByteBufferUtil.skipWithVIntLength(in);
            }

            @Override
            public long serializedSize(Update write, TableMetadatasAndKeys tablesAndKeys, Version version)
            {
                long size = 0;
                size += tablesAndKeys.serializedKeySize(write.key);
                size += TypeSizes.INT_SIZE;
                size += ByteBufferUtil.serializedSizeWithVIntLength(write.bytes(tablesAndKeys.tables, version));
                return size;
            }
        };
    }

    /**
     * Partition update that can later be supplemented with data from the read phase
     */
    public static class Fragment
    {
        public final PartitionKey key;
        public final int index;
        public final PartitionUpdate baseUpdate;
        public final TxnReferenceOperations referenceOps;
        public final long timestamp;

        public Fragment(PartitionKey key, int index, PartitionUpdate baseUpdate, TxnReferenceOperations referenceOps, long timestamp)
        {
            this.key = key;
            this.index = index;
            this.baseUpdate = baseUpdate;
            this.referenceOps = referenceOps;
            this.timestamp = timestamp;
        }

        public static int compareKeys(Fragment left, Fragment right)
        {
            return left.key.compareTo(right.key);
        }

        @Override
        public boolean equals(Object o)
        {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Fragment fragment = (Fragment) o;
            return index == fragment.index && key.equals(fragment.key) && baseUpdate.equals(fragment.baseUpdate) && referenceOps.equals(fragment.referenceOps);
        }

        @Override
        public int hashCode()
        {
            return Objects.hash(key, index, baseUpdate, referenceOps);
        }

        @Override
        public String toString()
        {
            return "Fragment{key=" + key + ", index=" + index + ", baseUpdate=" + baseUpdate + ", referenceOps=" + referenceOps + '}';
        }

        public boolean isComplete()
        {
            return referenceOps.isEmpty();
        }

        public Update toUpdate(TableMetadatas tables)
        {
            return new Update(key, index, baseUpdate, tables);
        }

        public Update complete(AccordUpdateParameters parameters, TableMetadatas tables)
        {
            if (isComplete())
                return toUpdate(tables);

            DecoratedKey key = baseUpdate.partitionKey();
            PartitionUpdate.Builder updateBuilder = new PartitionUpdate.Builder(baseUpdate.metadata(),
                                                                                key,
                                                                                columns(baseUpdate, referenceOps),
                                                                                baseUpdate.rowCount(),
                                                                                baseUpdate.canHaveShadowedData());

            RowUpdateBuilder up = parameters.updateBuilder(baseUpdate.metadata(), key, index, timestamp);
            TxnData data = parameters.getData();
            Row staticRow = applyUpdates(baseUpdate.staticRow(), referenceOps.statics, key, Clustering.STATIC_CLUSTERING, up, data);

            if (!staticRow.isEmpty())
                updateBuilder.add(staticRow);

            for (Clustering<?> clustering : referenceOps.clusterings)
            {
                Row existing = baseUpdate.hasRows() ? baseUpdate.getRow(clustering) : null;
                Row row = applyUpdates(existing, referenceOps.regulars, key, clustering, up, data);
                if (row != null)
                    updateBuilder.add(row);
            }

            return new Update(this.key, index, updateBuilder.build(), tables);
        }

        private static Columns columns(Columns current, List<TxnReferenceOperation> referenceOps)
        {
            if (referenceOps.isEmpty())
                return current;

            Set<ColumnMetadata> missing = null;
            for (int i = 0, mi = referenceOps.size() ; i < mi ; ++i)
            {
                ColumnMetadata cm = referenceOps.get(i).receiver();
                if (!current.contains(cm))
                {
                    if (missing == null)
                        missing = new HashSet<>();
                    missing.add(cm);
                }
            }
            if (missing == null)
                return current;
            return current.mergeTo(Columns.from(missing));
        }

        private static RegularAndStaticColumns columns(PartitionUpdate update, TxnReferenceOperations referenceOps)
        {
            checkState(!referenceOps.isEmpty());
            RegularAndStaticColumns current = update.columns();
            return new RegularAndStaticColumns(columns(current.statics, referenceOps.statics),
                                               columns(current.regulars, referenceOps.regulars));
        }

        private static Row applyUpdates(Row existing, List<TxnReferenceOperation> operations, DecoratedKey key, Clustering<?> clustering, RowUpdateBuilder up, TxnData data)
        {
            if (operations.isEmpty())
                return existing;

            if (existing != null && !existing.isEmpty())
            {
                checkState(existing.clustering().equals(clustering));
                up.addRow(existing);
            }
            else
                up.newRow(clustering);

            operations.forEach(op -> op.apply(data, key, up));
            return up.buildRow();
        }

        static final FragmentSerializer serializer = new FragmentSerializer();
        static class FragmentSerializer
        {
            public void serialize(Fragment fragment, TableMetadatas tables, DataOutputPlus out, Version version) throws IOException
            {
                out.writeUnsignedVInt32(fragment.index);
                PartitionUpdate.serializer.serializeWithoutKey(fragment.baseUpdate, tables, out, version.messageVersion());
                TxnReferenceOperations.serializer.serialize(fragment.referenceOps, tables, out, version);
                out.writeUnsignedVInt(fragment.timestamp);
            }

            public Fragment deserialize(PartitionKey key, TableMetadatas tables, DataInputPlus in, Version version) throws IOException
            {
                int idx = in.readUnsignedVInt32();
                // TODO (required): why FROM_REMOTE?
                PartitionUpdate baseUpdate = PartitionUpdate.serializer.deserialize(key, tables, in, version.messageVersion(), FROM_REMOTE);
                TxnReferenceOperations referenceOps = TxnReferenceOperations.serializer.deserialize(tables, in, version);
                long timestamp = in.readUnsignedVInt();
                return new Fragment(key, idx, baseUpdate, referenceOps, timestamp);
            }

            public long serializedSize(Fragment fragment, TableMetadatas tables, Version version)
            {
                long size = 0;
                size += TypeSizes.sizeofUnsignedVInt(fragment.index);
                size += PartitionUpdate.serializer.serializedSizeWithoutKey(fragment.baseUpdate, tables, version.messageVersion());
                size += TxnReferenceOperations.serializer.serializedSize(fragment.referenceOps, tables, version);
                size += TypeSizes.sizeofUnsignedVInt(fragment.timestamp);
                return size;
            }

            public static ByteBuffer serialize(Fragment fragment, TableMetadatas tables, Version version)
            {
                long size = TypeSizes.sizeofUnsignedVInt(version.version);
                size += Fragment.serializer.serializedSize(fragment, tables, version);

                try (DataOutputBuffer out = new DataOutputBuffer((int) size))
                {
                    out.writeUnsignedVInt32(version.version);
                    Fragment.serializer.serialize(fragment, tables, out, version);
                    return out.buffer(false);
                }
                catch (IOException e)
                {
                    throw new RuntimeException(e);
                }
            }

            private static Fragment deserialize(PartitionKey key, TableMetadatas tables, ByteBuffer bytes)
            {
                try (DataInputBuffer in = new DataInputBuffer(bytes, true))
                {
                    Version version = Version.fromVersion(in.readUnsignedVInt32());
                    return Fragment.serializer.deserialize(key, tables, in, version);
                }
                catch (IOException e)
                {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    public final TableMetadatas tables;
    private final SimpleBitSet conditionalBlockBitSet;

    private TxnWrite(TableMetadatas tables, Update[] items, SimpleBitSet conditionalBlockBitSet)
    {
        super(items, Domain.Key);
        this.tables = tables;
        this.conditionalBlockBitSet = conditionalBlockBitSet;
    }

    public TxnWrite(TableMetadatas tables, List<Update> items, SimpleBitSet conditionalBlockBitSet)
    {
        super(items, Domain.Key);
        this.tables = tables;
        this.conditionalBlockBitSet = conditionalBlockBitSet;
    }

    @Override
    int compareNonKeyFields(Update left, Update right)
    {
        return Integer.compare(left.index, right.index);
    }

    @Override
    Seekable getKey(Update item)
    {
        return item.key;
    }

    @Override
    Update[] newArray(int size)
    {
        return new Update[size];
    }

    public void unmemoize()
    {
        for (int i = 0 ; i < size() ; ++i)
            items[i].unmemoize();
    }

    @Override
    public AsyncChain<Void> apply(SafeCommandStore safeStore, Seekable key, TxnId txnId, Timestamp executeAt, PartialTxn txn)
    {
        return applyDirect(safeStore.commandStore(), key, txnId, executeAt, txn);
    }

    @Override
    public AsyncChain<Void> applyDirect(CommandStore commandStore, Seekable key, TxnId txnId, Timestamp executeAt, PartialTxn txn)
    {
        // UnrecoverableRepairUpdate will deserialize as null at other nodes
        // Accord should skip the Update for a read transaction, but handle it here anyways
        TxnUpdate txnUpdate = ((TxnUpdate)txn.update());
        if (txnUpdate == null)
            return AsyncChains.success(null);

        long timestamp = executeAt.uniqueHlc();

        // TODO (expected): optimise for the common single update case; lots of lists allocated
        List<AsyncChain<Void>> results = new ArrayList<>();
        if (!conditionalBlockBitSet.isEmpty())
        {
            AccordExecutor executor = ((AccordCommandStore) commandStore).executor();
            boolean preserveTimestamps = txnUpdate.preserveTimestamps().preserve;
            // Apply updates not specified fully by the client but built from fragments completed by data from reads.
            // This occurs, for example, when an UPDATE statement uses a value assigned by a LET statement.
            forEachWithKey(key, write -> results.add(write.write(executor, tables, preserveTimestamps, timestamp)));
            // Apply updates that are fully specified by the client and not reliant on data from reads.
            // ex. INSERT INTO tbl (a, b, c) VALUES (1, 2, 3)
            // These updates are persisted only in TxnUpdate and not in TxnWrite to avoid duplication.
            List<Update> updates = txnUpdate.completeUpdatesForKey(conditionalBlockBitSet, (RoutableKey) key);
            updates.forEach(write -> results.add(write.write(executor, tables, preserveTimestamps, timestamp)));
        }

        if (results.isEmpty())
            return AsyncChains.success(null);

        if (results.size() == 1)
            return results.get(0).mapToNull();

        return AsyncChains.reduce(results, (i1, i2) -> null, null);
    }

    public long estimatedSizeOnHeap()
    {
        long size = EMPTY_SIZE;
        for (int i = 0 ; i < size() ; ++i)
            size += get(i).estimatedSizeOnHeap();
        return size;
    }

    public static final ParameterisedVersionedSerializer<TxnWrite, Seekables, Version> serializer = new ParameterisedVersionedSerializer<>()
    {
        @Override
        public void serialize(TxnWrite write, Seekables keys, DataOutputPlus out, Version version) throws IOException
        {
            write.tables.serializeSelf(out);
            SimpleBitSetSerializers.any.serialize(write.conditionalBlockBitSet, out);
            serializeArray(write.items, new TableMetadatasAndKeys(write.tables, keys), out, version, Update.serializer);
        }

        @Override
        public TxnWrite deserialize(Seekables keys, DataInputPlus in, Version version) throws IOException
        {
            TableMetadatas tables = TableMetadatas.deserializeSelf(in);
            SimpleBitSet conditionalBlockBitSet = SimpleBitSetSerializers.any.deserialize(in);
            return new TxnWrite(tables, deserializeArray(new TableMetadatasAndKeys(tables, keys), in, version, Update.serializer, Update[]::new), conditionalBlockBitSet);
        }

        @Override
        public void skip(Seekables keys, DataInputPlus in, Version version) throws IOException
        {
            TableMetadatas tables = TableMetadatas.deserializeSelf(in);
            SimpleBitSetSerializers.any.deserialize(in);
            skipArray(new TableMetadatasAndKeys(tables, keys), in, version, Update.serializer);
        }

        @Override
        public long serializedSize(TxnWrite write, Seekables keys, Version version)
        {
            return write.tables.serializedSelfSize()
                   + SimpleBitSetSerializers.any.serializedSize(write.conditionalBlockBitSet)
                   + serializedArraySize(write.items, new TableMetadatasAndKeys(write.tables, keys), version, Update.serializer);
        }
    };
}
