/*
 * 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;


import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;

import org.junit.Test;
import org.mockito.Mockito;

import com.google.common.collect.Lists;

import nz.org.riskscape.dsl.Token;
import nz.org.riskscape.engine.function.FunctionResolver;
import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.function.JavaFunction;
import nz.org.riskscape.engine.function.NullSafeFunction;
import nz.org.riskscape.engine.function.OverloadedFunction;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.resource.Resource;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.types.CoverageType;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.engine.types.WithinRange;
import nz.org.riskscape.engine.types.WithinSet;
import nz.org.riskscape.engine.types.eqrule.Coercer;
import nz.org.riskscape.engine.types.eqrule.EquivalenceRule;
import nz.org.riskscape.engine.typeset.TypeRules;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.problem.Problem.Severity;
import nz.org.riskscape.rl.TokenTypes;
import nz.org.riskscape.rl.ast.FunctionCall;

@SuppressWarnings("unchecked")
public class DefaultFunctionResolverTest extends ProjectTest {

  List<Type> given = new ArrayList<>();
  FunctionSet functionSet = project.getFunctionSet();
  String functionIdent = "foo";
  Struct inputType = Struct.of();
  JavaFunction matched = JavaFunction.withId("foo");
  FunctionResolver resolver = new DefaultFunctionResolver();

  @Test
  public void returnsEmptyIfNoneFound() throws Exception {
    functionIdent = "bar";
    assertFalse(resolve().isPresent());
  }

  @Test
  public void returnsEmptyIfExtraArgsProvided() throws Exception {
    matched = matched.withArgumentTypes(Arrays.asList(Types.TEXT));
    given = Arrays.asList(Types.TEXT, Types.TEXT);
    assertFalse(resolve().isPresent());
  }

  @Test
  public void returnsAnUnadulteratedFunctionIfThereIsAnExactMatch() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.TEXT));
    given = Arrays.asList(Types.TEXT);
    assertSame(matched, resolve().orElse(null));
  }

  @Test
  public void returnsAnUnadulteratedFunctionIfTheReceiverAndGiverHaveNullableArgs() {
    matched = matched.withArgumentTypes(Arrays.asList(Nullable.of(Types.TEXT)));
    given = Arrays.asList(Nullable.of(Types.TEXT));
    assertSame(matched, resolve().orElse(null));
  }


  @Test
  public void returnsAnUnadulteratedFunctionIfTheReceiverHasNullableArgsAndGiverNot() {
    matched = matched.withArgumentTypes(Arrays.asList(Nullable.of(Types.TEXT)));
    given = Arrays.asList(Types.TEXT);
    assertSame(matched, resolve().orElse(null));
  }

  @Test
  public void returnsANullSafeFunctionIfThereIsMatchExcludingNullables() {
    matched = matched.withArgumentTypes(Types.TEXT);
    given = Arrays.asList(Nullable.TEXT);
    NullSafeFunction function =  resolve(NullSafeFunction.class).get();
    assertSame(function.getTarget(), matched);
  }


  @Test
  public void returnsAnUnadulteratedFunctionIfTheReceiverHasLessSpecificTypes() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.TEXT));
    given = Arrays.asList(new WithinSet(Types.TEXT, "foo", "bar", "baz"));
    assertSame(matched, resolve().orElse(null));
  }

  @Test
  public void returnsEmptyIfTheReceiverHasMoreSpecificTypes() {
    matched = matched.withArgumentTypes(Arrays.asList(new WithinSet(Types.TEXT, "foo", "bar", "baz")));
    given = Arrays.asList(Types.TEXT);
    assertNull(resolve().orElse(null));
  }

  @Test
  public void returnsOriginalFunctionIfThereAreMissingOptionalArgs() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.TEXT, Nullable.INTEGER));
    given = Arrays.asList(Types.TEXT);
    assertSame(matched, resolve().orElse(null));
  }

  @Test
  public void returnsANullSafeFunctionIfThereAreMissingOptionalsAndNonNullArgs() throws Exception {
    matched = matched.withArgumentTypes(Arrays.asList(Types.TEXT, Nullable.INTEGER));
    given = Arrays.asList(Nullable.TEXT);
    assertSame(matched, resolve(NullSafeFunction.class).map(nsf -> nsf.getTarget()).orElse(null));
  }

  @Test
  public void noCoercionIfTheReceiverHasFewerArgs() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.FLOATING));
    given = Arrays.asList(Types.INTEGER, Types.INTEGER);
    assertFalse(resolve().isPresent());
  }

  @Test
  public void wrapsFunctionInNumberAdaptor_ForIntegerToFloating() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.FLOATING));
    given = Arrays.asList(Types.INTEGER);

    CoercingFunctionWrapper numberAdapting = resolve(CoercingFunctionWrapper.class).get();
    assertNotNull(numberAdapting);
    assertEquals(Lists.newArrayList(Types.INTEGER), numberAdapting.getArgumentTypes());
    assertSame(matched, numberAdapting.getWrapped());
  }

  @Test
  public void wrapsFunctionInNumberAdaptor_ForWrappedIntegerToFloating() {
    matched = matched.withArgumentTypes(Arrays.asList(Types.FLOATING));
    Type myInputType = new WithinRange(Types.INTEGER, 0, 10);
    given = Arrays.asList(myInputType);

    CoercingFunctionWrapper numberAdapting = resolve(CoercingFunctionWrapper.class).get();
    assertNotNull(numberAdapting);
    assertEquals(Lists.newArrayList(myInputType), numberAdapting.getArgumentTypes());
    assertSame(matched, numberAdapting.getWrapped());
  }

  @Test
  public void willNotNumberAdaptedToWrappedFloat() {
    matched = matched.withArgumentTypes(Arrays.asList(new WithinRange(Types.FLOATING, 0, 10)));
    given = Arrays.asList(Types.INTEGER);

    assertFalse(resolve(JavaFunction.class).isPresent());
  }

  @Test
  public void delegatesToRealizableFunctionIfPresent() throws Exception {
    functionIdent = "bar";
    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(Types.TEXT), Types.FLOATING);
    Mockito.when(mock.getRealizable()).thenReturn(Optional.of((a, b, c) -> ResultOrProblems.of(matched)));
    matched = matched.withArgumentTypes(Types.TEXT).withReturnType(Types.FLOATING);
    given = Arrays.asList(Types.TEXT);
    functionSet.add(mock);

    assertSame(matched, resolve(RiskscapeFunction.class).get());

  }

  @Test
  public void aRealizedFunctionCanStillBeAdaptedForNullHandling() throws Exception {
    functionIdent = "bar";
    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(Types.TEXT), Types.INTEGER);
    Mockito.when(mock.getRealizable()).thenReturn(Optional.of((a, b, c) -> ResultOrProblems.of(matched)));
    matched = matched.withArgumentTypes(Types.TEXT).withReturnType(Types.FLOATING);
    given = Arrays.asList(Nullable.TEXT);
    functionSet.add(mock);

    assertSame(matched, resolve(NullSafeFunction.class).get().getTarget());
  }

  @Test
  public void ifARealizableFunctionCanNotBeAdaptedAnErrorIsGiven() throws Exception {
    functionIdent = "bar";
    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(Types.ANYTHING), Types.FLOATING);
    Mockito.when(mock.getRealizable()).thenReturn(Optional.of((a, b, c) -> ResultOrProblems.of(matched)));
    matched = matched.withArgumentTypes(Types.TEXT).withReturnType(Types.FLOATING);
    given = Arrays.asList(Types.INTEGER);
    functionSet.add(mock);

    ResultOrProblems<RiskscapeFunction> result = resolve();
    assertThat(
        result,
        Matchers.failedResult(Matchers.isProblem(
            Severity.ERROR, ArgsProblems.class, "realizableDidNotMatch"
        ))
    );

    assertThat(result.getAsSingleProblem(),
        isProblemShown(Severity.ERROR, is("Could not apply 'bar' function to given arguments [Integer]")));
  }

  @Test
  public void resolvesNothingIfRealizableFunctionGivesErrors() throws Exception {
    matched = null;
    functionIdent = "bar";
    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(), Types.ANYTHING);
    Mockito.when(mock.getRealizable()).thenReturn(Optional.of((a, b, c) -> ResultOrProblems.error("sad time")));
    given = Arrays.asList(Types.TEXT);
    functionSet.add(mock);

    assertFalse(resolve().isPresent());
    assertThat(
        resolve(),
        Matchers.failedResult(
            allOf(
                Matchers.isProblemAffecting(Severity.ERROR, RiskscapeFunction.class),
                Matchers.hasProblems(
                    Matchers.isProblem(Severity.ERROR, containsString("sad time"))
                )
            )
        )
    );
  }

  @Test
  public void willAttemptToMatchOverloadedAlternativesIfPresent() throws Exception {
    matched = null;
    functionIdent = "bar";
    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(Types.TEXT), Types.ANYTHING);
    functionSet.add(mock);

    RiskscapeFunction option1 = mockFunction("option 1", Arrays.asList(Types.FLOATING), Types.FLOATING);
    RiskscapeFunction option2 = mockFunction("option 2", Arrays.asList(Types.GEOMETRY), Types.GEOMETRY);

    OverloadedFunction mockOverloaded = Mockito.mock(OverloadedFunction.class);
    Mockito.when(mockOverloaded.getAlternatives()).thenReturn(Arrays.asList(option1, option2));
    Mockito.when(mock.getOverloaded()).thenReturn(Optional.of(mockOverloaded));

    given = Arrays.asList(Types.FLOATING);
    assertSame(option1, resolve().get());

    given = Arrays.asList(Types.GEOMETRY);
    assertSame(option2, resolve().get());

    given = Arrays.asList(Types.TEXT);
    assertSame(mock, resolve().get());

    // wrapping still takes effect
    given = Arrays.asList(Nullable.FLOATING);
    assertSame(option1, resolve(NullSafeFunction.class).get().getTarget());
  }


  @Test
  public void anOverloadedFunctionIsIgnoredIfIgnoreThisIsTrue() throws Exception {
    AtomicBoolean bool = new AtomicBoolean(true);
    functionIdent = "bar";
    matched = null;

    IdentifiedFunction mock = mockIdentified(functionIdent, Arrays.asList(Types.TEXT), Types.TEXT);
    functionSet.add(mock);

    RiskscapeFunction option1 = mockFunction("option 1", Arrays.asList(Types.FLOATING), Types.FLOATING);

    OverloadedFunction mockOverloaded = Mockito.mock(OverloadedFunction.class);
    Mockito.when(mockOverloaded.getAlternatives()).thenReturn(Arrays.asList(option1));
    Mockito.when(mock.getOverloaded()).thenReturn(Optional.of(mockOverloaded));
    Mockito.when(mockOverloaded.ignoreThis()).thenAnswer(i -> bool.get());

    given = Arrays.asList(Types.FLOATING);
    assertSame(option1, resolve().get());

    given = Arrays.asList(Types.TEXT);
    bool.set(true);
    assertFalse(resolve().isPresent());
    bool.set(false);
    assertSame(mock, resolve().get());
  }

  @Test
  public void declaredLinkedTypesOnAFunctionAreStrippedOffBeforeArgumentAdapting() throws Exception {
    Type positiveInteger = new WithinRange(Types.INTEGER, 0, Long.MAX_VALUE);
    matched = matched.withArgumentTypes(project.getTypeSet().add("footype", positiveInteger));

    given = Arrays.asList(positiveInteger);
    assertSame(matched, resolve(RiskscapeFunction.class).orElse(null));
  }

  @Test
  public void aSingleValueStructCanBeTreatedLikeTheSingleValue() throws Exception {
    Type struct = Struct.of("foo", Types.INTEGER);
    Type integer = Types.INTEGER;

    matched = matched.withArgumentTypes(struct);
    given = Arrays.asList(integer);

    CoercingFunctionWrapper wrapper = resolve(CoercingFunctionWrapper.class).orElse(null);
    assertSame(matched, wrapper.getWrapped());
  }

  @Test
  public void aSimpleValueCanBeTreatedLikeASingleValueStruct() throws Exception {
    Type struct = Struct.of("foo", Types.INTEGER);
    Type integer = Types.INTEGER;

    matched = matched.withArgumentTypes(integer);
    given = Arrays.asList(struct);

    CoercingFunctionWrapper wrapper = resolve(CoercingFunctionWrapper.class).orElse(null);
    assertSame(matched, wrapper.getWrapped());
  }

  @Test
  public void aNullableSimpleValueCanBeTreatedLikeANullableSingleValueStruct() throws Exception {
    Type struct = Nullable.of(Struct.of("foo", Types.INTEGER));
    Type integer = Nullable.INTEGER;

    matched = matched.withArgumentTypes(integer);
    given = Arrays.asList(struct);

    CoercingFunctionWrapper function = resolve(CoercingFunctionWrapper.class).get();
    assertSame(matched, function.getWrapped());
  }

  @Test
  public void aMisbehavingCoercionRuleDoesNotCauseAStackOverflow() throws Exception {
    EquivalenceRule naughtyRule = new EquivalenceRule() {

      @Override
      public Optional<Coercer> getCoercer(TypeRules typeSet, Type sourceType, Type targetType) {
        if (sourceType == targetType) {
          return Optional.of(Coercer.build(sourceType, targetType, Function.identity()));
        } else {
          return Optional.empty();
        }
      }

    };
    project.getTypeSet().getTypeRegistry().addEquivalenceRule(naughtyRule);
    Type struct = Struct.of("foo", Types.TEXT);
    Type linkedType = project.getTypeSet().add("linked", struct);

    matched = matched
        .withArgumentTypes(linkedType, Types.TEXT);

    // there used to be a bug where the partialstructrule was erroneously giving a coercer for the same struct,
    // which meant the function resolving routine was wrapping over and over again, giving a stack overflow - this
    // was happening in the real world because of a linked type yielding the exact same struct, which the second arg
    // required a null safe function
    given = Arrays.asList(linkedType, Nullable.TEXT);

    assertNotNull(resolve(NullSafeFunction.class).orElse(null));
  }

  @Test
  public void checkRealizationOccursAgainstAdaptedAndOriginalTypes() throws Exception {

    // This tests the case where realization was adapting a realized function's advertised args before realizing and
    // making sure that RealizableFunction#realize gets given the adapted types, not the given types
    matched = null;
    functionIdent = "sample_centroid";
    // the given args in our example contains one coerceable type (the struct with the geometry) and one covariant type
    // ( the coverage)
    given = Arrays.asList(
        Struct.of("geom", Types.GEOMETRY, "name", Types.TEXT),
        new CoverageType(Nullable.of(Types.FLOATING))
    );

    // ... and as the realizable function declares geometry as the first arg, the adapting routine finds the coercer to
    // go from struct -> geom, but it must tell the realization routine that this is happening by passing the target
    // type of the coercer, not the declared type or the given type
    List<Type> declaredArgumentTypes = Arrays.asList(
        Types.GEOMETRY,
        new CoverageType(Nullable.ANYTHING)
    );

    // we want to see the adapted geometry - target of the coercion
    // but the given covariant coverage type - from given
    List<Type> expectedTypesForRealization = Arrays.asList(
        Types.GEOMETRY,
        new CoverageType(Nullable.of(Types.FLOATING))
    );
    RiskscapeFunction realizationResult =
        mockFunction("realization result", expectedTypesForRealization, Types.ANYTHING);

    IdentifiedFunction function = mockIdentified(functionIdent, declaredArgumentTypes, Types.ANYTHING);
    RealizableFunction realizable = Mockito.mock(RealizableFunction.class, "realizable function");
    Mockito.when(realizable.isDoTypeAdaptation()).thenReturn(true);
    Mockito.when(function.getRealizable()).thenReturn(Optional.of(realizable));


    // this is happening out of order, but shouldn't be an issue for the test.  In "real life", the sampling function
    // returns a function that has argument types matching those given at realization
    // NB maybe we should be warning when realize doesn't return a function that matches this input types?
    Mockito.when(realizable.realize(
        Mockito.same(this.realizationContext),
        Mockito.any(),
        Mockito.eq(expectedTypesForRealization)
    ))
      .thenReturn(ResultOrProblems.of(realizationResult));

    functionSet.add(function);
    ResultOrProblems<RiskscapeFunction> resolvedOr = resolve();

    assertTrue(resolvedOr.isPresent());
    RiskscapeFunction resolved = resolvedOr.get();
    CoercingFunctionWrapper wrapper = (CoercingFunctionWrapper) resolved;
    assertSame(wrapper.getWrapped(), realizationResult);
  }

  @Test
  public void checkRealizationOccursAgainstAdaptedAndOriginalTypes1() throws Exception {

    // This tests the case where realization was adapting a realized function's advertised args before realizing and
    // making sure that RealizableFunction#realize gets given the adapted types, not the given types
    matched = null;
    functionIdent = "test_function";
    // the given args in our example contains structs with coerceable types
    Struct arg1Type = Struct.of("x", Types.INTEGER, "y", Types.FLOATING);
    Struct arg2Type = Struct.of("a", Types.TEXT, "b", Types.INTEGER);
    given = Arrays.asList(arg1Type, arg2Type);

    // declared are sub-set of those given, and x needs int -> float coercion
    List<Type> declaredArgumentTypes = Arrays.asList(
        Struct.of("x", Types.FLOATING),
        Struct.of("a", Types.TEXT)
    );

    // x has int -> float coercion
    List<Type> expectedTypesForRealization = Arrays.asList(
        Struct.of("x", Types.FLOATING),
        Struct.of("a", Types.TEXT)
    );

    RiskscapeFunction realizationResult =
        mockFunction("realization result", expectedTypesForRealization, Types.ANYTHING);

    IdentifiedFunction function = mockIdentified(functionIdent, declaredArgumentTypes, Types.ANYTHING);
    RealizableFunction realizable = Mockito.mock(RealizableFunction.class, "realizable function");
    Mockito.when(realizable.isDoTypeAdaptation()).thenReturn(true);
    Mockito.when(function.getRealizable()).thenReturn(Optional.of(realizable));

    List<Types> realizedTypes = new ArrayList<>();
    Mockito.when(realizationResult.call(Mockito.any())).thenAnswer(invocation -> {
      List args = invocation.getArgument(0, List.class);
      // Now check that the function is called with the same types as was realized
      for (int i = 0; i < realizedTypes.size(); i++) {
        Tuple arg = (Tuple)args.get(i);
        assertSame(arg.getStruct(), realizedTypes.get(i));
      }
      return "good";
    });

    Mockito.when(realizable.realize(
        Mockito.same(this.realizationContext),
        Mockito.any(),
        Mockito.eq(expectedTypesForRealization)
    ))
      .thenAnswer(invocation -> {
        // Save the actuall types that were realized
        realizedTypes.addAll(invocation.getArgument(2));
        return ResultOrProblems.of(realizationResult);
      });

    functionSet.add(function);
    ResultOrProblems<RiskscapeFunction> resolvedOr = resolve();

    assertTrue(resolvedOr.isPresent());
    RiskscapeFunction resolved = resolvedOr.get();
    CoercingFunctionWrapper wrapper = (CoercingFunctionWrapper) resolved;
    assertSame(wrapper.getWrapped(), realizationResult);

    resolved.call(Arrays.asList(
        Tuple.ofValues(arg1Type, 10L, 3.0D),
        Tuple.ofValues(arg2Type, "bar", 4L)
    ));
  }

  @Test
  public void aReturnedStructWillGetNormalizedForYou() throws Exception {
    Struct toBeNormalized = Struct.of("foo", Types.TEXT);
    Struct notNormalized = Struct.of("foo", Types.TEXT);

    functionIdent = "bar";
    IdentifiedFunction function = mockIdentified(functionIdent, Arrays.asList(), toBeNormalized);
    functionSet.add(function);

    assertSame(function, resolve().get());

    // call it with the struct we didn't expect to become the canonical one - it should return the one that the resolver
    // was supposed to normalize
    assertSame(toBeNormalized, realizationContext.normalizeStruct(notNormalized));
  }

  private ResultOrProblems<RiskscapeFunction> resolve() {
    return resolve(RiskscapeFunction.class);
  }

  private <T extends RiskscapeFunction> ResultOrProblems<T> resolve(Class<T> expected) {
    if (matched != null) {
      functionSet.add(matched);
    }
    FunctionCall fc = new FunctionCall(Token.token(TokenTypes.IDENTIFIER, functionIdent), Collections.emptyList());
    return functionSet.resolve(this.realizationContext, fc, inputType, given, resolver).map(found -> {
      if (!expected.isInstance(found)) {
        fail(found + " not of expected type " + expected);
      }
      return expected.cast(found);
    });
  }

  private IdentifiedFunction mockIdentified(String id, List<Type> argumentTypes, Type returnType) {
    IdentifiedFunction mock = mockFunction(IdentifiedFunction.class, id, argumentTypes, returnType);

    Mockito.when(mock.getId()).thenReturn(id);
    Mockito.when(mock.getSourceURI()).thenReturn(Resource.UNKNOWN_URI);

    return mock;
  }

  private RiskscapeFunction mockFunction(String desc, List<Type> argumentTypes, Type returnType) {
    return mockFunction(RiskscapeFunction.class, desc, argumentTypes, returnType);
  }

  private <T extends RiskscapeFunction> T mockFunction(Class<T> clazz, String desc, List<Type> args, Type returnType) {
    T mock = Mockito.mock(clazz, desc);

    Mockito.when(mock.getArgumentTypes()).thenReturn(args);
    Mockito.when(mock.getReturnType()).thenReturn(returnType);

    return mock;
  }

}
