package com.sony.etl;

import java.util.Map;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.hadoop.hbase.client.Mutation;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.util.Bytes;

import com.google.api.services.bigquery.model.TableRow;
import com.google.cloud.bigtable.beam.CloudBigtableIO;
import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
import com.google.gson.Gson;

public class BigQueryBigtableTransfer {
	private static final byte[] FAMILY = Bytes.toBytes("cf");

	  static final DoFn<TableRow, Mutation> MUTATION_TRANSFORM = new DoFn<TableRow, Mutation>() {
	    private static final long serialVersionUID = 1L;

	    @ProcessElement
	    public void processElement(DoFn<TableRow, Mutation>.ProcessContext c) throws Exception {
	      Gson gson = new Gson();
	      TableRow row = c.element();

	      String rowKey = row.get("newRowkey").toString();

	      Put p = new Put(rowKey.getBytes());

	      for (Map.Entry<String, Object> field : row.entrySet()) {
	        p.addColumn(FAMILY, field.getKey().getBytes(), ((String) field.getValue()).getBytes());
	      }

	      p.addColumn("json".getBytes(), "data".getBytes(), gson.toJson(row, TableRow.class).getBytes());

	      c.output(p);

	    }
	  };
	  
	  public static void main(String[] args) {
		    // CloudBigtableOptions is one way to retrieve the options.  It's not required.
		  CloudBigtableOptions options =
		        PipelineOptionsFactory.fromArgs(args).withValidation().as(CloudBigtableOptions.class);
		    
		    // CloudBigtableTableConfiguration contains the project, instance and table to connect to.
		    CloudBigtableTableConfiguration config =
		        new CloudBigtableTableConfiguration.Builder()
		        .withProjectId(options.getBigtableProjectId())
		        .withInstanceId(options.getBigtableInstanceId())
		        .withTableId(options.getBigtableTableId())
		        .build();

		    Pipeline p = Pipeline.create(options);
		    /*String query = null;
		    if (options.getBqQuery().isAccessible() && options.getReportDate().isAccessible()) {
		    	query = options.getBqQuery().get()  + 
	        		" WHERE date = DATE('" + options.getReportDate().get() + "')";
		    	System.out.println("Query == " + query);
		    }*/
		    p
		        .apply(BigQueryIO.read().withTemplateCompatibility().fromQuery(options.getBqQuery()).withoutValidation()
		            .usingStandardSql())
		        .apply(ParDo.of(MUTATION_TRANSFORM))
		        .apply(CloudBigtableIO.writeToTable(config));

		    p.run().waitUntilFinish();

		  }
}