package com.thinkaurelius.titan.hadoop.formats.util;

import com.google.common.base.Preconditions;
import com.thinkaurelius.titan.diskstorage.Entry;
import com.thinkaurelius.titan.diskstorage.StaticBuffer;
import com.thinkaurelius.titan.hadoop.formats.util.input.TitanHadoopSetup;
import com.thinkaurelius.titan.util.system.ConfigurationUtil;
import org.apache.tinkerpop.gremlin.hadoop.structure.io.VertexWritable;
import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.mapreduce.*;

import java.io.IOException;
import java.util.List;
import java.util.function.Supplier;

import static com.thinkaurelius.titan.hadoop.formats.util.input.TitanHadoopSetupCommon.SETUP_CLASS_NAME;
import static com.thinkaurelius.titan.hadoop.formats.util.input.TitanHadoopSetupCommon.SETUP_PACKAGE_PREFIX;

public abstract class GiraphInputFormat extends InputFormat<NullWritable, VertexWritable> implements Configurable {

    private final InputFormat<StaticBuffer, Iterable<Entry>> inputFormat;
    private static RefCountedCloseable<TitanVertexDeserializer> refCounter;

    public GiraphInputFormat(InputFormat<StaticBuffer, Iterable<Entry>> inputFormat) {
        this.inputFormat = inputFormat;
        Preconditions.checkState(Configurable.class.isAssignableFrom(inputFormat.getClass()));
    }

    @Override
    public List<InputSplit> getSplits(JobContext context) throws IOException, InterruptedException {
        return inputFormat.getSplits(context);
    }

    @Override
    public RecordReader<NullWritable, VertexWritable> createRecordReader(InputSplit split, TaskAttemptContext context) throws IOException, InterruptedException {
        return new GiraphRecordReader(refCounter, inputFormat.createRecordReader(split, context));
    }

    public void setConf(Configuration conf, boolean reset) {
        ((Configurable)inputFormat).setConf(conf);

        if (refCounter != null && !reset) {
            return;
        }

        refCounter = new RefCountedCloseable<>(() -> {
            final String titanVersion = "current";

            String className = SETUP_PACKAGE_PREFIX + titanVersion + SETUP_CLASS_NAME;

            TitanHadoopSetup ts = ConfigurationUtil.instantiate(className, new Object[]{conf}, new Class[]{Configuration.class});

            return new TitanVertexDeserializer(ts);
        });
    }

    @Override
    public void setConf(Configuration conf) {
        setConf(conf, false);
    }

    @Override
    public Configuration getConf() {
        return ((Configurable)inputFormat).getConf();
    }

    public static class RefCountedCloseable<T extends AutoCloseable> {

        private T current;
        private long refCount;
        private final Supplier<T> builder;

        public RefCountedCloseable(Supplier<T> builder) {
            this.builder = builder;
        }

        public synchronized T acquire() {
            if (null == current) {
                Preconditions.checkState(0 == refCount);
                current = builder.get();
            }

            refCount++;

            return current;
        }

        public synchronized void release() throws Exception {
            Preconditions.checkState(null != current);
            Preconditions.checkState(0 < refCount);

            refCount--;

            if (0 == refCount) {
                current.close();
                current = null;
            }
        }
    }
}

@@ -47,16 +47,56 @@
 .util;%0A%0A
+import com.google.common.base.Function;%0A
 import c
@@ -672,44 +672,8 @@
 ist;
-%0Aimport java.util.function.Supplier;
 %0A%0Aim
@@ -1095,16 +1095,22 @@
  static 
+final 
 RefCount
@@ -1155,24 +1155,465 @@
 efCounter;%0A%0A
+    static %7B%0A        refCounter = new RefCountedCloseable%3C%3E((conf) -%3E %7B%0A            final String titanVersion = %22current%22;%0A%0A            String className = SETUP_PACKAGE_PREFIX + titanVersion + SETUP_CLASS_NAME;%0A%0A            TitanHadoopSetup ts = ConfigurationUtil.instantiate(%0A                    className, new Object%5B%5D%7B conf %7D, new Class%5B%5D%7B Configuration.class %7D);%0A%0A            return new TitanVertexDeserializer(ts);%0A        %7D);%0A    %7D%0A%0A%0A%0A
     public G
@@ -2272,32 +2272,46 @@
 ntext));%0A    %7D%0A%0A
+    @Override%0A
     public void 
@@ -2310,32 +2310,38 @@
 ic void setConf(
+final 
 Configuration co
@@ -2346,23 +2346,8 @@
 conf
-, boolean reset
 ) %7B%0A
@@ -2406,20 +2406,16 @@
         
-if (
 refCount
@@ -2420,545 +2420,37 @@
 nter
- != null && !reset) %7B%0A            return;%0A        %7D%0A%0A        refCounter = new RefCountedCloseable%3C%3E(() -%3E %7B%0A            final String titanVersion = %22current%22;%0A%0A            String className = SETUP_PACKAGE_PREFIX + titanVersion + SETUP_CLASS_NAME;%0A%0A            TitanHadoopSetup ts = ConfigurationUtil.instantiate(className, new Object%5B%5D%7Bconf%7D, new Class%5B%5D%7BConfiguration.class%7D);%0A%0A            return new TitanVertexDeserializer(ts);%0A        %7D);%0A    %7D%0A%0A    @Override%0A    public void setConf(Configuration conf) %7B%0A        setConf(conf, false
+.setBuilderConfiguration(conf
 );%0A 
@@ -2719,35 +2719,95 @@
 e final 
-Supplier%3CT%3E builder
+Function%3CConfiguration, T%3E builder;%0A        private Configuration configuration
 ;%0A%0A     
@@ -2840,17 +2840,32 @@
 ble(
-Supplier%3C
+Function%3CConfiguration, 
 T%3E b
@@ -2913,32 +2913,179 @@
 der;%0A        %7D%0A%0A
+        public synchronized void setBuilderConfiguration(Configuration configuration) %7B%0A            this.configuration = configuration;%0A        %7D%0A%0A
         public s
@@ -3240,12 +3240,27 @@
 der.
-get(
+apply(configuration
 );%0A 

Fix test failure caused by #1197