import java.io.*;
import java.util.Properties;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileUtil;
import org.apache.hadoop.fs.LocalFileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hdfs.HdfsConfiguration;
import org.apache.hadoop.hdfs.MiniDFSCluster;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.server.MiniYARNCluster ;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fifo.FifoScheduler;
import org.apache.pig.ExecType;


public class MiniCluster {
    private MiniDFSCluster hdfsCluster = null;
    private MiniYARNCluster yarnCluster = null;
    private Configuration hdfsConf;
    private YarnConfiguration yarnConf;
    private String testName = null;
    private FileSystem fs = null;
    private File baseDir = null;

     public MiniCluster(String testName) {
        this.testName = testName;
    }

    public void setupMiniDfsAndYarnClusters() {
        try {
            final int noOfNodeManagers = 1;
            final int numLocalDirs = 1;
            final int numLogDirs = 1;
            baseDir = new File("./target/hdfs/" + testName).getAbsoluteFile();
            FileUtil.fullyDelete(baseDir);
            hdfsConf = new HdfsConfiguration();
            hdfsConf.set(MiniDFSCluster.HDFS_MINIDFS_BASEDIR, baseDir.getAbsolutePath());
            MiniDFSCluster.Builder builder = new MiniDFSCluster.Builder(hdfsConf);
            hdfsCluster = builder.build();
            String hdfsURI = "hdfs://localhost:" + hdfsCluster.getNameNodePort() + "/";
            fs = FileSystem.get(hdfsConf);

            yarnConf = new YarnConfiguration();
            hdfsConf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 64);
            hdfsConf.setClass(YarnConfiguration.RM_SCHEDULER,
                    FifoScheduler.class, ResourceScheduler.class);
            yarnCluster = new MiniYARNCluster(testName, noOfNodeManagers, numLocalDirs, numLogDirs);
            yarnCluster.init(hdfsConf);
            yarnCluster.start();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

//    abstract public ExecType getExecType();

    public void shutdownMiniDfsAndYarnClusters() {
        try {
            if (fs != null) {
                Path dataDir = new Path(baseDir.getParentFile().getParentFile().getParent());
                fs.delete(dataDir, true);
                File rootTestFile = new File(baseDir.getParentFile().getParentFile().getParent());
                String rootTestDir = rootTestFile.getAbsolutePath();
                Path rootTestPath = new Path(rootTestDir);
                LocalFileSystem localFileSystem = FileSystem.getLocal(hdfsConf);
                localFileSystem.delete(rootTestPath, true);
                fs.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        if (hdfsCluster != null) { hdfsCluster.shutdown(); }
        if (yarnCluster != null) { yarnCluster.stop(); }
    }
}