package com.marklogic.client.test.dataservices;

import com.fasterxml.jackson.databind.node.ObjectNode;
import com.marklogic.client.dataservices.IOEndpoint;
import com.marklogic.client.dataservices.OutputEndpoint;
import com.marklogic.client.document.JSONDocumentManager;
import com.marklogic.client.io.DocumentMetadataHandle;
import com.marklogic.client.io.JacksonHandle;
import com.marklogic.client.io.StringHandle;
import com.marklogic.client.query.DeleteQueryDefinition;
import com.marklogic.client.query.QueryManager;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

import java.io.ByteArrayInputStream;
import java.util.ArrayList;
import java.util.List;

import static org.junit.Assert.assertTrue;

public class BulkIOOutputCallerTest {
    static ObjectNode apiObj;
    static String apiName = "bulkIOOutputCaller.api";
    static String scriptPath;
    static String apiPath;
    static JSONDocumentManager docMgr;
    static List<String> list = new ArrayList<>();

    private static String collectionName_1 = "bulkOutputTest_1";
    private static String collectionName_2 = "bulkOutputTest_2";

    @Rule
    public ExpectedException expectedException = ExpectedException.none();

    @BeforeClass
    public static void setup() throws Exception {
        docMgr = IOTestUtil.db.newJSONDocumentManager();
        apiObj = IOTestUtil.readApi(apiName);
        scriptPath = IOTestUtil.getScriptPath(apiObj);
        apiPath = IOTestUtil.getApiPath(scriptPath);
        IOTestUtil.load(apiName, apiObj, scriptPath, apiPath);
        writeDocuments(10,20, collectionName_1);
        writeDocuments(30,40, collectionName_2);
    }

    @Test
    public void bulkOutputCallerTestWithMultipleCallContexts() throws Exception {

        String endpointState1 = "{\"next\":"+1+"}";
        String workUnit1      = "{\"max\":6,\"limit\":5,\"collection\":\"bulkOutputTest_1\"}";

        String endpointState2 = "{\"next\":"+1+"}";
        String workUnit2      =  "{\"max\":6,\"limit\":5,\"collection\":\"bulkOutputTest_2\"}";
        OutputEndpoint endpoint = OutputEndpoint.on(IOTestUtil.db, new JacksonHandle(apiObj));
        IOEndpoint.CallContext[] callContextArray = {endpoint.newCallContext()
                .withEndpointState(new ByteArrayInputStream(endpointState2.getBytes()))
                .withWorkUnit(new ByteArrayInputStream(workUnit2.getBytes())), endpoint.newCallContext()
                .withEndpointState(new ByteArrayInputStream(endpointState1.getBytes()))
                .withWorkUnit(new ByteArrayInputStream(workUnit1.getBytes()))};
        OutputEndpoint.BulkOutputCaller bulkCaller = endpoint.bulkCaller(callContextArray);
        class Output {
            int counter =0;
        }
        Output output = new Output();
        bulkCaller.setOutputListener(i-> {
            try {
                assertTrue(list.contains(IOTestUtil.mapper.readValue(i, ObjectNode.class).toString()));
                output.counter++;
            } catch (Exception e) {
                e.printStackTrace();
            }
        });

        bulkCaller.awaitCompletion();
        assertTrue(output.counter == 20);
    }

    @AfterClass
    public static void cleanup() {

        QueryManager queryMgr = IOTestUtil.db.newQueryManager();
        DeleteQueryDefinition deletedef = queryMgr.newDeleteDefinition();
        deletedef.setCollections(collectionName_1);
        queryMgr.delete(deletedef);
        deletedef.setCollections(collectionName_2);
        queryMgr.delete(deletedef);

    }

    private static void writeDocuments(int startCount,int endCount, String collection) {
        JSONDocumentManager manager = IOTestUtil.db.newJSONDocumentManager();
        DocumentMetadataHandle metadata = new DocumentMetadataHandle().withCollections(collection);

        for(int i=startCount;i<endCount;i++) {
            StringHandle data = new StringHandle("{\"docNum\":"+i+",\"docName\":\"doc"+i+"\"}");
            manager.write("/test/"+collection+"/"+i+".json", metadata, data);
            list.add(data.toString());
        }
    }
}

@@ -722,42 +722,185 @@
 til.
-ArrayList;%0Aimport java.util.List;%0A
+HashSet;%0Aimport java.util.Set;%0Aimport java.util.concurrent.ConcurrentSkipListSet;%0Aimport java.util.concurrent.atomic.AtomicBoolean;%0A%0Aimport static org.junit.Assert.assertEquals;
 %0Aimp
@@ -1169,19 +1169,18 @@
  static 
-Lis
+Se
 t%3CString
@@ -1185,27 +1185,29 @@
 ng%3E 
-list = new ArrayLis
+expected = new HashSe
 t%3C%3E(
@@ -2901,96 +2901,192 @@
     
-class Output %7B%0A            int counter =0;%0A        %7D%0A        Output output = new Output(
+Set%3CString%3E actual = new ConcurrentSkipListSet%3C%3E();%0A        final AtomicBoolean duplicated = new AtomicBoolean(false);%0A        final AtomicBoolean exceptional = new AtomicBoolean(false
 );%0A 
@@ -3121,17 +3121,23 @@
 istener(
-i
+output 
 -%3E %7B%0A   
@@ -3171,33 +3171,28 @@
     
-assertTrue(list.contains(
+String serialized = 
 IOTe
@@ -3215,17 +3215,22 @@
 adValue(
-i
+output
 , Object
@@ -3255,45 +3255,206 @@
 ng()
-));%0A                output.counter++;
+;%0A                if (actual.contains(serialized)) %7B%0A                    duplicated.compareAndSet(false, true);%0A                %7D else %7B%0A                    actual.add(serialized);%0A                %7D
 %0A   
@@ -3519,24 +3519,80 @@
 ackTrace();%0A
+                exceptional.compareAndSet(false, true);%0A
             
@@ -3662,33 +3662,278 @@
 sert
-True(
+Equals(%22exceptions on calls%22, false, exceptional.get());%0A        assertEquals(%22duplicate output%22, false, duplicated.get());%0A        assertEquals(%22unexpected 
 output
-.
+ 
 count
-er == 20
+%22, expected.size(), actual.size());%0A        assertEquals(%22unexpected output values%22, expected, actual
 );%0A 
@@ -4819,12 +4819,16 @@
     
-list
+expected
 .add

improve tolerance of concurrency for unit test