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

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;

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

import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.lifecycle.LifecycleTransaction;
import org.apache.cassandra.dht.Range;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.io.sstable.ISSTableScanner;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.repair.consistent.admin.CleanupSummary;
import org.apache.cassandra.schema.CompactionParams;
import org.apache.cassandra.service.ActiveRepairService;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.TimeUUID;

/**
 * Companion to CompactionStrategyManager which manages the sstables marked pending repair.
 *
 * SSTables are classified as pending repair by the anti-compaction performed at the beginning
 * of an incremental repair, or when they're streamed in with a pending repair id. This prevents
 * unrepaired / pending repaired sstables from being compacted together. Once the repair session
 * has completed, or failed, sstables will be re-classified as part of the compaction process.
 */
class PendingRepairManager
{
    private static final Logger logger = LoggerFactory.getLogger(PendingRepairManager.class);

    private final ColumnFamilyStore cfs;
    private final CompactionParams params;
    private final boolean isTransient;
    private volatile ImmutableMap<TimeUUID, AbstractCompactionStrategy> strategies = ImmutableMap.of();

    /**
     * Indicates we're being asked to do something with an sstable that isn't marked pending repair
     */
    public static class IllegalSSTableArgumentException extends IllegalArgumentException
    {
        public IllegalSSTableArgumentException(String s)
        {
            super(s);
        }
    }

    PendingRepairManager(ColumnFamilyStore cfs, CompactionParams params, boolean isTransient)
    {
        this.cfs = cfs;
        this.params = params;
        this.isTransient = isTransient;
    }

    private ImmutableMap.Builder<TimeUUID, AbstractCompactionStrategy> mapBuilder()
    {
        return ImmutableMap.builder();
    }

    AbstractCompactionStrategy get(TimeUUID id)
    {
        return strategies.get(id);
    }

    AbstractCompactionStrategy get(SSTableReader sstable)
    {
        assert sstable.isPendingRepair();
        return get(sstable.getSSTableMetadata().pendingRepair);
    }

    AbstractCompactionStrategy getOrCreate(TimeUUID id)
    {
        checkPendingID(id);
        assert id != null;
        AbstractCompactionStrategy strategy = get(id);
        if (strategy == null)
        {
            synchronized (this)
            {
                strategy = get(id);

                if (strategy == null)
                {
                    logger.debug("Creating {}.{} compaction strategy for pending repair: {}", cfs.metadata.keyspace, cfs.metadata.name, id);
                    strategy = cfs.createCompactionStrategyInstance(params);
                    strategies = mapBuilder().putAll(strategies).put(id, strategy).build();
                }
            }
        }
        return strategy;
    }

    private static void checkPendingID(TimeUUID pendingID)
    {
        if (pendingID == null)
        {
            throw new IllegalSSTableArgumentException("sstable is not pending repair");
        }
    }

    AbstractCompactionStrategy getOrCreate(SSTableReader sstable)
    {
        return getOrCreate(sstable.getSSTableMetadata().pendingRepair);
    }

    private synchronized void removeSessionIfEmpty(TimeUUID sessionID)
    {
        if (!strategies.containsKey(sessionID) || !strategies.get(sessionID).getSSTables().isEmpty())
            return;

        logger.debug("Removing compaction strategy for pending repair {} on  {}.{}", sessionID, cfs.metadata.keyspace, cfs.metadata.name);
        strategies = ImmutableMap.copyOf(Maps.filterKeys(strategies, k -> !k.equals(sessionID)));
    }

    synchronized void removeSSTable(SSTableReader sstable)
    {
        for (Map.Entry<TimeUUID, AbstractCompactionStrategy> entry : strategies.entrySet())
        {
            entry.getValue().removeSSTable(sstable);
            removeSessionIfEmpty(entry.getKey());
        }
    }


    void removeSSTables(Iterable<SSTableReader> removed)
    {
        for (SSTableReader sstable : removed)
            removeSSTable(sstable);
    }

    synchronized void addSSTable(SSTableReader sstable)
    {
        Preconditions.checkArgument(sstable.isTransient() == isTransient);
        getOrCreate(sstable).addSSTable(sstable);
    }

    void addSSTables(Iterable<SSTableReader> added)
    {
        for (SSTableReader sstable : added)
            addSSTable(sstable);
    }

    synchronized void replaceSSTables(Set<SSTableReader> removed, Set<SSTableReader> added)
    {
        if (removed.isEmpty() && added.isEmpty())
            return;

        // left=removed, right=added
        Map<TimeUUID, Pair<Set<SSTableReader>, Set<SSTableReader>>> groups = new HashMap<>();
        for (SSTableReader sstable : removed)
        {
            TimeUUID sessionID = sstable.getSSTableMetadata().pendingRepair;
            if (!groups.containsKey(sessionID))
            {
                groups.put(sessionID, Pair.create(new HashSet<>(), new HashSet<>()));
            }
            groups.get(sessionID).left.add(sstable);
        }

        for (SSTableReader sstable : added)
        {
            TimeUUID sessionID = sstable.getSSTableMetadata().pendingRepair;
            if (!groups.containsKey(sessionID))
            {
                groups.put(sessionID, Pair.create(new HashSet<>(), new HashSet<>()));
            }
            groups.get(sessionID).right.add(sstable);
        }

        for (Map.Entry<TimeUUID, Pair<Set<SSTableReader>, Set<SSTableReader>>> entry : groups.entrySet())
        {
            AbstractCompactionStrategy strategy = getOrCreate(entry.getKey());
            Set<SSTableReader> groupRemoved = entry.getValue().left;
            Set<SSTableReader> groupAdded = entry.getValue().right;

            if (!groupRemoved.isEmpty())
                strategy.replaceSSTables(groupRemoved, groupAdded);
            else
                strategy.addSSTables(groupAdded);

            removeSessionIfEmpty(entry.getKey());
        }
    }

    synchronized void startup()
    {
        strategies.values().forEach(AbstractCompactionStrategy::startup);
    }

    synchronized void shutdown()
    {
        strategies.values().forEach(AbstractCompactionStrategy::shutdown);
    }

    private int getEstimatedRemainingTasks(TimeUUID sessionID, AbstractCompactionStrategy strategy)
    {
        return getEstimatedRemainingTasks(sessionID, strategy, 0, 0);
    }

    private int getEstimatedRemainingTasks(TimeUUID sessionID, AbstractCompactionStrategy strategy, int additionalSSTables, long additionalBytes)
    {
        return canCleanup(sessionID) ? 0 : strategy.getEstimatedRemainingTasks();
    }

    int getEstimatedRemainingTasks()
    {
        return getEstimatedRemainingTasks(0, 0);
    }

    int getEstimatedRemainingTasks(int additionalSSTables, long additionalBytes)
    {
        int tasks = 0;
        for (Map.Entry<TimeUUID, AbstractCompactionStrategy> entry : strategies.entrySet())
        {
            tasks += getEstimatedRemainingTasks(entry.getKey(), entry.getValue(), additionalSSTables, additionalBytes);
        }
        return tasks;
    }

    /**
     * @return the highest max remaining tasks of all contained compaction strategies
     */
    int getMaxEstimatedRemainingTasks()
    {
        int tasks = 0;
        for (Map.Entry<TimeUUID, AbstractCompactionStrategy> entry : strategies.entrySet())
        {
            tasks = Math.max(tasks, getEstimatedRemainingTasks(entry.getKey(), entry.getValue()));
        }
        return tasks;
    }

    private RepairFinishedCompactionTask getRepairFinishedCompactionTask(TimeUUID sessionID)
    {
        Preconditions.checkState(canCleanup(sessionID));
        AbstractCompactionStrategy compactionStrategy = get(sessionID);
        if (compactionStrategy == null)
            return null;
        Set<SSTableReader> sstables = compactionStrategy.getSSTables();
        long repairedAt = ActiveRepairService.instance().consistent.local.getFinalSessionRepairedAt(sessionID);
        LifecycleTransaction txn = cfs.getTracker().tryModify(sstables, OperationType.COMPACTION);
        return txn == null ? null : new RepairFinishedCompactionTask(cfs, txn, sessionID, repairedAt);
    }

    public static class CleanupTask
    {
        private final ColumnFamilyStore cfs;
        private final List<Pair<TimeUUID, RepairFinishedCompactionTask>> tasks;

        public CleanupTask(ColumnFamilyStore cfs, List<Pair<TimeUUID, RepairFinishedCompactionTask>> tasks)
        {
            this.cfs = cfs;
            this.tasks = tasks;
        }

        public CleanupSummary cleanup()
        {
            Set<TimeUUID> successful = new HashSet<>();
            Set<TimeUUID> unsuccessful = new HashSet<>();
            for (Pair<TimeUUID, RepairFinishedCompactionTask> pair : tasks)
            {
                TimeUUID session = pair.left;
                RepairFinishedCompactionTask task = pair.right;

                if (task != null)
                {
                    try
                    {
                        task.run();
                        successful.add(session);
                    }
                    catch (Throwable t)
                    {
                        t = task.transaction.abort(t);
                        logger.error("Failed cleaning up " + session, t);
                        unsuccessful.add(session);
                    }
                }
                else
                {
                    unsuccessful.add(session);
                }
            }
            return new CleanupSummary(cfs, successful, unsuccessful);
        }

        public Throwable abort(Throwable accumulate)
        {
            for (Pair<TimeUUID, RepairFinishedCompactionTask> pair : tasks)
                accumulate = pair.right.transaction.abort(accumulate);
            return accumulate;
        }
    }

    public CleanupTask releaseSessionData(Collection<TimeUUID> sessionIDs)
    {
        List<Pair<TimeUUID, RepairFinishedCompactionTask>> tasks = new ArrayList<>(sessionIDs.size());
        for (TimeUUID session : sessionIDs)
        {
            if (hasDataForSession(session))
            {
                tasks.add(Pair.create(session, getRepairFinishedCompactionTask(session)));
            }
        }
        return new CleanupTask(cfs, tasks);
    }

    synchronized int getNumPendingRepairFinishedTasks()
    {
        int count = 0;
        for (TimeUUID sessionID : strategies.keySet())
        {
            if (canCleanup(sessionID))
            {
                count++;
            }
        }
        return count;
    }

    synchronized Collection<AbstractCompactionTask> getNextRepairFinishedTasks()
    {
        List<AbstractCompactionTask> tasks = new ArrayList<>();
        for (TimeUUID sessionID : strategies.keySet())
        {
            if (canCleanup(sessionID))
            {
                RepairFinishedCompactionTask repairFinishedTask = getRepairFinishedCompactionTask(sessionID);
                if (repairFinishedTask != null)
                    tasks.add(repairFinishedTask);
            }
        }
        return tasks;
    }

    synchronized Collection<AbstractCompactionTask> getNextBackgroundTasks(long gcBefore)
    {
        if (strategies.isEmpty())
            return Collections.emptyList();

        Map<TimeUUID, Integer> numTasks = new HashMap<>(strategies.size());
        ArrayList<TimeUUID> sessions = new ArrayList<>(strategies.size());
        for (Map.Entry<TimeUUID, AbstractCompactionStrategy> entry : strategies.entrySet())
        {
            if (canCleanup(entry.getKey()))
            {
                continue;
            }
            numTasks.put(entry.getKey(), getEstimatedRemainingTasks(entry.getKey(), entry.getValue()));
            sessions.add(entry.getKey());
        }

        if (sessions.isEmpty())
            return Collections.emptyList();

        // we want the session with the most compactions at the head of the list
        sessions.sort((o1, o2) -> numTasks.get(o2) - numTasks.get(o1));

        TimeUUID sessionID = sessions.get(0);
        return get(sessionID).getNextBackgroundTasks(gcBefore);
    }

    synchronized Collection<AbstractCompactionTask> getMaximalTasks(long gcBefore, boolean splitOutput)
    {
        if (strategies.isEmpty())
            return null;

        List<AbstractCompactionTask> maximalTasks = new ArrayList<>(strategies.size());
        for (Map.Entry<TimeUUID, AbstractCompactionStrategy> entry : strategies.entrySet())
        {
            if (canCleanup(entry.getKey()))
            {
                maximalTasks.add(getRepairFinishedCompactionTask(entry.getKey()));
            }
            else
            {
                Collection<AbstractCompactionTask> tasks = entry.getValue().getMaximalTasks(gcBefore, splitOutput);
                if (tasks != null)
                    maximalTasks.addAll(tasks);
            }
        }
        return !maximalTasks.isEmpty() ? maximalTasks : null;
    }

    Collection<AbstractCompactionStrategy> getStrategies()
    {
        return strategies.values();
    }

    Set<TimeUUID> getSessions()
    {
        return strategies.keySet();
    }

    boolean canCleanup(TimeUUID sessionID)
    {
        return !ActiveRepairService.instance().consistent.local.isSessionInProgress(sessionID);
    }

    synchronized Set<ISSTableScanner> getScanners(Collection<SSTableReader> sstables, Collection<Range<Token>> ranges)
    {
        if (sstables.isEmpty())
        {
            return Collections.emptySet();
        }

        Map<TimeUUID, Set<SSTableReader>> sessionSSTables = new HashMap<>();
        for (SSTableReader sstable : sstables)
        {
            TimeUUID sessionID = sstable.getSSTableMetadata().pendingRepair;
            checkPendingID(sessionID);
            sessionSSTables.computeIfAbsent(sessionID, k -> new HashSet<>()).add(sstable);
        }

        Set<ISSTableScanner> scanners = new HashSet<>(sessionSSTables.size());
        try
        {
            for (Map.Entry<TimeUUID, Set<SSTableReader>> entry : sessionSSTables.entrySet())
            {
                scanners.addAll(getOrCreate(entry.getKey()).getScanners(entry.getValue(), ranges).scanners);
            }
        }
        catch (Throwable t)
        {
            ISSTableScanner.closeAllAndPropagate(scanners, t);
        }
        return scanners;
    }

    public boolean hasStrategy(AbstractCompactionStrategy strategy)
    {
        return strategies.values().contains(strategy);
    }

    public synchronized boolean hasDataForSession(TimeUUID sessionID)
    {
        return strategies.containsKey(sessionID);
    }

    boolean containsSSTable(SSTableReader sstable)
    {
        if (!sstable.isPendingRepair())
            return false;

        AbstractCompactionStrategy strategy = strategies.get(sstable.getPendingRepair());
        return strategy != null && strategy.getSSTables().contains(sstable);
    }

    public Collection<AbstractCompactionTask> createUserDefinedTasks(Collection<SSTableReader> sstables, long gcBefore)
    {
        Map<TimeUUID, List<SSTableReader>> group = sstables.stream().collect(Collectors.groupingBy(s -> s.getSSTableMetadata().pendingRepair));
        return group.entrySet().stream().map(g -> strategies.get(g.getKey()).getUserDefinedTask(g.getValue(), gcBefore)).collect(Collectors.toList());
    }

    @VisibleForTesting
    public synchronized boolean hasPendingRepairSSTable(TimeUUID sessionID, SSTableReader sstable)
    {
        AbstractCompactionStrategy strat = strategies.get(sessionID);
        if (strat == null)
            return false;
        return strat.getSSTables().contains(sstable);
    }

    /**
     * promotes/demotes sstables involved in a consistent repair that has been finalized, or failed
     */
    class RepairFinishedCompactionTask extends AbstractCompactionTask
    {
        private final TimeUUID sessionID;
        private final long repairedAt;

        RepairFinishedCompactionTask(ColumnFamilyStore cfs, LifecycleTransaction transaction, TimeUUID sessionID, long repairedAt)
        {
            super(cfs, transaction);
            this.sessionID = sessionID;
            this.repairedAt = repairedAt;
        }

        @VisibleForTesting
        TimeUUID getSessionID()
        {
            return sessionID;
        }

        protected void runMayThrow() throws Exception
        {
            boolean completed = false;
            boolean obsoleteSSTables = isTransient && repairedAt > 0;
            try
            {
                if (obsoleteSSTables)
                {
                    logger.info("Obsoleting transient repaired sstables for {}", sessionID);
                    Preconditions.checkState(Iterables.all(transaction.originals(), SSTableReader::isTransient));
                    transaction.obsoleteOriginals();
                }
                else
                {
                    logger.info("Moving {} from pending to repaired with repaired at = {} and session id = {}", transaction.originals(), repairedAt, sessionID);
                    cfs.getCompactionStrategyManager().mutateRepaired(transaction.originals(), repairedAt, ActiveRepairService.NO_PENDING_REPAIR, false);
                }
                completed = true;
            }
            finally
            {
                if (obsoleteSSTables)
                {
                    transaction.prepareToCommit();
                    transaction.commit();
                }
                else
                {
                    // we abort here because mutating metadata isn't guarded by LifecycleTransaction, so this won't roll
                    // anything back. Also, we don't want to obsolete the originals. We're only using it to prevent other
                    // compactions from marking these sstables compacting, and unmarking them when we're done
                    transaction.abort();
                }
                if (completed)
                {
                    removeSessionIfEmpty(sessionID);
                }
            }
        }

        protected void executeInternal(ActiveCompactionsTracker activeCompactions)
        {
            run();
        }
    }

}
