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

import org.apache.commons.lang3.StringUtils;
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.fs.RawLocalFileSystem;
import org.apache.hadoop.fs.permission.FsPermission;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.zeppelin.conf.ZeppelinConfiguration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.file.attribute.PosixFilePermission;
import java.nio.file.attribute.PosixFilePermissions;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;


/**
 * Hadoop FileSystem wrapper. Support both secure and no-secure mode
 */
public class FileSystemStorage {

  private static final Logger LOGGER = LoggerFactory.getLogger(FileSystemStorage.class);
  private static final String S3A = "s3a";
  private static final String FS_DEFAULTFS = "fs.defaultFS";

  // only do UserGroupInformation.loginUserFromKeytab one time, otherwise you will still get
  // your ticket expired.
  static {
    if (UserGroupInformation.isSecurityEnabled()) {
      String keytab = ZeppelinConfiguration.getStaticString(
          ZeppelinConfiguration.ConfVars.ZEPPELIN_SERVER_KERBEROS_KEYTAB);
      String principal = ZeppelinConfiguration.getStaticString(
          ZeppelinConfiguration.ConfVars.ZEPPELIN_SERVER_KERBEROS_PRINCIPAL);
      if (StringUtils.isBlank(keytab) || StringUtils.isBlank(principal)) {
        throw new RuntimeException("keytab and principal can not be empty, keytab: " + keytab
            + ", principal: " + principal);
      }
      try {
        UserGroupInformation.loginUserFromKeytab(principal, keytab);
      } catch (IOException e) {
        throw new RuntimeException("Fail to login via keytab:" + keytab +
            ", principal:" + principal, e);
      }
    }
  }

  private ZeppelinConfiguration zConf;
  private Configuration hadoopConf;
  private boolean isSecurityEnabled;
  private FileSystem fs;

  public FileSystemStorage(ZeppelinConfiguration zConf, String path) throws IOException {
    this.zConf = zConf;
    this.hadoopConf = new Configuration();
    URI zepConfigURI;
    URI defaultFSURI;

    try {
      zepConfigURI = new URI(path);
    } catch (URISyntaxException e) {
      LOGGER.error("Failed to get Zeppelin config URI");
      throw new IOException(e);
    }
    // disable checksum for local file system. because interpreter.json may be updated by
    // non-hadoop filesystem api
    // disable caching for file:// scheme to avoid getting LocalFS which does CRC checks.

    this.hadoopConf.setBoolean("fs.file.impl.disable.cache", true);
    String defaultFS = this.hadoopConf.get(FS_DEFAULTFS);
    try {
      defaultFSURI = new URI(defaultFS);
    } catch (URISyntaxException e) {
      LOGGER.error("Failed to get defaultFS URI");
      throw new IOException(e);
    }

    // to check whether underlying fileSystemStorage is S3A or not
    if (!isS3AFileSystem(defaultFSURI, zepConfigURI)) {
      this.hadoopConf.set("fs.file.impl", RawLocalFileSystem.class.getName());
    }

    this.isSecurityEnabled = UserGroupInformation.isSecurityEnabled();

    this.fs = FileSystem.get(zepConfigURI, this.hadoopConf);
  }

  public boolean isS3AFileSystem(URI defaultFSURI, URI zepConfigURI) {
    return defaultFSURI.getScheme().equals(S3A)
      || (StringUtils.isNotEmpty(zepConfigURI.getScheme())
      && zepConfigURI.getScheme().equals(S3A));
  }

  public FileSystem getFs() {
    return fs;
  }

  public Path makeQualified(Path path) {
    return fs.makeQualified(path);
  }

  public boolean exists(final Path path) throws IOException {
    return callHdfsOperation(new HdfsOperation<Boolean>() {

      @Override
      public Boolean call() throws IOException {
        return fs.exists(path);
      }
    });
  }

  public void tryMkDir(final Path dir) throws IOException {
    callHdfsOperation(new HdfsOperation<Void>() {
      @Override
      public Void call() throws IOException {
        if (!fs.exists(dir)) {
          fs.mkdirs(dir);
          LOGGER.info("Create dir {} in hdfs", dir);
        }
        if (fs.getFileStatus(dir).isFile()) {
          throw new IOException(dir.toString() + " is file instead of directory, please remove " +
              "it or specify another directory");
        }
        fs.mkdirs(dir);
        return null;
      }
    });
  }

  public List<Path> list(final Path path) throws IOException {
    return callHdfsOperation(new HdfsOperation<List<Path>>() {
      @Override
      public List<Path> call() throws IOException {
        List<Path> paths = new ArrayList<>();
        for (FileStatus status : fs.globStatus(path)) {
          paths.add(status.getPath());
        }
        return paths;
      }
    });
  }

  // recursive search path, (TODO zjffdu, list folder in sub folder on demand, instead of load all
  // data when zeppelin server start)
  public List<Path> listAll(final Path path) throws IOException {
    return callHdfsOperation(new HdfsOperation<List<Path>>() {
      @Override
      public List<Path> call() throws IOException {
        List<Path> paths = new ArrayList<>();
        collectNoteFiles(path, paths);
        return paths;
      }

      private void collectNoteFiles(Path folder, List<Path> noteFiles) throws IOException {
        FileStatus[] paths = fs.listStatus(folder);
        for (FileStatus path : paths) {
          if (path.isDirectory()) {
            collectNoteFiles(path.getPath(), noteFiles);
          } else {
            if (path.getPath().getName().endsWith(".zpln")) {
              noteFiles.add(path.getPath());
            } else {
              LOGGER.warn("Unknown file: {}", path.getPath());
            }
          }
        }
      }
    });
  }

  public boolean delete(final Path path) throws IOException {
    return callHdfsOperation(new HdfsOperation<Boolean>() {
      @Override
      public Boolean call() throws IOException {
        return fs.delete(path, true);
      }
    });
  }

  public String readFile(final Path file) throws IOException {
    return callHdfsOperation(new HdfsOperation<String>() {
      @Override
      public String call() throws IOException {
        LOGGER.debug("Read from file: {}", file);
        ByteArrayOutputStream noteBytes = new ByteArrayOutputStream();
        IOUtils.copyBytes(fs.open(file), noteBytes, hadoopConf);
        return noteBytes.toString(zConf.getString(ZeppelinConfiguration.ConfVars.ZEPPELIN_ENCODING));
      }
    });
  }

  public void writeFile(final String content, final Path file, boolean writeTempFileFirst)
      throws IOException {
      writeFile(content, file, writeTempFileFirst, null);
  }

  public void writeFile(final String content, final Path file, boolean writeTempFileFirst, Set<PosixFilePermission> permissions)
      throws IOException {
    FsPermission fsPermission;
    if (permissions == null || permissions.isEmpty()) {
      fsPermission = FsPermission.getFileDefault();
    } else {
      // FsPermission expects a 10-character string because of the leading
      // directory indicator, i.e. "drwx------". The JDK toString method returns
      // a 9-character string, so prepend a leading character.
      fsPermission = FsPermission.valueOf("-" + PosixFilePermissions.toString(permissions));
    }
    callHdfsOperation(new HdfsOperation<Void>() {
      @Override
      public Void call() throws IOException {
        InputStream in = new ByteArrayInputStream(content.getBytes(
            zConf.getString(ZeppelinConfiguration.ConfVars.ZEPPELIN_ENCODING)));
        Path tmpFile = new Path(file.toString() + ".tmp");
        IOUtils.copyBytes(in, fs.create(tmpFile), hadoopConf);
        fs.setPermission(tmpFile, fsPermission);
        fs.delete(file, true);
        fs.rename(tmpFile, file);
        return null;
      }
    });
  }

  public void move(Path src, Path dest) throws IOException {
    callHdfsOperation(() -> {
      fs.rename(src, dest);
      return null;
    });
  }

  private interface HdfsOperation<T> {
    T call() throws IOException;
  }

  public synchronized <T> T callHdfsOperation(final HdfsOperation<T> func) throws IOException {
    if (isSecurityEnabled) {
      try {
        return UserGroupInformation.getCurrentUser().doAs(new PrivilegedExceptionAction<T>() {
          @Override
          public T run() throws Exception {
            return func.call();
          }
        });
      } catch (InterruptedException e) {
        throw new IOException(e);
      }
    } else {
      return func.call();
    }
  }

}
