/*
 * 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.hadoop.hbase.snapshot;

import static org.apache.hadoop.util.ToolRunner.run;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Stream;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hbase.HBaseTestingUtility;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Admin;
import org.apache.hadoop.hbase.client.RegionInfo;
import org.apache.hadoop.hbase.mob.MobUtils;
import org.apache.hadoop.hbase.regionserver.StoreFileInfo;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.CommonFSUtils;
import org.apache.hadoop.hbase.util.EnvironmentEdgeManager;
import org.apache.hadoop.hbase.util.Pair;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.TestInfo;
import org.junit.jupiter.params.provider.Arguments;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos.SnapshotDescription;
import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos.SnapshotRegionManifest;

public class ExportSnapshotTestBase {

  private static final Logger LOG = LoggerFactory.getLogger(ExportSnapshotTestBase.class);

  protected final static HBaseTestingUtility TEST_UTIL = new HBaseTestingUtility();

  protected final static byte[] FAMILY = Bytes.toBytes("cf");

  protected String methodName;

  protected TableName tableName;

  protected String emptySnapshotName;

  protected String snapshotName;

  protected int tableNumFiles;

  protected Admin admin;

  protected boolean mob;

  public static Stream<Arguments> parameters() {
    return Stream.of(Arguments.of(false), Arguments.of(true));
  }

  protected ExportSnapshotTestBase(boolean mob) {
    this.mob = mob;
  }

  @AfterAll
  public static void tearDownAfterClass() throws Exception {
    TEST_UTIL.shutdownMiniMapReduceCluster();
    TEST_UTIL.shutdownMiniCluster();
  }

  /**
   * Create a table and take a snapshot of the table used by the export test.
   */
  @BeforeEach
  public void setUp(TestInfo testInfo) throws Exception {
    this.admin = TEST_UTIL.getAdmin();

    methodName = testInfo.getTestMethod().map(Method::getName).orElse(getClass().getSimpleName());

    String suffix = mob ? methodName + "-mob" : methodName;
    tableName = TableName.valueOf("testtb-" + suffix);
    snapshotName = "snaptb0-" + suffix;
    emptySnapshotName = "emptySnaptb0-" + suffix;

    // create Table
    createTable(this.tableName);

    // Take an empty snapshot
    admin.snapshot(emptySnapshotName, tableName);

    // Add some rows
    SnapshotTestingUtils.loadData(TEST_UTIL, tableName, 50, FAMILY);
    tableNumFiles = admin.getRegions(tableName).size();

    // take a snapshot
    admin.snapshot(snapshotName, tableName);
  }

  protected final void createTable(TableName tableName) throws Exception {
    if (mob) {
      MobSnapshotTestingUtils.createPreSplitMobTable(TEST_UTIL, tableName, 2, FAMILY);
    } else {
      SnapshotTestingUtils.createPreSplitTable(TEST_UTIL, tableName, 2, FAMILY);
    }
  }

  protected interface RegionPredicate {
    boolean evaluate(final RegionInfo regionInfo);
  }

  protected final RegionPredicate getBypassRegionPredicate() {
    if (mob) {
      return MobUtils::isMobRegionInfo;
    } else {
      return null;
    }
  }

  @AfterEach
  public void tearDown() throws Exception {
    for (TableName tn : TEST_UTIL.getAdmin().listTableNames()) {
      TEST_UTIL.deleteTable(tn);
    }
    SnapshotTestingUtils.deleteAllSnapshots(TEST_UTIL.getAdmin());
    SnapshotTestingUtils.deleteArchiveDirectory(TEST_UTIL);
  }

  protected final void testExportFileSystemState(final TableName tableName,
    final String snapshotName, final String targetName, int filesExpected) throws Exception {
    testExportFileSystemState(tableName, snapshotName, targetName, filesExpected,
      getHdfsDestinationDir(), false);
  }

  protected final void testExportFileSystemState(final TableName tableName,
    final String snapshotName, final String targetName, int filesExpected, Path copyDir,
    boolean overwrite) throws Exception {
    testExportFileSystemState(tableName, snapshotName, targetName, filesExpected, copyDir,
      overwrite, false);
  }

  protected final void testExportFileSystemState(final TableName tableName,
    final String snapshotName, final String targetName, int filesExpected, Path copyDir,
    boolean overwrite, boolean resetTtl) throws Exception {
    testExportFileSystemState(TEST_UTIL.getConfiguration(), tableName, snapshotName, targetName,
      filesExpected, TEST_UTIL.getDefaultRootDirPath(), copyDir, overwrite, resetTtl,
      getBypassRegionPredicate(), true, false);
  }

  /**
   * Creates destination directory, runs ExportSnapshot() tool, and runs some verifications.
   */
  protected static void testExportFileSystemState(final Configuration conf,
    final TableName tableName, final String snapshotName, final String targetName,
    final int filesExpected, final Path srcDir, Path rawTgtDir, final boolean overwrite,
    final boolean resetTtl, final RegionPredicate bypassregionPredicate, final boolean success,
    final boolean checksumVerify) throws Exception {
    FileSystem tgtFs = rawTgtDir.getFileSystem(conf);
    FileSystem srcFs = srcDir.getFileSystem(conf);
    Path tgtDir = rawTgtDir.makeQualified(tgtFs.getUri(), tgtFs.getWorkingDirectory());

    // Export Snapshot
    int res = runExportSnapshot(conf, snapshotName, targetName, srcDir, rawTgtDir, overwrite,
      resetTtl, checksumVerify, true, true);
    assertEquals(success ? 0 : 1, res, "success " + success + ", res=" + res);
    if (!success) {
      final Path targetDir = new Path(HConstants.SNAPSHOT_DIR_NAME, targetName);
      assertFalse(tgtFs.exists(new Path(tgtDir, targetDir)),
        tgtDir.toString() + " " + targetDir.toString());
      return;
    }
    LOG.info("Exported snapshot");

    // Verify File-System state
    FileStatus[] rootFiles = tgtFs.listStatus(tgtDir);
    assertEquals(filesExpected > 0 ? 2 : 1, rootFiles.length);
    for (FileStatus fileStatus : rootFiles) {
      String name = fileStatus.getPath().getName();
      assertTrue(fileStatus.isDirectory(), fileStatus.toString());
      assertTrue(name.equals(HConstants.SNAPSHOT_DIR_NAME)
        || name.equals(HConstants.HFILE_ARCHIVE_DIRECTORY), name.toString());
    }
    LOG.info("Verified filesystem state");

    // Compare the snapshot metadata and verify the hfiles
    final Path snapshotDir = new Path(HConstants.SNAPSHOT_DIR_NAME, snapshotName);
    final Path targetDir = new Path(HConstants.SNAPSHOT_DIR_NAME, targetName);
    verifySnapshotDir(srcFs, new Path(srcDir, snapshotDir), tgtFs, new Path(tgtDir, targetDir));
    Set<String> snapshotFiles =
      verifySnapshot(conf, tgtFs, tgtDir, tableName, targetName, resetTtl, bypassregionPredicate);
    assertEquals(filesExpected, snapshotFiles.size());
  }

  /*
   * verify if the snapshot folder on file-system 1 match the one on file-system 2
   */
  protected static void verifySnapshotDir(final FileSystem fs1, final Path root1,
    final FileSystem fs2, final Path root2) throws IOException {
    assertEquals(listFiles(fs1, root1, root1), listFiles(fs2, root2, root2));
  }

  /*
   * Verify if the files exists
   */
  protected static Set<String> verifySnapshot(final Configuration conf, final FileSystem fs,
    final Path rootDir, final TableName tableName, final String snapshotName,
    final boolean resetTtl, final RegionPredicate bypassregionPredicate) throws IOException {
    final Path exportedSnapshot =
      new Path(rootDir, new Path(HConstants.SNAPSHOT_DIR_NAME, snapshotName));
    final Set<String> snapshotFiles = new HashSet<>();
    final Path exportedArchive = new Path(rootDir, HConstants.HFILE_ARCHIVE_DIRECTORY);
    SnapshotReferenceUtil.visitReferencedFiles(conf, fs, exportedSnapshot,
      new SnapshotReferenceUtil.SnapshotVisitor() {
        @Override
        public void storeFile(final RegionInfo regionInfo, final String family,
          final SnapshotRegionManifest.StoreFile storeFile) throws IOException {
          if (bypassregionPredicate != null && bypassregionPredicate.evaluate(regionInfo)) {
            return;
          }

          if (!storeFile.hasReference() && !StoreFileInfo.isReference(storeFile.getName())) {
            String hfile = storeFile.getName();
            snapshotFiles.add(hfile);
            verifyNonEmptyFile(new Path(exportedArchive,
              new Path(CommonFSUtils.getTableDir(new Path("./"), tableName),
                new Path(regionInfo.getEncodedName(), new Path(family, hfile)))));
          } else {
            Pair<String, String> referredToRegionAndFile =
              StoreFileInfo.getReferredToRegionAndFile(storeFile.getName());
            String region = referredToRegionAndFile.getFirst();
            String hfile = referredToRegionAndFile.getSecond();
            snapshotFiles.add(hfile);
            verifyNonEmptyFile(new Path(exportedArchive,
              new Path(CommonFSUtils.getTableDir(new Path("./"), tableName),
                new Path(region, new Path(family, hfile)))));
          }
        }

        private void verifyNonEmptyFile(final Path path) throws IOException {
          assertTrue(fs.exists(path), path + " should exists");
          assertTrue(fs.getFileStatus(path).getLen() > 0, path + " should not be empty");
        }
      });

    // Verify Snapshot description
    SnapshotDescription desc = SnapshotDescriptionUtils.readSnapshotInfo(fs, exportedSnapshot);
    assertTrue(desc.getName().equals(snapshotName));
    assertTrue(desc.getTable().equals(tableName.getNameAsString()));
    if (resetTtl) {
      assertEquals(HConstants.DEFAULT_SNAPSHOT_TTL, desc.getTtl());
    }
    return snapshotFiles;
  }

  private static Set<String> listFiles(final FileSystem fs, final Path root, final Path dir)
    throws IOException {
    Set<String> files = new HashSet<>();
    LOG.debug("List files in {} in root {} at {}", fs, root, dir);
    int rootPrefix = root.makeQualified(fs.getUri(), fs.getWorkingDirectory()).toString().length();
    FileStatus[] list = CommonFSUtils.listStatus(fs, dir);
    if (list != null) {
      for (FileStatus fstat : list) {
        LOG.debug(Objects.toString(fstat.getPath()));
        if (fstat.isDirectory()) {
          files.addAll(listFiles(fs, root, fstat.getPath()));
        } else {
          files.add(fstat.getPath().makeQualified(fs.getUri(), fs.getWorkingDirectory()).toString()
            .substring(rootPrefix));
        }
      }
    }
    return files;
  }

  protected final Path getHdfsDestinationDir() {
    Path rootDir = TEST_UTIL.getHBaseCluster().getMaster().getMasterFileSystem().getRootDir();
    Path path =
      new Path(new Path(rootDir, "export-test"), "export-" + EnvironmentEdgeManager.currentTime());
    LOG.info("HDFS export destination path: " + path);
    return path;
  }

  protected static Path getLocalDestinationDir(HBaseTestingUtility htu) {
    Path path = htu.getDataTestDir("local-export-" + EnvironmentEdgeManager.currentTime());
    try {
      FileSystem fs = FileSystem.getLocal(htu.getConfiguration());
      LOG.info("Local export destination path: " + path);
      return path.makeQualified(fs.getUri(), fs.getWorkingDirectory());
    } catch (IOException ioe) {
      throw new RuntimeException(ioe);
    }
  }

  protected static void removeExportDir(final Path path) throws IOException {
    FileSystem fs = FileSystem.get(path.toUri(), new Configuration());
    fs.delete(path, true);
  }

  protected static int runExportSnapshot(final Configuration conf, final String sourceSnapshotName,
    final String targetSnapshotName, final Path srcDir, Path rawTgtDir, final boolean overwrite,
    final boolean resetTtl, final boolean checksumVerify, final boolean noSourceVerify,
    final boolean noTargetVerify) throws Exception {
    FileSystem tgtFs = rawTgtDir.getFileSystem(conf);
    FileSystem srcFs = srcDir.getFileSystem(conf);
    Path tgtDir = rawTgtDir.makeQualified(tgtFs.getUri(), tgtFs.getWorkingDirectory());
    LOG.info("tgtFsUri={}, tgtDir={}, rawTgtDir={}, srcFsUri={}, srcDir={}", tgtFs.getUri(), tgtDir,
      rawTgtDir, srcFs.getUri(), srcDir);
    List<String> opts = new ArrayList<>();
    opts.add("--snapshot");
    opts.add(sourceSnapshotName);
    opts.add("--copy-to");
    opts.add(tgtDir.toString());
    if (!targetSnapshotName.equals(sourceSnapshotName)) {
      opts.add("--target");
      opts.add(targetSnapshotName);
    }
    if (overwrite) {
      opts.add("--overwrite");
    }
    if (resetTtl) {
      opts.add("--reset-ttl");
    }
    if (!checksumVerify) {
      opts.add("--no-checksum-verify");
    }
    if (!noSourceVerify) {
      opts.add("--no-source-verify");
    }
    if (!noTargetVerify) {
      opts.add("--no-target-verify");
    }

    // Export Snapshot
    return run(conf, new ExportSnapshot(), opts.toArray(new String[opts.size()]));
  }
}
