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

import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.FunctionArgument;
import nz.org.riskscape.engine.function.FunctionMetadata;
import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.function.IdentifiedFunction.Category;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.function.UntypedFunction;
import nz.org.riskscape.engine.types.Bool;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.ScopedLambdaType;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.engine.types.eqrule.Coercer;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.ExpressionProblems;
import nz.org.riskscape.rl.ast.FunctionCall;
import nz.org.riskscape.rl.ast.FunctionCall.Argument;
import nz.org.riskscape.rl.ast.Lambda;
import nz.org.riskscape.rl.ast.MinimalVisitor;

/**
 * Implementation of https://docs.geoserver.org/latest/en/user/filter/function_reference.html#control-functions
  */
public class IfThenElse implements RealizableFunction {

  public static final ArgumentList ARGUMENTS = ArgumentList.create(
      "condition", Nullable.BOOLEAN,
      "then", Nullable.ANYTHING,
      "else", Nullable.ANYTHING
  );

  public static final IdentifiedFunction IF = asFunction("if", "");
  public static final IdentifiedFunction IF_THEN_ELSE = asFunction("if_then_else", "");

  private static IdentifiedFunction asFunction(String id, String description) {
    FunctionMetadata metadata = new FunctionMetadata(id, ARGUMENTS, Types.ANYTHING, description,
        Category.LOGICAL, RiskscapeFunction.BUILT_IN);
    return RealizableFunction.identified(new IfThenElse(), metadata);
  }

  @Override
  public ResultOrProblems<RiskscapeFunction> realize(
      RealizationContext context,
      FunctionCall functionCall,
      List<Type> givenTypes
  ) {
    if (givenTypes.size() < 2 || givenTypes.size() > ARGUMENTS.size()) {
      return ResultOrProblems.failed(
          Problems.get(ArgsProblems.class).wrongNumberRange(2, ARGUMENTS.size(), givenTypes.size()));
    }

    return ProblemException.catching(() -> {
      ARGUMENTS.getRequiredArgument(functionCall, "condition").getOrThrow();
      ARGUMENTS.getRequiredArgument(functionCall, "then").getOrThrow();


      // ensure the condition is a boolean.
      // the function realizer will apply coercers if possible before falling back to whatever was provided.
      // so we only need to make sure we get a boolean here (since the other args are anything)
      givenTypes.get(0).findAllowNull(Bool.class)
          .orElseThrow(() -> new ProblemException(ArgsProblems.mismatch(ARGUMENTS.get(0), givenTypes.get(0))));

      Boolean constantConditionArg = checkConstantConditionArg(context, functionCall);

      if (constantConditionArg != null) {
        return realizeConstant(context, functionCall, givenTypes, constantConditionArg);
      } else {
        return realizeNotConstant(context, functionCall, givenTypes);
      }
    });
  }

  private RiskscapeFunction realizeConstant(RealizationContext context, FunctionCall functionCall,
      List<Type> givenTypes, Boolean constantConditionArg) throws ProblemException {

    // return cut-down functions that don't bother testing the condition
    if (constantConditionArg) {
      RealizedExpression thenLambda = realizeLambdaExpressions(givenTypes.get(1),
          functionCall.getArguments().get(1), context);
      Type thenType = thenLambda != null ? thenLambda.getResultType() : givenTypes.get(1);

      return RiskscapeFunction.create(this, givenTypes, thenType, args ->
        evaluateThen(args, thenLambda)
      );
    } else {

      if (givenTypes.size() == 2) {
        // no else condition, test failed, so return null/nothing

        // try and infer the type for the then case, but allow it to fail (and return nothing) - this makes it easier
        // to use if in constant conditions where most of the time you'll want to leave behind some value in the
        // pipeline that'll always be null without having to have acrobatics to deal with a nothing type instead of a
        // nullable(thing) where the if case would have succeeded.  I think the nothing return type is the 20% case and
        // can be still supported by returning an explicit null from else.
        Type thenType = Types.NOTHING;
        try {
          RealizedExpression thenLambda = realizeLambdaExpressions(givenTypes.get(1),
            functionCall.getArguments().get(1), context);
          thenType = Nullable.of(thenLambda != null ? thenLambda.getResultType() : givenTypes.get(1));
        } catch (ProblemException e) {}

        return RiskscapeFunction.create(this, givenTypes, thenType, args -> null);
      } else {
        RealizedExpression elseLambda = realizeLambdaExpressions(givenTypes.get(2),
            functionCall.getArguments().get(2), context);
        Type elseType = elseLambda != null ? elseLambda.getResultType() : givenTypes.get(2);

        return RiskscapeFunction.create(this, givenTypes, elseType, args ->
          evaluateElse(args, elseLambda)
        );
      }
    }
  }

  private RiskscapeFunction realizeNotConstant(
      RealizationContext context,
      FunctionCall functionCall,
      List<Type> givenTypes
  ) throws ProblemException {

    RealizedExpression thenLambda = realizeLambdaExpressions(givenTypes.get(1),
        functionCall.getArguments().get(1), context);
    Type thenType = thenLambda != null ? thenLambda.getResultType() : givenTypes.get(1);

    RealizedExpression elseLambda;
    Type returnType;
    Type elseType = null;
    if (givenTypes.size() == 2) {
      elseLambda = null;
      // no else condition given. we'll be returning then or null
      returnType = Nullable.of(thenType);
    } else {
      elseLambda = realizeLambdaExpressions(givenTypes.get(2),
          functionCall.getArguments().get(2), context);
      elseType = elseLambda != null ? elseLambda.getResultType() : givenTypes.get(2);

      returnType = context.getTypeSet().computeAncestorNoConversion(thenType, elseType);
    }

   return RiskscapeFunction.create(
      this, givenTypes, returnType,
      buildFunction(thenLambda, elseLambda),
      // make sure these are cleaned up if realized
      thenLambda, elseLambda
    );
  }

  private Boolean checkConstantConditionArg(RealizationContext context, FunctionCall functionCall)
      throws ProblemException{
    FunctionArgument conditionArgument = ARGUMENTS.get("condition");

    // first up we need to test whether the argument contains any function calls - if it does, there's the possibility
    // that it contains a random function, which while constant in the sense that it doesn't depend on the input scope,
    // is not really suitable for constantization for the purposes of the if method.
    // See https://gitlab.catalyst.net.nz/riskscape/riskscape/-/issues/1568 for more info and how we might fix this.
    Expression expression = functionCall.getArguments().get(conditionArgument.getIndex()).getExpression();

    AtomicBoolean conditionContainsFunctionCall = new AtomicBoolean(false);
    expression.accept(new MinimalVisitor<Void>() {
      @Override
      public Void visit(FunctionCall expression, Void data) {
        conditionContainsFunctionCall.set(true);
        return null;
      }
    }, null);


    // nope, can't risk it until we have better support for annotating methods as constant or not
    if (conditionContainsFunctionCall.get()) {
      return null;
    }

    // We roll our own expression here because using FunctionArgument#evaluateConstant is not able
    // to return a null value (which we want to use with the the elseType path.)
    RealizedExpression constantConditionExpr = conditionArgument
        .getFunctionCallArgument(functionCall)
        .map(expr -> context.getExpressionRealizer().realizeConstant(expr.getExpression()).orElse(null))
        .orElse(null);

    // not constant
    if (constantConditionExpr == null) {
      return null;
    }

    Object uncoercedConditionArg = constantConditionExpr.evaluateConstant();
    // even through we've checked that the given condition arg is a bool we may need to use a coercer
    // because the function realizer may have some arg mapping in play, which doesn't apply to the
    // expression we've realized.
    // this makes allows a const condition like `{test: true}` to work in the same way that would
    // in a non-const expression.
    Coercer constConditionCoercer = context.getTypeSet()
        .findEquivalenceCoercer(Nullable.strip(constantConditionExpr.getResultType()), Types.BOOLEAN)
        // given that the function realizer must have a coercer for the cond expr to boolean we'd
        // never expect to throw this problem, be we keep it for completeness.
        .orElseThrow(() -> new ProblemException(
            ArgsProblems.mismatch(ARGUMENTS.get(0), constantConditionExpr.getResultType())));

    return (Boolean) constConditionCoercer.apply(uncoercedConditionArg);
  }

  /**
   * Realizes the lambda expression if one was provided.
   *
   * The then/else arguments do not have to be lambda expressions. But if they are they will be
   * validated (don't accept any args) and realized for later use.
   *
   * @return realized lambda expression or null if provided arg is not a lambda
   */
  private RealizedExpression realizeLambdaExpressions(Type givenType, Argument arg, RealizationContext context)
      throws ProblemException {
    ScopedLambdaType lambdaType = givenType.find(ScopedLambdaType.class).orElse(null);
    if (lambdaType == null) {
      // the argument is not a lambda expression so bail from here (lambdas are optional)
      return null;
    }
    if (lambdaType.getArity() != 0) {
      throw new ProblemException(ExpressionProblems.get().lambdaArityError(
          arg.getExpression(),
          lambdaType.getArity(),
          0
      ));
    }

    Lambda lambda = (Lambda) arg.getExpression();
    return context.getExpressionRealizer()
        .realize(lambdaType.getScopeType(), lambda.getExpression())
        .getOrThrow();
  }

  /**
   * Return an UntypedFunction for the if else, with optional then and else expressions
   */
  private UntypedFunction buildFunction(
      RealizedExpression thenLambda,
      RealizedExpression elseLambda
  ) {
    return (args) -> {
      Object condition = args.get(0);

      if (condition != null && condition.equals(Boolean.TRUE)) {
        return evaluateThen(args, thenLambda);
      } else {
        if (args.size() == 2) {
          // there is no else argument so we just return null
          return null;
        }
        return evaluateElse(args, elseLambda);
      }
    };
  }

  private Object evaluateThen(List<Object> args, RealizedExpression thenLambda) {
    if (thenLambda == null) {
      return args.get(1);
    }
    ScopedLambdaExpression givenLambda = (ScopedLambdaExpression) args.get(1);
    return thenLambda.evaluate(givenLambda.getScope());
  }

  private Object evaluateElse(List<Object> args, RealizedExpression elseLambda) {
    if (elseLambda == null) {
      return args.get(2);
    }
    ScopedLambdaExpression givenLambda = (ScopedLambdaExpression) args.get(2);
    return elseLambda.evaluate(givenLambda.getScope());
  }
}
