/*
 * RiskScape™ Copyright New Zealand Institute for Earth Science Limited
 * (Earth Sciences New Zealand) is distributed for research purposes only
 * under the terms of AGPLv3.
 *
 * RiskScape™ Copyright 2026 New Zealand Institute for Earth Science
 * Limited (Earth Sciences New Zealand). All rights reserved. Source code
 * available under the AGPLv3.
 * 
 * This program is free software: you can redistribute it and/or modify it under
 *  the terms of the GNU Affero General Public License as published by the Free
 *  Software Foundation, either version 3 of the License, or (at your option) any
 *  later version.
 * 
 * This program is distributed for RESEARCH PURPOSES ONLY, in the hope that it will
 * be useful for research and education initiatives.
 * 
 * If you are not a researcher, or you are a researcher who wishes to use this
 * program on terms other than AGPLv3 (including those who wish to restrict the
 * distribution of any source code created using this program), please contact:
 * https://riskscape.org.nz
 * 
 * This program is distributed WITHOUT ANY WARRANTY; without even the implied
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Affero General Public License for more details.  You should have received a copy
 * of the GNU Affero General Public License along with this program.  If not, see
 * <http://www.gnu.org/licenses/>.
 * 
 * By way of summary only, under the AGPLv3:
 *     • Permissions of this strongest copyleft license are conditioned
 *       on making available complete source code of licensed works and
 *       modifications, which include larger works using a licensed work,
 *       under the same license.
 *     • Copyright and license notices must be preserved.
 *     • Contributors provide an express grant of patent rights.
 *     • When a modified version is used to provide a service over a
 *       network, the complete source code of the modified version must be made
 *       available.
 */
package nz.org.riskscape.engine.function.agg;

import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import org.hamcrest.Matchers;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;

import lombok.Getter;
import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.FunctionArgument;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.rl.agg.Accumulator;
import nz.org.riskscape.engine.rl.agg.AggregationFunction;
import nz.org.riskscape.engine.rl.agg.RealizedAggregateExpression;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.TypeProblems;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ExpressionParser;
import nz.org.riskscape.rl.ast.Constant;
import nz.org.riskscape.rl.ast.FunctionCall;
import nz.org.riskscape.rl.ast.PropertyAccess;

public class AggregationFunctionAdapterTest {

  private RealizationContext context;
  private ArgumentCaptor<Type> inputTypeCaptor;
  private ArgumentCaptor<FunctionCall> functionCallCaptor;
  private AggregationFunction wrapped;
  private RealizedAggregateExpression realizedAggExpression;
  private Accumulator accumulator;
  private String processedValue;
  private FunctionCall functionCall;
  private RiskscapeFunction subject;
  private List<Type> argumentTypes;

  @Before
  public void setup() {
    context = mock(RealizationContext.class);
    argumentTypes = Arrays.asList(RSList.create(Types.INTEGER));

    inputTypeCaptor = ArgumentCaptor.forClass(Type.class);
    functionCallCaptor = ArgumentCaptor.forClass(FunctionCall.class);

    wrapped = mock(AggregationFunction.class);
    when(wrapped.getArguments()).thenReturn(ArgumentList.anonymous(argumentTypes));

    realizedAggExpression = mock(RealizedAggregateExpression.class);
    accumulator = mock(Accumulator.class);
    processedValue = "it works!";
    when(accumulator.process()).thenReturn(processedValue);
    when(realizedAggExpression.newAccumulator()).thenReturn(accumulator);
    when(realizedAggExpression.getResultType()).thenReturn(Types.INTEGER);

    functionCall = ExpressionParser.parseString("foo(['foo', 'bar', 'baz'])").isA(FunctionCall.class).get();
    subject = AggregationFunction.asFunction(wrapped);
  }

  @Test
  public void canWrapAGivenAggregationFunctionAndAccumulateAllMembersOfTheGivenList() {

    // realization succeeds...
    when(wrapped.realize(same(context), inputTypeCaptor.capture(), functionCallCaptor.capture()))
      .thenReturn(ResultOrProblems.of(realizedAggExpression));

    // it must be realizable to work
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    RiskscapeFunction realized = realizable.realize(context, functionCall, argumentTypes).orElse(null);
    assertNotNull(realized);

    // realized has the correct arg types
    assertEquals(Arrays.asList(RSList.create(Types.INTEGER)), realized.getArgumentTypes());

    // return type is preserved
    assertEquals(Types.INTEGER, realized.getReturnType());

    // remember which values get given to the accumulator
    List<Object> values = Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L);
    List<Object> capturedValues = new ArrayList<>();
    doAnswer(inv -> capturedValues.add((((Tuple)inv.getArgument(0)).clone()))).when(accumulator)
      .accumulate(any(Tuple.class));


    // the accumulator's process result is what gets returned
    assertSame(processedValue, realized.call(Arrays.asList(values)));

    // the implementation kludges together a struct an expression for use with realizing the agg function, check it
    // looks OK
    Struct capturedStruct = inputTypeCaptor.getValue().find(Struct.class).orElse(null);
    FunctionCall capturedFunctionCall = functionCallCaptor.getValue();
    assertEquals(
        capturedStruct.getMembers().get(0).getKey(),
        capturedFunctionCall.getArguments().get(0).getExpression().isA(PropertyAccess.class)
        .map(PropertyAccess::getAccessString).orElse(null)
    );

    // check that the accumulator was given all of the values
    assertEquals(
        values.stream().map(val -> Tuple.ofValues(capturedStruct, val)).collect(Collectors.toList()),
        capturedValues
    );

    values = Collections.emptyList();
    capturedValues.clear();
    // even no values means we get the accumulator's result - empty always false in this case
    assertSame(processedValue, realized.call(Arrays.asList(values)));
    assertTrue(capturedValues.isEmpty());

  }

  @Test
  public void canWrapAGivenAggregationFunctionThatRequiresExtraArgsAndAccumulateAllMembersOfTheGivenList() {

    subject = new AggregationFunctionAdapter(wrapped) {
      @Getter
        private final ArgumentList arguments = ArgumentList.fromArray(
            new FunctionArgument("items", RSList.LIST_ANYTHING),
            new FunctionArgument("a1", Types.ANYTHING));

        @Override
        public List<Type> getArgumentTypes() {
          return getArguments().getArgumentTypes();
        }
    };
    argumentTypes = Arrays.asList(RSList.create(Types.INTEGER), Types.TEXT);
    functionCall = ExpressionParser.parseString("foo(['foo', 'bar', 'baz'], 'bob')").isA(FunctionCall.class).get();

    // realization succeeds...
    when(wrapped.realize(same(context), inputTypeCaptor.capture(), functionCallCaptor.capture()))
      .thenReturn(ResultOrProblems.of(realizedAggExpression));

    // it must be realizable to work
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    RiskscapeFunction realized = realizable.realize(context, functionCall, argumentTypes).orElse(null);
    assertNotNull(realized);

    // realized has the correct arg types
    assertEquals(Arrays.asList(RSList.create(Types.INTEGER), Types.TEXT), realized.getArgumentTypes());

    // return type is preserved
    assertEquals(Types.INTEGER, realized.getReturnType());

    // remember which values get given to the accumulator
    List<Object> values = Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L);
    List<Object> capturedValues = new ArrayList<>();
    doAnswer(inv -> capturedValues.add((((Tuple)inv.getArgument(0)).clone()))).when(accumulator)
      .accumulate(any(Tuple.class));


    // the accumulator's process result is what gets returned
    assertSame(processedValue, realized.call(Arrays.asList(values, "bob")));

    // the implementation kludges together a struct an expression for use with realizing the agg function
    // and it's extra args, check it looks OK
    Struct capturedStruct = inputTypeCaptor.getValue().find(Struct.class).orElse(null);
    FunctionCall capturedFunctionCall = functionCallCaptor.getValue();
    assertEquals(
        capturedStruct.getMembers().get(0).getKey(),
        capturedFunctionCall.getArguments().get(0).getExpression().isA(PropertyAccess.class)
        .map(PropertyAccess::getAccessString).orElse(null)
    );

    // The second argument expression should be passed to the aggregation function, as the input expression
    assertEquals(
        Constant.string("bob"),
        capturedFunctionCall.getArguments().get(1).getExpression()
    );

    // check that the accumulator was given all of the values
    assertEquals(
        values.stream().map(val -> Tuple.ofValues(capturedStruct, val)).collect(Collectors.toList()),
        capturedValues
    );

    values = Collections.emptyList();
    capturedValues.clear();
    // even no values means we get the accumulator's result - empty always false in this case
    assertSame(processedValue, realized.call(Arrays.asList(values)));
    assertTrue(capturedValues.isEmpty());
  }

  @Test
  public void willReturnAnErrorWhenExtraArgsAreNotSupplied() {
    subject = new AggregationFunctionAdapter(wrapped) {
      @Getter
        private final ArgumentList arguments = ArgumentList.fromArray(
            new FunctionArgument("items", RSList.LIST_ANYTHING),
            new FunctionArgument("a1", Types.ANYTHING));

        @Override
        public List<Type> getArgumentTypes() {
          return getArguments().getArgumentTypes();
        }
    };

    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    assertThat(realizable.realize(context, functionCall, argumentTypes).getProblems(), Matchers.contains(
        Matchers.is(
            ArgsProblems.get().wrongNumber(2, 1)
        )));
  }

  @Test
  public void canWrapAGivenAggregationFunctionWithAnEmptyAccumulator() {
    // the accumulator is empty until it gets given some values
    List<Object> capturedValues = new ArrayList<>();
    when(accumulator.isEmpty()).thenAnswer(a -> capturedValues.isEmpty());

    // realization succeeds...
    when(wrapped.realize(same(context), inputTypeCaptor.capture(), functionCallCaptor.capture()))
      .thenReturn(ResultOrProblems.of(realizedAggExpression));

    // it must be realizable to work
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    RiskscapeFunction realized = realizable.realize(context, functionCall, argumentTypes).orElse(null);
    assertNotNull(realized);

    // because a new accumulator is empty, the return type can be null
    assertEquals(realized.getReturnType(), Nullable.INTEGER);

    // remember which values get given to the accumulator
    List<Object> values = Collections.emptyList();

    // no values means null
    assertNull(realized.call(Arrays.asList(values)));

    // now again with some values
    values = Arrays.asList(1L, 2L, 3L, 4L, 5L, 6L);
    // process should get called, since the accumulator won't be empty
    when(accumulator.process()).thenReturn(processedValue);

    doAnswer(inv -> capturedValues.add((((Tuple)inv.getArgument(0)).clone()))).when(accumulator)
      .accumulate(any(Tuple.class));

    // the accumulator's process result is what gets returned
    assertSame(processedValue, realized.call(Arrays.asList(values)));
  }

  @Test
  public void willReturnAnErrorFromTheUnderlyingFunction() throws Exception {
    // it must be realizable to work
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);

    Problem problem = Problem.error("bad ting");
    // realization fails...
    when(wrapped.realize(same(context), inputTypeCaptor.capture(), functionCallCaptor.capture()))
      .thenReturn(ResultOrProblems.failed(problem));

    ResultOrProblems<RiskscapeFunction> realized = realizable.realize(context, functionCall, argumentTypes);
    assertThat(
      realized.getProblems(),
      Matchers.contains(problem)
    );
  }

  @Test
  public void willReturnAnErrorIfNotExactlyOneArg() throws Exception {
    argumentTypes = Arrays.asList(RSList.create(Types.INTEGER), Types.TEXT);
    functionCall = ExpressionParser.parseString("foo(['foo', 'bar', 'baz'], 1)").isA(FunctionCall.class).get();
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    ResultOrProblems<RiskscapeFunction> realized = realizable.realize(context, functionCall, argumentTypes);

    assertThat(
        realized.getProblems(),
        Matchers.contains(ArgsProblems.get().wrongNumber(1, 2))
    );
  }

  @Test
  public void willReturnAnErrorIfGivenTypeNotAList() throws Exception {
    argumentTypes = Arrays.asList(Types.TEXT);
    functionCall = ExpressionParser.parseString("foo(1)").isA(FunctionCall.class).get();
    RealizableFunction realizable = subject.getRealizable().orElse(null);
    assertNotNull(realizable);
    ResultOrProblems<RiskscapeFunction> realized = realizable.realize(context, functionCall, argumentTypes);

    assertThat(
        realized.getProblems(),
        Matchers.contains(TypeProblems.get().mismatch(
            functionCall.getArguments().get(0), RSList.create(Nullable.ANYTHING), Types.TEXT)
        )
    );
  }

}
