package com.thinkbiganalytics.nifi.provenance.v2;

import com.google.common.collect.Lists;
import com.thinkbiganalytics.nifi.provenance.v2.writer.ProvenanceEventActiveMqWriter;
import com.thinkbiganalytics.util.SpringInitializer;
import org.apache.nifi.events.EventReporter;
import org.apache.nifi.provenance.PersistentProvenanceRepository;
import org.apache.nifi.provenance.ProvenanceEventBuilder;
import org.apache.nifi.provenance.ProvenanceEventRecord;
import org.apache.nifi.provenance.ProvenanceEventRepository;
import org.apache.nifi.provenance.lineage.ComputeLineageSubmission;
import org.apache.nifi.provenance.search.Query;
import org.apache.nifi.provenance.search.QuerySubmission;
import org.apache.nifi.provenance.search.SearchableField;

import java.io.IOException;
import java.util.List;
import java.util.Timer;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Custom Event Repository that will take the PovenanceEvents and write them to some output.
 */
public class ThinkbigProvenanceEventRepository implements ProvenanceEventRepository {

    private PersistentProvenanceRepository repository;


    private ProvenanceEventActiveMqWriter provenanceEventRecordWriter;

    private Timer timer;
    private AtomicLong eventCounter = new AtomicLong(0);

    public ThinkbigProvenanceEventRepository() {
        try {

            repository = new PersistentProvenanceRepository();
            provenanceEventRecordWriter = new ProvenanceEventActiveMqWriter();
            System.out.println(" new event recorder " + provenanceEventRecordWriter);
            //  provenanceEventRecordWriter.setEventId(repository.getMaxEventId());
            SpringInitializer.getInstance().initializeSpring();
            SpringApplicationListener listener = new SpringApplicationListener();
            listener.addObjectToAutowire("provenanceEventRecordWriter", provenanceEventRecordWriter);
            listener.autowire("provenanceEventRecordWriter", provenanceEventRecordWriter);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }


    @Override
    public void initialize(EventReporter eventReporter) throws IOException {
        repository.initialize(eventReporter);
    }

    @Override
    public ProvenanceEventBuilder eventBuilder() {
        return repository.eventBuilder();

    }

    @Override
    public void registerEvent(ProvenanceEventRecord provenanceEventRecord) {
        System.out.println("ThinkbigProvenanceEventRepository recording 1 single provenance event: " + provenanceEventRecord.getComponentId() + ", " + provenanceEventRecord.getComponentType());
        checkAndSetEventId();
        repository.registerEvent(provenanceEventRecord);
        provenanceEventRecordWriter.writeEvent(provenanceEventRecord);
    }

    /**
     * get the latest event id
     * if the getMaxEventId is null use -1 as a base id
     * the writer will do an incrementAndGet on it to start it off a 0 which is what Nifi starts with
     */
    private void checkAndSetEventId(){
        Long maxId = getMaxEventId();
        Long eventId = provenanceEventRecordWriter.checkAndSetMaxEventId(maxId == null ? -1L : maxId) ;
        if(maxId == null){
            System.out.println("Register Events ... getMaxEventId is null  using "+eventId+" from internal incrementer");
        }
    }

    @Override
    public void registerEvents(Iterable<ProvenanceEventRecord> iterable) {
        List<ProvenanceEventRecord> events = Lists.newArrayList(iterable);
        checkAndSetEventId();
        repository.registerEvents(events);
        if (events != null && events.size() > 0) {
            System.out.println("ThinkbigProvenanceEventRepository recording " + events.size() + " provenance events ");
            for (ProvenanceEventRecord event : events) {
                provenanceEventRecordWriter.writeEvent(event);
            }
        }
    }


    @Override
    public List<ProvenanceEventRecord> getEvents(long firstRecordId, int maxRecords) throws IOException {
        List<ProvenanceEventRecord> events = repository.getEvents(firstRecordId, maxRecords);
        return events;
    }

    @Override
    public Long getMaxEventId() {
        Long maxEventId = repository.getMaxEventId();
        return maxEventId;
    }

    @Override
    public QuerySubmission submitQuery(Query query) {
        return repository.submitQuery(query);
    }

    @Override
    public QuerySubmission retrieveQuerySubmission(String s) {
        return repository.retrieveQuerySubmission(s);
    }

    @Override
    public ComputeLineageSubmission submitLineageComputation(String s) {
        return repository.submitLineageComputation(s);
    }

    @Override
    public ComputeLineageSubmission retrieveLineageSubmission(String s) {
        return repository.retrieveLineageSubmission(s);
    }

    @Override
    public ProvenanceEventRecord getEvent(long l) throws IOException {
        return repository.getEvent(l);
    }

    @Override
    public ComputeLineageSubmission submitExpandParents(long eventId) {
        return repository.submitExpandParents(eventId);
    }

    @Override
    public ComputeLineageSubmission submitExpandChildren(long eventId) {
        return repository.submitExpandChildren(eventId);
    }

    @Override
    public void close() throws IOException {
        repository.close();
        timer.cancel();
    }

    @Override
    public List<SearchableField> getSearchableFields() {
        return repository.getSearchableFields();
    }

    @Override
    public List<SearchableField> getSearchableAttributes() {
        return repository.getSearchableAttributes();
    }


}


@@ -767,24 +767,54 @@
 OException;%0A
+import java.util.Collections;%0A
 import java.
@@ -894,16 +894,114 @@
 micLong;
+%0Aimport java.util.concurrent.locks.Lock;%0Aimport java.util.concurrent.locks.ReentrantReadWriteLock;
 %0A%0A/**%0A *
@@ -1238,16 +1238,61 @@
 itory;%0A%0A
+    private  Lock registerEventLock = null;%0A%0A
 %0A    pri
@@ -1365,134 +1365,129 @@
    p
-rivate Timer timer;%0A    private AtomicLong eventCounter = new AtomicLong(0);%0A%0A    public ThinkbigProvenanceEventRepository() %7B
+ublic ThinkbigProvenanceEventRepository() %7B%0A        this.registerEventLock = new ReentrantReadWriteLock(true).readLock();
 %0A   
@@ -1494,24 +1494,25 @@
      try %7B%0A%0A
+%0A
             
@@ -2792,133 +2792,44 @@
     
-checkAndSetEventId();%0A        repository.registerEvent(provenanceEventRecord);%0A        provenanceEventRecordWriter.writeEvent
+registerEvents(Collections.singleton
 (pro
@@ -2847,16 +2847,17 @@
 tRecord)
+)
 ;%0A    %7D%0A
@@ -3494,24 +3494,77 @@
 iterable) %7B%0A
+        this.registerEventLock.lock();%0A        try %7B%0A
         List
@@ -4006,32 +4006,107 @@
     %7D%0A        %7D%0A
+        %7D finally %7B%0A            this.registerEventLock.unlock();%0A        %7D%0A
     %7D%0A%0A%0A    @Ove
@@ -5558,32 +5558,8 @@
 ();%0A
-        timer.cancel();%0A
     

fix provenance id