package nl.esciencecenter.newsreader.hbase;

import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import org.apache.commons.compress.compressors.bzip2.BZip2CompressorInputStream;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.client.HTable;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.RetriesExhaustedWithDetailsException;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.nio.file.*;
import java.nio.file.attribute.BasicFileAttributes;

/**
 * Loads a directory of naf files into HBase
 */
@Parameters(separators="=", commandDescription="Load directory of naf files into HBase")
public class DirectoryLoader {

    @Parameter(names="--root", description="Root directory", required=true)
    private String rootDirectory;

    @Parameter(names="--table", description="HBase table name")
    private String tableName = "documents";

    @Parameter(names="--family", description="HBase column family name")
    private String familyName = "naf";

    @Parameter(names="--column", description="HBase column qualifier name")
    private String columnName = "annotated";

    @Parameter(names="--buffer", description="Write buffer size in bytes")
    private long writeBufferSize = 100 * 1024 * 1024;

    private HTable table;

    public String getRootDirectory() {
        return rootDirectory;
    }

    public void setRootDirectory(String rootDirectory) {
        this.rootDirectory = rootDirectory;
    }

    public String getTableName() {
        return tableName;
    }

    public String getFamilyName() {
        return familyName;
    }

    public String getColumnName() {
        return columnName;
    }

    public void setTable(HTable table) {
        this.table = table;
    }

    private void setup() throws IOException {
        if (table == null) {
            Configuration config = HBaseConfiguration.create();
            table = HTableHelper.getTable(tableName, familyName, config);
            // default write buffer size is 2097152, but a annotated documents are >1Mb,
            // so a flush/commit will be done each document
            // increase it so several documents are flushed
            table.setWriteBufferSize(writeBufferSize);
        }
    }

    public void run() throws IOException {
        setup();

        Path rootPath = Paths.get(rootDirectory);
        walkDirectory(rootPath);

        cleanup();
    }

    private void cleanup() throws IOException {
        table.close();
    }

    private void walkDirectory(Path rootPath) throws IOException {
        FileVisitor<Path> visitor = new SimpleFileVisitor<Path>() {
            @Override
            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                addDocument(file);
                return FileVisitResult.CONTINUE;
            }
        };

        System.out.println("Loading directory: " + rootDirectory);
        Files.walkFileTree(rootPath, visitor);
    }

    private void addDocument(Path file) {
        String filename = file.getFileName().toString().toLowerCase();
        if (filename.endsWith(".bz2")) {
            try {
                InputStream in = Files.newInputStream(file);
                BZip2CompressorInputStream bzIn = new BZip2CompressorInputStream(in);
                ByteArrayOutputStream bos = new ByteArrayOutputStream();
                int n = 0;
                byte[] buffer = new byte[1024];
                while (-1 != (n = bzIn.read(buffer))) {
                    bos.write(buffer, 0, n);
                }
                byte[] content = bos.toByteArray();
                addDocumentContent(filename, content);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        if (filename.endsWith(".naf")) {
            try {
                byte[] content = Files.readAllBytes(file);
                addDocumentContent(filename, content);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private void addDocumentContent(String filename, byte[] content) throws InterruptedIOException, RetriesExhaustedWithDetailsException {
        System.out.println("Inserting = " + filename);
        Put put = new Put(filename.getBytes());
        put.add(familyName.getBytes(), columnName.getBytes(), content);
        table.put(put);
    }
}

@@ -668,16 +668,55 @@
 ributes;
+%0Aimport java.security.PrivilegedAction;
 %0A%0A/**%0A *
@@ -879,16 +879,50 @@
 ryLoader
+ implements PrivilegedAction%3CLong%3E
  %7B%0A%0A    
@@ -1527,16 +1527,59 @@
 e table;
+%0A    private long addedDocumentCounter = 0;
 %0A%0A    pu
@@ -2526,32 +2526,46 @@
        %7D%0A    %7D%0A%0A
+    @Override%0A
     public void 
@@ -2563,40 +2563,39 @@
 lic 
-void
+Long
  run() 
-throws IOException %7B%0A
+%7B%0A        try %7B%0A    
     
@@ -2604,24 +2604,28 @@
   setup();%0A%0A
+    
         Path
@@ -2670,16 +2670,20 @@
         
+  
+  
 walkDire
@@ -2708,16 +2708,20 @@
         
+    
 cleanup(
@@ -2719,24 +2719,141 @@
  cleanup();%0A
+        %7D catch (IOException e) %7B%0A            System.err.println(e);%0A        %7D%0A%0A        return addedDocumentCounter;%0A
     %7D%0A%0A    p
@@ -4822,20 +4822,53 @@
 e.put(put);%0A
+        addedDocumentCounter ++;%0A
     %7D%0A%7D%0A

Make dumper a PrivilegedAction