/*
 * 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.druid.segment.loading;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.annotations.VisibleForTesting;
import com.google.inject.Inject;
import org.apache.druid.guice.annotations.Json;
import org.apache.druid.java.util.common.FileUtils;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.emitter.EmittingLogger;
import org.apache.druid.timeline.DataSegment;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import java.io.File;
import java.io.IOException;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 *
 */
public class SegmentLocalCacheManager implements SegmentCacheManager
{
  @VisibleForTesting
  static final String DOWNLOAD_START_MARKER_FILE_NAME = "downloadStartMarker";

  private static final EmittingLogger log = new EmittingLogger(SegmentLocalCacheManager.class);

  private final SegmentLoaderConfig config;
  private final ObjectMapper jsonMapper;

  private final List<StorageLocation> locations;

  // This directoryWriteRemoveLock is used when creating or removing a directory
  private final Object directoryWriteRemoveLock = new Object();

  /**
   * A map between segment and referenceCountingLocks.
   *
   * These locks should be acquired whenever getting or deleting files for a segment.
   * If different threads try to get or delete files simultaneously, one of them creates a lock first using
   * {@link #createOrGetLock}. And then, all threads compete with each other to get the lock.
   * Finally, the lock should be released using {@link #unlock}.
   *
   * An example usage is:
   *
   * final ReferenceCountingLock lock = createOrGetLock(segment);
   * synchronized (lock) {
   *   try {
   *     doSomething();
   *   }
   *   finally {
   *     unlock(lock);
   *   }
   * }
   */
  private final ConcurrentHashMap<DataSegment, ReferenceCountingLock> segmentLocks = new ConcurrentHashMap<>();

  private final StorageLocationSelectorStrategy strategy;

  // Note that we only create this via injection in historical and realtime nodes. Peons create these
  // objects via SegmentCacheManagerFactory objects, so that they can store segments in task-specific
  // directories rather than statically configured directories.
  @Inject
  public SegmentLocalCacheManager(
      List<StorageLocation> locations,
      SegmentLoaderConfig config,
      @Nonnull StorageLocationSelectorStrategy strategy,
      @Json ObjectMapper mapper
  )
  {
    this.config = config;
    this.jsonMapper = mapper;
    this.locations = locations;
    this.strategy = strategy;
    log.info("Using storage location strategy: [%s]", this.strategy.getClass().getSimpleName());
  }

  @VisibleForTesting
  SegmentLocalCacheManager(
      SegmentLoaderConfig config,
      @Nonnull StorageLocationSelectorStrategy strategy,
      @Json ObjectMapper mapper
  )
  {
    this(config.toStorageLocations(), config, strategy, mapper);
  }

  /**
   * creates instance with default storage location selector strategy
   *
   * This ctor is mainly for test cases, including test cases in other modules
   */
  @VisibleForTesting
  public SegmentLocalCacheManager(
      SegmentLoaderConfig config,
      @Json ObjectMapper mapper
  )
  {
    this.config = config;
    this.jsonMapper = mapper;
    this.locations = config.toStorageLocations();
    this.strategy = new LeastBytesUsedStorageLocationSelectorStrategy(locations);
    log.info("Using storage location strategy: [%s]", this.strategy.getClass().getSimpleName());
  }


  static String getSegmentDir(DataSegment segment)
  {
    return DataSegmentPusher.getDefaultStorageDir(segment, false);
  }

  @Override
  public boolean isSegmentCached(final DataSegment segment)
  {
    return findStoragePathIfCached(segment) != null;
  }

  /**
   * This method will try to find if the segment is already downloaded on any location. If so, the segment path
   * is returned. Along with that, location state is also updated with the segment location. Refer to
   * {@link StorageLocation#maybeReserve(String, DataSegment)} for more details.
   * If the segment files are damaged in any location, they are removed from the location.
   * @param segment - Segment to check
   * @return - Path corresponding to segment directory if found, null otherwise.
   */
  @Nullable
  private File findStoragePathIfCached(final DataSegment segment)
  {
    for (StorageLocation location : locations) {
      String storageDir = getSegmentDir(segment);
      File localStorageDir = location.segmentDirectoryAsFile(storageDir);
      if (localStorageDir.exists()) {
        if (checkSegmentFilesIntact(localStorageDir)) {
          log.warn(
              "[%s] may be damaged. Delete all the segment files and pull from DeepStorage again.",
              localStorageDir.getAbsolutePath()
          );
          cleanupCacheFiles(location.getPath(), localStorageDir);
          location.removeSegmentDir(localStorageDir, segment);
          break;
        } else {
          // Before returning, we also reserve the space. Refer to the StorageLocation#maybeReserve documentation for details.
          location.maybeReserve(storageDir, segment);
          return localStorageDir;
        }
      }
    }
    return null;
  }

  /**
   * check data intact.
   * @param dir segments cache dir
   * @return true means segment files may be damaged.
   */
  private boolean checkSegmentFilesIntact(File dir)
  {
    return checkSegmentFilesIntactWithStartMarker(dir);
  }

  /**
   * If there is 'downloadStartMarker' existed in localStorageDir, the segments files might be damaged.
   * Because each time, Druid will delete the 'downloadStartMarker' file after pulling and unzip the segments from DeepStorage.
   * downloadStartMarker existed here may mean something error during download segments and the segment files may be damaged.
   */
  private boolean checkSegmentFilesIntactWithStartMarker(File localStorageDir)
  {
    final File downloadStartMarker = new File(localStorageDir.getPath(), DOWNLOAD_START_MARKER_FILE_NAME);
    return downloadStartMarker.exists();
  }

  /**
   * Make sure segments files in loc is intact, otherwise function like loadSegments will failed because of segment files is damaged.
   * @param segment
   * @return
   * @throws SegmentLoadingException
   */
  @Override
  public File getSegmentFiles(DataSegment segment) throws SegmentLoadingException
  {
    final ReferenceCountingLock lock = createOrGetLock(segment);
    synchronized (lock) {
      try {
        File segmentDir = findStoragePathIfCached(segment);
        if (segmentDir != null) {
          return segmentDir;
        }

        return loadSegmentWithRetry(segment);
      }
      finally {
        unlock(segment, lock);
      }
    }
  }

  /**
   * If we have already reserved a location before, probably via {@link #reserve(DataSegment)}, then only that location
   * should be tried. Otherwise, we would fetch locations using {@link StorageLocationSelectorStrategy} and try all
   * of them one by one till there is success.
   * Location may fail because of IO failure, most likely in two cases:<p>
   * 1. druid don't have the write access to this location, most likely the administrator doesn't config it correctly<p>
   * 2. disk failure, druid can't read/write to this disk anymore
   * <p>
   * Locations are fetched using {@link StorageLocationSelectorStrategy}.
   */
  private File loadSegmentWithRetry(DataSegment segment) throws SegmentLoadingException
  {
    String segmentDir = getSegmentDir(segment);

    // Try the already reserved location. If location has been reserved outside, then we do not release the location
    // here and simply delete any downloaded files. That is, we revert anything we do in this function and nothing else.
    for (StorageLocation loc : locations) {
      if (loc.isReserved(segmentDir)) {
        File storageDir = loc.segmentDirectoryAsFile(segmentDir);
        boolean success = loadInLocationWithStartMarkerQuietly(loc, segment, storageDir, false);
        if (!success) {
          throw new SegmentLoadingException("Failed to load segment %s in reserved location [%s]", segment.getId(), loc.getPath().getAbsolutePath());
        }
        return storageDir;
      }
    }

    // No location was reserved so we try all the locations
    Iterator<StorageLocation> locationsIterator = strategy.getLocations();
    while (locationsIterator.hasNext()) {

      StorageLocation loc = locationsIterator.next();

      // storageDir is the file path corresponding to segment dir
      File storageDir = loc.reserve(segmentDir, segment);
      if (storageDir != null) {
        boolean success = loadInLocationWithStartMarkerQuietly(loc, segment, storageDir, true);
        if (success) {
          return storageDir;
        }
      }
    }
    throw new SegmentLoadingException("Failed to load segment %s in all locations.", segment.getId());
  }

  /**
   * A helper method over {@link #loadInLocationWithStartMarker(DataSegment, File)} that catches the {@link SegmentLoadingException}
   * and emits alerts.
   * @param loc - {@link StorageLocation} where segment is to be downloaded in.
   * @param segment - {@link DataSegment} to download
   * @param storageDir - {@link File} pointing to segment directory
   * @param releaseLocation - Whether to release the location in case of failures
   * @return - True if segment was downloaded successfully, false otherwise.
   */
  private boolean loadInLocationWithStartMarkerQuietly(StorageLocation loc, DataSegment segment, File storageDir, boolean releaseLocation)
  {
    try {
      loadInLocationWithStartMarker(segment, storageDir);
      return true;
    }
    catch (SegmentLoadingException e) {
      try {
        log.makeAlert(
            e,
            "Failed to load segment in current location [%s], try next location if any",
            loc.getPath().getAbsolutePath()
        ).addData("location", loc.getPath().getAbsolutePath()).emit();
      }
      finally {
        if (releaseLocation) {
          loc.removeSegmentDir(storageDir, segment);
        }
        cleanupCacheFiles(loc.getPath(), storageDir);
      }
    }
    return false;
  }

  private void loadInLocationWithStartMarker(DataSegment segment, File storageDir) throws SegmentLoadingException
  {
    // We use a marker to prevent the case where a segment is downloaded, but before the download completes,
    // the parent directories of the segment are removed
    final File downloadStartMarker = new File(storageDir, DOWNLOAD_START_MARKER_FILE_NAME);
    synchronized (directoryWriteRemoveLock) {
      if (!storageDir.mkdirs()) {
        log.debug("Unable to make parent file[%s]", storageDir);
      }
      try {
        if (!downloadStartMarker.createNewFile()) {
          throw new SegmentLoadingException("Was not able to create new download marker for [%s]", storageDir);
        }
      }
      catch (IOException e) {
        throw new SegmentLoadingException(e, "Unable to create marker file for [%s]", storageDir);
      }
    }
    loadInLocation(segment, storageDir);

    if (!downloadStartMarker.delete()) {
      throw new SegmentLoadingException("Unable to remove marker file for [%s]", storageDir);
    }
  }

  private void loadInLocation(DataSegment segment, File storageDir) throws SegmentLoadingException
  {
    // LoadSpec isn't materialized until here so that any system can interpret Segment without having to have all the
    // LoadSpec dependencies.
    final LoadSpec loadSpec = jsonMapper.convertValue(segment.getLoadSpec(), LoadSpec.class);
    final LoadSpec.LoadSpecResult result = loadSpec.loadSegment(storageDir);
    if (result.getSize() != segment.getSize()) {
      log.warn(
          "Segment [%s] is different than expected size. Expected [%d] found [%d]",
          segment.getId(),
          segment.getSize(),
          result.getSize()
      );
    }
  }

  @Override
  public boolean reserve(final DataSegment segment)
  {
    final ReferenceCountingLock lock = createOrGetLock(segment);
    synchronized (lock) {
      try {
        // May be the segment was already loaded [This check is required to account for restart scenarios]
        if (null != findStoragePathIfCached(segment)) {
          return true;
        }

        String storageDirStr = getSegmentDir(segment);

        // check if we already reserved the segment
        for (StorageLocation location : locations) {
          if (location.isReserved(storageDirStr)) {
            return true;
          }
        }

        // Not found in any location, reserve now
        for (Iterator<StorageLocation> it = strategy.getLocations(); it.hasNext(); ) {
          StorageLocation location = it.next();
          if (null != location.reserve(storageDirStr, segment)) {
            return true;
          }
        }
      }
      finally {
        unlock(segment, lock);
      }
    }

    return false;
  }

  @Override
  public boolean release(final DataSegment segment)
  {
    final ReferenceCountingLock lock = createOrGetLock(segment);
    synchronized (lock) {
      try {
        String storageDir = getSegmentDir(segment);

        // Release the first location encountered
        for (StorageLocation location : locations) {
          if (location.isReserved(storageDir)) {
            File localStorageDir = location.segmentDirectoryAsFile(storageDir);
            if (localStorageDir.exists()) {
              throw new ISE(
                  "Asking to release a location '%s' while the segment directory '%s' is present on disk. Any state on disk must be deleted before releasing",
                  location.getPath().getAbsolutePath(),
                  localStorageDir.getAbsolutePath()
              );
            }
            return location.release(storageDir, segment.getSize());
          }
        }
      }
      finally {
        unlock(segment, lock);
      }
    }

    return false;
  }

  @Override
  public void cleanup(DataSegment segment)
  {
    if (!config.isDeleteOnRemove()) {
      return;
    }

    final ReferenceCountingLock lock = createOrGetLock(segment);
    synchronized (lock) {
      try {
        File loc = findStoragePathIfCached(segment);

        if (loc == null) {
          log.warn("Asked to cleanup something[%s] that didn't exist.  Skipping.", segment.getId());
          return;
        }
        // If storageDir.mkdirs() success, but downloadStartMarker.createNewFile() failed,
        // in this case, findStorageLocationIfLoaded() will think segment is located in the failed storageDir which is actually not.
        // So we should always clean all possible locations here
        for (StorageLocation location : locations) {
          File localStorageDir = new File(location.getPath(), getSegmentDir(segment));
          if (localStorageDir.exists()) {
            // Druid creates folders of the form dataSource/interval/version/partitionNum.
            // We need to clean up all these directories if they are all empty.
            cleanupCacheFiles(location.getPath(), localStorageDir);
            location.removeSegmentDir(localStorageDir, segment);
          }
        }
      }
      finally {
        unlock(segment, lock);
      }
    }
  }

  private void cleanupCacheFiles(File baseFile, File cacheFile)
  {
    if (cacheFile.equals(baseFile)) {
      return;
    }

    synchronized (directoryWriteRemoveLock) {
      log.info("Deleting directory[%s]", cacheFile);
      try {
        FileUtils.deleteDirectory(cacheFile);
      }
      catch (Exception e) {
        log.error(e, "Unable to remove directory[%s]", cacheFile);
      }
    }

    File parent = cacheFile.getParentFile();
    if (parent != null) {
      File[] children = parent.listFiles();
      if (children == null || children.length == 0) {
        cleanupCacheFiles(baseFile, parent);
      }
    }
  }

  private ReferenceCountingLock createOrGetLock(DataSegment dataSegment)
  {
    return segmentLocks.compute(
        dataSegment,
        (segment, lock) -> {
          final ReferenceCountingLock nonNullLock;
          if (lock == null) {
            nonNullLock = new ReferenceCountingLock();
          } else {
            nonNullLock = lock;
          }
          nonNullLock.increment();
          return nonNullLock;
        }
    );
  }

  @SuppressWarnings("ObjectEquality")
  private void unlock(DataSegment dataSegment, ReferenceCountingLock lock)
  {
    segmentLocks.compute(
        dataSegment,
        (segment, existingLock) -> {
          if (existingLock == null) {
            throw new ISE("Lock has already been removed");
          } else if (existingLock != lock) {
            throw new ISE("Different lock instance");
          } else {
            if (existingLock.numReferences == 1) {
              return null;
            } else {
              existingLock.decrement();
              return existingLock;
            }
          }
        }
    );
  }

  @VisibleForTesting
  private static class ReferenceCountingLock
  {
    private int numReferences;

    private void increment()
    {
      ++numReferences;
    }

    private void decrement()
    {
      --numReferences;
    }
  }

  @VisibleForTesting
  public ConcurrentHashMap<DataSegment, ReferenceCountingLock> getSegmentLocks()
  {
    return segmentLocks;
  }

  @VisibleForTesting
  public List<StorageLocation> getLocations()
  {
    return locations;
  }
}