/*
 * Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
 * license agreements.  See the NOTICE file distributed with this work for
 * additional information regarding copyright ownership.  Crate licenses
 * this file to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.  You may
 * obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
 * License for the specific language governing permissions and limitations
 * under the License.
 *
 * However, if you have executed another commercial license agreement
 * with Crate these terms will supersede the license and you may use the
 * software solely pursuant to the terms of the relevant commercial agreement.
 */

package io.crate.operation.aggregation;

import com.google.common.collect.ImmutableList;
import io.crate.action.sql.SessionContext;
import io.crate.breaker.RamAccountingContext;
import io.crate.data.ArrayBucket;
import io.crate.data.Row;
import io.crate.execution.engine.aggregation.AggregationFunction;
import io.crate.execution.engine.collect.InputCollectExpression;
import io.crate.expression.symbol.Function;
import io.crate.expression.symbol.Literal;
import io.crate.expression.symbol.Symbol;
import io.crate.metadata.CoordinatorTxnCtx;
import io.crate.metadata.FunctionIdent;
import io.crate.metadata.Functions;
import io.crate.metadata.SearchPath;
import io.crate.test.integration.CrateUnitTest;
import io.crate.types.DataType;
import org.elasticsearch.Version;
import org.elasticsearch.common.breaker.CircuitBreaker;
import org.elasticsearch.common.breaker.NoopCircuitBreaker;
import org.elasticsearch.common.util.BigArrays;
import org.junit.Before;

import java.util.Arrays;
import java.util.List;

import static io.crate.testing.TestingHelpers.getFunctions;

public abstract class AggregationTest extends CrateUnitTest {

    protected static final RamAccountingContext ramAccountingContext =
        new RamAccountingContext("dummy", new NoopCircuitBreaker(CircuitBreaker.FIELDDATA));

    protected Functions functions;

    @Before
    public void prepare() throws Exception {
        functions = getFunctions();
    }

    public Object executeAggregation(String name, DataType dataType, Object[][] data) throws Exception {
        if (dataType == null) {
            return executeAggregation(name, dataType, data, ImmutableList.of());
        } else {
            return executeAggregation(name, dataType, data, ImmutableList.of(dataType));
        }
    }

    public Object executeAggregation(String name, DataType dataType, Object[][] data, List<DataType> argumentTypes) throws Exception {
        FunctionIdent fi;
        InputCollectExpression[] inputs;
        if (dataType != null) {
            fi = new FunctionIdent(name, argumentTypes);
            inputs = new InputCollectExpression[argumentTypes.size()];
            for (int i = 0; i < argumentTypes.size(); i++) {
                inputs[i] = new InputCollectExpression(i);
            }
        } else {
            fi = new FunctionIdent(name, ImmutableList.of());
            inputs = new InputCollectExpression[0];
        }
        AggregationFunction impl = (AggregationFunction) functions.getQualified(fi);
        Object state = impl.newState(ramAccountingContext, Version.CURRENT, BigArrays.NON_RECYCLING_INSTANCE);

        ArrayBucket bucket = new ArrayBucket(data);

        for (Row row : bucket) {
            for (InputCollectExpression i : inputs) {
                i.setNextRow(row);
            }
            state = impl.iterate(ramAccountingContext, state, inputs);

        }
        return impl.terminatePartial(ramAccountingContext, state);
    }

    protected Symbol normalize(String functionName, Object value, DataType type) {
        return normalize(functionName, Literal.of(type, value));
    }

    protected Symbol normalize(String functionName, Symbol... args) {
        List<Symbol> arguments = Arrays.asList(args);
        AggregationFunction function =
            (AggregationFunction) functions.get(null, functionName, arguments, SearchPath.pathWithPGCatalogAndDoc());
        return function.normalizeSymbol(
            new Function(function.info(), arguments),
            new CoordinatorTxnCtx(SessionContext.systemSessionContext()));
    }
}

@@ -1973,16 +1973,44 @@
 efore;%0A%0A
+import java.util.ArrayList;%0A
 import j
@@ -3550,23 +3550,68 @@
     
+List%3C
 Object
+%3E
  state
+s
  = 
+new ArrayList%3C%3E();%0A        states.add(
 impl
@@ -3692,18 +3692,18 @@
 NSTANCE)
+)
 ;
-%0A
 %0A       
@@ -3707,28 +3707,22 @@
     
-ArrayBucket bucket =
+for (Row row :
  new
@@ -3743,40 +3743,8 @@
 ata)
-;%0A%0A        for (Row row : bucket
 ) %7B%0A
@@ -3784,16 +3784,20 @@
 ession i
+nput
  : input
@@ -3818,16 +3818,20 @@
        i
+nput
 .setNext
@@ -3870,15 +3870,241 @@
     
-state =
+if (randomIntBetween(1, 4) == 1) %7B%0A                states.add(impl.newState(ramAccountingContext, Version.CURRENT, BigArrays.NON_RECYCLING_INSTANCE));%0A            %7D%0A            int idx = states.size() - 1;%0A            states.set(idx,
  imp
@@ -4140,16 +4140,26 @@
 t, state
+s.get(idx)
 , inputs
@@ -4159,18 +4159,193 @@
  inputs)
+)
 ;%0A
+        %7D%0A        Object state = states.get(0);%0A        for (int i = 1; i %3C states.size(); i++) %7B%0A            state = impl.reduce(ramAccountingContext, state, states.get(i));
 %0A       

Add occasional reduce steps to aggregation test framework