/*
 * 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.db.index.sasi.disk;

import java.io.File;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.*;

import org.apache.cassandra.concurrent.JMXEnabledThreadPoolExecutor;
import org.apache.cassandra.concurrent.NamedThreadFactory;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.Column;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.index.sasi.analyzer.AbstractAnalyzer;
import org.apache.cassandra.db.index.sasi.conf.ColumnIndex;
import org.apache.cassandra.db.index.sasi.utils.CombinedTermIterator;
import org.apache.cassandra.db.index.sasi.utils.TypeUtil;
import org.apache.cassandra.db.marshal.AbstractType;
import org.apache.cassandra.io.sstable.Descriptor;
import org.apache.cassandra.io.sstable.SSTableWriterListener;
import org.apache.cassandra.io.util.FileUtils;
import org.apache.cassandra.utils.Pair;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.Uninterruptibles;

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

public class PerSSTableIndexWriter implements SSTableWriterListener
{
    private static final Logger logger = LoggerFactory.getLogger(PerSSTableIndexWriter.class);

    private static final ThreadPoolExecutor INDEX_FLUSHER_MEMTABLE;
    private static final ThreadPoolExecutor INDEX_FLUSHER_GENERAL;

    static
    {
        INDEX_FLUSHER_GENERAL = new JMXEnabledThreadPoolExecutor(1, 8, 60, TimeUnit.SECONDS,
                                                                 new LinkedBlockingQueue<Runnable>(),
                                                                 new NamedThreadFactory("SASI-General"),
                                                                 "internal");
        INDEX_FLUSHER_GENERAL.allowCoreThreadTimeOut(true);

        INDEX_FLUSHER_MEMTABLE = new JMXEnabledThreadPoolExecutor(1, 8, 60, TimeUnit.SECONDS,
                                                                  new LinkedBlockingQueue<Runnable>(),
                                                                  new NamedThreadFactory("SASI-Memtable"),
                                                                  "internal");
        INDEX_FLUSHER_MEMTABLE.allowCoreThreadTimeOut(true);
    }

    private final Descriptor descriptor;
    private final Source source;

    private final AbstractType<?> keyValidator;
    private final Map<ByteBuffer, ColumnIndex> supportedIndexes;

    @VisibleForTesting
    protected final Map<ByteBuffer, Index> indexes;

    private DecoratedKey currentKey;
    private long currentKeyPosition;
    private boolean isComplete;

    private final long start = System.currentTimeMillis();

    public PerSSTableIndexWriter(AbstractType<?> keyValidator,
                                 Descriptor descriptor,
                                 Source source,
                                 Map<ByteBuffer, ColumnIndex> supportedIndexes)
    {
        this.keyValidator = keyValidator;
        this.descriptor = descriptor;
        this.source = source;
        this.supportedIndexes = supportedIndexes;
        this.indexes = new HashMap<>();
    }

    @Override
    public void begin()
    {}

    @Override
    public void startRow(DecoratedKey key, long curPosition)
    {
        currentKey = key;
        currentKeyPosition = curPosition;
    }

    @Override
    public void nextColumn(Column column)
    {
        // avoid adding already expired or deleted columns to the index
        if (column.isMarkedForDelete(start))
            return;

        ColumnIndex columnIndex = supportedIndexes.get(column.name());
        if (columnIndex == null)
            return;

        Index index = indexes.get(column.name());
        if (index == null)
            indexes.put(column.name(), (index = new Index(columnIndex)));

        index.add(column.value().duplicate(), currentKey, currentKeyPosition);
    }

    public Index getIndex(ByteBuffer columnName)
    {
        return indexes.get(columnName);
    }

    public Descriptor getDescriptor()
    {
        return descriptor;
    }

    @Override
    public void complete()
    {
        if (isComplete)
            return;

        currentKey = null;

        try
        {
            CountDownLatch latch = new CountDownLatch(indexes.size());
            for (Index index : indexes.values())
                index.complete(latch);

            Uninterruptibles.awaitUninterruptibly(latch);
        }
        finally
        {
            indexes.clear();
            isComplete = true;
        }
    }

    @VisibleForTesting
    protected class Index
    {
        private final ColumnIndex columnIndex;
        private final String outputFile;
        private final AbstractAnalyzer analyzer;
        private final long maxMemorySize;

        @VisibleForTesting
        protected final Set<Future<OnDiskIndex>> segments;
        private int segmentNumber = 0;

        private OnDiskIndexBuilder currentBuilder;

        public Index(ColumnIndex columnIndex)
        {
            this.columnIndex = columnIndex;
            this.outputFile = descriptor.filenameFor(columnIndex.getComponent());
            this.analyzer = columnIndex.getAnalyzer();
            this.segments = new HashSet<>();
            this.maxMemorySize = maxMemorySize(columnIndex);
            this.currentBuilder = newIndexBuilder();
        }

        public void add(ByteBuffer term, DecoratedKey key, long keyPosition)
        {
            if (term.remaining() == 0)
                return;

            boolean isAdded = false;

            analyzer.reset(term);
            while (analyzer.hasNext())
            {
                ByteBuffer token = analyzer.next();
                int size = token.remaining();

                if (token.remaining() >= OnDiskIndexBuilder.MAX_TERM_SIZE)
                {
                    logger.info("Rejecting value (size {}, maximum {} bytes) for column {} (analyzed {}) at {} SSTable.",
                            term.remaining(),
                            OnDiskIndexBuilder.MAX_TERM_SIZE,
                            columnIndex.getColumnName(),
                            columnIndex.getMode().isAnalyzed,
                            descriptor);
                    continue;
                }

                if (!TypeUtil.isValid(token, columnIndex.getValidator()))
                {
                    if ((token = TypeUtil.tryUpcast(token, columnIndex.getValidator())) == null)
                    {
                        logger.info("({}) Failed to add {} to index for key: {}, value size was {} bytes, validator is {}.",
                                    outputFile,
                                    columnIndex.getColumnName(),
                                    keyValidator.getString(key.key),
                                    size,
                                    columnIndex.getValidator());
                        continue;
                    }
                }

                currentBuilder.add(token, key, keyPosition);
                isAdded = true;
            }

            if (!isAdded || currentBuilder.estimatedMemoryUse() < maxMemorySize)
                return; // non of the generated tokens were added to the index or memory size wasn't reached

            segments.add(getExecutor().submit(scheduleSegmentFlush(false)));
        }

        @VisibleForTesting
        protected Callable<OnDiskIndex> scheduleSegmentFlush(final boolean isFinal)
        {
            final OnDiskIndexBuilder builder = currentBuilder;
            currentBuilder = newIndexBuilder();

            final String segmentFile = filename(isFinal);

            return new Callable<OnDiskIndex>()
            {
                @Override
                public OnDiskIndex call()
                {
                    long start = System.nanoTime();

                    try
                    {
                        File index = new File(segmentFile);
                        return builder.finish(index) ? new OnDiskIndex(index, columnIndex.getValidator(), null) : null;
                    }
                    finally
                    {
                        if (!isFinal)
                            logger.info("Flushed index segment {}, took {} ms.", segmentFile, TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start));
                    }
                }
            };
        }

        public void complete(final CountDownLatch latch)
        {
            logger.info("Scheduling index flush to {}", outputFile);

            getExecutor().submit(new Runnable()
            {
                @Override
                public void run()
                {
                    long start = System.nanoTime();

                    OnDiskIndex[] parts = new OnDiskIndex[segments.size() + 1];

                    try
                    {
                        // no parts present, build entire index from memory
                        if (segments.isEmpty())
                        {
                            scheduleSegmentFlush(true).call();
                            return;
                        }

                        // parts are present but there is something still in memory, let's flush that inline
                        if (!currentBuilder.isEmpty())
                        {
                            OnDiskIndex last = scheduleSegmentFlush(false).call();
                            segments.add(Futures.immediateFuture(last));
                        }

                        int index = 0;
                        ByteBuffer combinedMin = null, combinedMax = null;

                        for (Future<OnDiskIndex> f : segments)
                        {
                            OnDiskIndex part = Futures.getUnchecked(f);
                            if (part == null)
                                continue;

                            parts[index++] = part;
                            combinedMin = (combinedMin == null || keyValidator.compare(combinedMin, part.minKey()) > 0) ? part.minKey() : combinedMin;
                            combinedMax = (combinedMax == null || keyValidator.compare(combinedMax, part.maxKey()) < 0) ? part.maxKey() : combinedMax;
                        }

                        OnDiskIndexBuilder builder = newIndexBuilder();
                        builder.finish(Pair.create(combinedMin, combinedMax),
                                       new File(outputFile),
                                       new CombinedTermIterator(parts));
                    }
                    catch (Exception e)
                    {
                        logger.error("Failed to flush index {}.", outputFile, e);
                        FileUtils.delete(outputFile);
                    }
                    finally
                    {
                        logger.info("Index flush to {} took {} ms.", outputFile, TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start));

                        for (OnDiskIndex part : parts)
                        {
                            if (part == null)
                                continue;

                            FileUtils.closeQuietly(part);
                            FileUtils.delete(part.getIndexPath());
                        }

                        latch.countDown();
                    }
                }
            });
        }

        private ExecutorService getExecutor()
        {
            return source == Source.MEMTABLE ? INDEX_FLUSHER_MEMTABLE : INDEX_FLUSHER_GENERAL;
        }

        private OnDiskIndexBuilder newIndexBuilder()
        {
            return new OnDiskIndexBuilder(keyValidator, columnIndex.getValidator(), columnIndex.getMode().mode);
        }

        public String filename(boolean isFinal)
        {
            return outputFile + (isFinal ? "" : "_" + segmentNumber++);
        }
    }

    protected long maxMemorySize(ColumnIndex columnIndex)
    {
        final long def = DatabaseDescriptor.getTotalMemtableSpaceInMB() * 1048576L;
        return source == Source.MEMTABLE ? def : columnIndex.getMode().maxCompactionFlushMemoryInMb;
    }

    @Override
    public int hashCode()
    {
        return descriptor.hashCode();
    }

    @Override
    public boolean equals(Object o)
    {
        return !(o == null || !(o instanceof PerSSTableIndexWriter)) && descriptor.equals(((PerSSTableIndexWriter) o).descriptor);
    }
}
