package org.embulk.input;

import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

import com.google.common.base.Optional;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.*;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;
import org.embulk.config.CommitReport;
import org.embulk.config.Config;
import org.embulk.config.ConfigDiff;
import org.embulk.config.ConfigSource;
import org.embulk.config.Task;
import org.embulk.config.TaskSource;
import org.embulk.exec.PooledBufferAllocator;
import org.embulk.spi.*;
import org.embulk.spi.type.Type;
import org.embulk.spi.type.Types;

public class HBaseInputPlugin
        implements InputPlugin
{
    public interface PluginTask
            extends Task
    {
        @Config("zookeeper")
        public String getZooKeeperHosts();

        @Config("table")
        public String getTableName();

        @Config("start_row")
        public Optional<String> getStartRow();

        @Config("stop_row")
        public Optional<String> getStopRow();

        @Config("columns")
        public SchemaConfig getColumns();
    }

    byte[] getStartRowKey(PluginTask task) {
        if (task.getStartRow().isPresent()) {
            return Bytes.toBytes(task.getStartRow().get());
        }
        return HConstants.EMPTY_START_ROW;
    }

    byte[] getStopRowKey(PluginTask task) {
        if (task.getStopRow().isPresent()) {
            return Bytes.toBytes(task.getStopRow().get());
        }
        return HConstants.EMPTY_END_ROW;
    }

    HConnection connect(PluginTask task) throws IOException {
        Configuration config = HBaseConfiguration.create();
        config.set("hbase.zookeeper.quorum", task.getZooKeeperHosts());
        return HConnectionManager.createConnection(config);
    }

    List<byte[]> regionIds;

    @Override
    public ConfigDiff transaction(ConfigSource config,
            InputPlugin.Control control)
    {
        PluginTask task = config.loadConfig(PluginTask.class);
        Schema schema = task.getColumns().toSchema();
        try (HConnection connection = connect(task)) {
            HTable table = (HTable) connection.getTable(task.getTableName());
            List<HRegionLocation> locations = table.getRegionsInRange(getStartRowKey(task), getStopRowKey(task));
            regionIds = locations.stream().map(l -> l.getRegionInfo().getRegionName()).collect(Collectors.toList());
            return resume(task.dump(), schema, regionIds.size(), control);
        } catch (IOException e) {
            return null;
        }
    }

    @Override
    public ConfigDiff resume(TaskSource taskSource,
            Schema schema, int taskCount,
            InputPlugin.Control control)
    {
        control.run(taskSource, schema, taskCount);
        return Exec.newConfigDiff();
    }

    @Override
    public void cleanup(TaskSource taskSource,
            Schema schema, int taskCount,
            List<CommitReport> successCommitReports)
    {
    }

    @Override
    public CommitReport run(TaskSource taskSource,
            Schema schema, int taskIndex,
            PageOutput output)
    {
        PluginTask task = taskSource.loadTask(PluginTask.class);
        try (HConnection connection = connect(task);
             PageBuilder pageBuilder = new PageBuilder(new PooledBufferAllocator(), schema, output);
             HTableInterface table = connection.getTable(task.getTableName())) {
            Scan scan = new Scan();

            // row range
            HRegionInfo regionInfo = connection.locateRegion(regionIds.get(taskIndex)).getRegionInfo();
            byte[] startRow = getStartRowKey(task);
            if (Bytes.compareTo(startRow, regionInfo.getStartKey()) < 0) {
                startRow = regionInfo.getStartKey();
            }
            byte[] stopRow = getStopRowKey(task);
            if (Bytes.compareTo(regionInfo.getEndKey(), stopRow) < 0) {
                stopRow = regionInfo.getEndKey();
            }
            scan.setStartRow(startRow);
            scan.setStopRow(stopRow);

            for (Result result : table.getScanner(scan)) {
                for (Column column : schema.getColumns()) {
                    String[] token = column.getName().split(":", 2);
                    byte[] family = Bytes.toBytes(token[0]);
                    byte[] qualifier = Bytes.toBytes(token[1]);
                    byte[] value;
                    if (result.containsColumn(family, qualifier)) {
                        value = CellUtil.cloneValue(result.getColumnLatestCell(family, qualifier));
                        Type type = column.getType();
                        if (type.equals(Types.BOOLEAN)) {
                            pageBuilder.setBoolean(column, Bytes.toBoolean(value));
                        }
                        else if (type.equals(Types.DOUBLE)) {
                            pageBuilder.setDouble(column, Bytes.toDouble(value));
                        }
                        else if (type.equals(Types.LONG)) {
                            pageBuilder.setLong(column, Bytes.toLong(value));
                        }
                        else if (type.equals(Types.STRING)) {
                            pageBuilder.setString(column, Bytes.toString(value));
                        }
                    }
                }
                pageBuilder.addRecord();
            }
            pageBuilder.finish();
            return Exec.newCommitReport().setAll(taskSource);
        } catch (IOException e) {
            Exec.getLogger(this.getClass()).error("failed to connect to HBase", e);
            return null;
        }
    }

    @Override
    public ConfigDiff guess(ConfigSource config)
    {
        // TODO
        throw new UnsupportedOperationException("'hbase' input plugin does not support guessing.");
        //ConfigDiff diff = Exec.newConfigDiff();
        //diff.set("property1", "value");
        //return diff;
    }
}

@@ -24,96 +24,8 @@
 t;%0A%0A
-import java.io.IOException;%0Aimport java.util.List;%0Aimport java.util.stream.Collectors;%0A%0A
 impo
@@ -60,16 +60,16 @@
 tional;%0A
+
 import o
@@ -252,243 +252,9 @@
 fig.
-CommitReport;%0Aimport org.embulk.config.Config;%0Aimport org.embulk.config.ConfigDiff;%0Aimport org.embulk.config.ConfigSource;%0Aimport org.embulk.config.Task;%0Aimport org.embulk.config.TaskSource;%0Aimport org.embulk.exec.PooledBufferAllocator
+*
 ;%0Aim
@@ -344,16 +344,104 @@
 Types;%0A%0A
+import java.io.IOException;%0Aimport java.util.List;%0Aimport java.util.stream.Collectors;%0A%0A
 public c
@@ -720,24 +720,55 @@
 start_row%22)%0A
+        @ConfigDefault(%22null%22)%0A
         publ
@@ -827,24 +827,55 @@
 %22stop_row%22)%0A
+        @ConfigDefault(%22null%22)%0A
         publ
@@ -936,16 +936,16 @@
 lumns%22)%0A
-
         
@@ -977,16 +977,92 @@
 lumns();
+%0A%0A        @ConfigInject%0A        public BufferAllocator getBufferAllocator();
 %0A    %7D%0A%0A
@@ -3257,18 +3257,16 @@
 der(
-new Pooled
+task.get
 Buff

Fix configs