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

import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.FailedObjectException;
import nz.org.riskscape.engine.NoSuchObjectException;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.ParameterSet;
import nz.org.riskscape.engine.data.Bookmark;
import nz.org.riskscape.engine.data.BookmarkFactory;
import nz.org.riskscape.engine.data.ResolvedBookmark;
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.problem.ProblemFactory;
import nz.org.riskscape.engine.problem.SeverityLevel;
import nz.org.riskscape.engine.resource.UriHelper;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Struct.StructMember;
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.engine.types.WithMetadata;
import nz.org.riskscape.engine.typeset.TypeSet;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problem.Severity;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.FunctionCall;
import nz.org.riskscape.rl.ast.FunctionCall.Argument;

/**
 * Inject a source of data in to a pipeline, either as a constant (during realization) or dynamically (during execution)
 */
public class LookupBookmark implements RealizableFunction {

  public interface LocalProblems extends ProblemFactory {
    Problem optionValueMustBeTextNumericOrBoolean(String optionKey, Type found);

    @SeverityLevel(Severity.INFO)
    Problem typeRequiredHint();

    @SeverityLevel(Severity.INFO)
    Problem badPlaceholderHint();
  }

  static final LocalProblems PROBLEMS = Problems.get(LocalProblems.class);

  @Getter
  private final ArgumentList arguments = ArgumentList.create(
      "id", Types.TEXT,
      "options", Types.ANYTHING,  // options is anything, rather than an emtpy struct as all sorts of
                                  // type equivilance rules (partial struct, single value struct) may coerce to {}
                                  // and the default function resolver does that before realizing with given args
      "type", Types.TEXT
  );

  /**
   * Realization helper method - we allocate one of these for each call to realize to close over all our realization
   * state
   */
  @RequiredArgsConstructor
  class Instance {
    final RealizationContext context;
    final FunctionCall functionCall;
    final List<Type> givenTypes;

    Bookmark constantBookmarkFromIdArg;
    Tuple constantOptions;
    Type constantType;

    // record these to make it easier to construct problems
    FunctionCall.Argument idArg;
    FunctionCall.Argument optionsArg;
    FunctionCall.Argument typeArg;

    // problems collected during realization/validation
    List<Problem> problems = new ArrayList<>();

    /**
     * Do basic checks that the given args are valid, e.g. arity, types are correct, plus memoize constant values.
     */
    boolean validateArgs() {
      // arity check
      int givenArgs = functionCall.getArguments().size();
      if (givenArgs > arguments.size() || givenArgs < 1) {
        problems.add(ArgsProblems.get().wrongNumberRange(
            1, arguments.size(), givenArgs));
        return false;
      }

      // build a bookmark from our idArg
      idArg = arguments.getArgument(functionCall, "id").get();
      idArg
        .evaluateConstant(context, String.class, Types.TEXT)
        .ifElse(
            id -> {
              // if it's a constant string, we lookup a bookmark (or build one from a URI)
              lookupBookmark(id).ifElse(
                  bookmark -> this.constantBookmarkFromIdArg = bookmark,
                  probs -> problems.addAll(probs));
            },
            probs -> {
              // it might be a nested bookmark function call, let's see
              Bookmark nestedBookmark = WithMetadata.find(givenTypes.get(0), Bookmark.class).orElse(null);

              if (nestedBookmark != null) {
                this.constantBookmarkFromIdArg = nestedBookmark;
              } else {
                idArg.evaluateConstant(context, Object.class, Types.ANYTHING).ifPresent(constantValue -> {
                  // we have a constant, but it's not a string and it's not a bookmark expression
                  problems.add(TypeProblems.get().mismatch(idArg, Types.TEXT, givenTypes.get(0)));
                });
              }
            }
        );

      // build options tuple (if constant)
      arguments
          .getArgument(functionCall, "options")
          .ifPresent(arg -> {
            optionsArg = arg;
            if (validateOptions()) {
              arg.evaluateConstant(context, Tuple.class, Struct.EMPTY_STRUCT).ifPresent(optionsTuple -> {
                constantOptions = optionsTuple;
              });
            }
          });


      // parse third type argument (if given)
      arguments.getArgument(functionCall, "type").ifPresent(arg -> {
        typeArg = arg;

        arg.evaluateConstant(context, String.class, Types.TEXT).ifElse(
            typeExp -> {
              try {
                constantType = context.getProject().getTypeBuilder().build(typeExp);
              } catch (RiskscapeException e) {
                problems.add(Problems.caught(e));
              }
            },
            probs -> problems.addAll(probs));
      });

      // validate the type arg (or type inference) for a non-constant bookmark
      if (!isConstant() && !hasErrors()) {
        if (constantType == null) {
          // there's no constant type, let's see if we have a constant bookmark we can use to infer the type
          if (constantBookmarkFromIdArg != null) {
            Bookmark templateBookmark = this.constantBookmarkFromIdArg;

            resolveBookmark(templateBookmark).ifElse(
                resolved -> this.constantType = resolved.getScalarType(),
                probs -> {
                  problems.addAll(probs);
                  // the constant bookmark ID is not a valid bookmark. The simplest thing the user can do here
                  // is use a valid placeholder bookmark
                  problems.add(PROBLEMS.badPlaceholderHint());
                }
            );
          }
        }

        // check again and see if we were able to infer a type
        if (constantType == null && !hasErrors()) {
          // this should add the 'type arg missing' error
          // NB: the bookmark ID must be non-constant to hit this case
          boolean added = problems.add(
              Problems
                .toSingleProblem(arguments.getRequiredArgument(functionCall, "type").getProblems())
                .withChildren(PROBLEMS.typeRequiredHint())
          );
          assert added; // if the list didn't change, that's a bug
        }
      }

      return !hasErrors();
    }

    /**
     * @return a bookmark built from constant args (with options merged)
     */
    ResultOrProblems<Bookmark> getConstantBookmark() {
      assert constantBookmarkFromIdArg != null;

      if (constantOptions != null) {
        return mergeBookmarkWithCustomOptions(constantBookmarkFromIdArg, constantOptions);
      } else {
        return ResultOrProblems.of(constantBookmarkFromIdArg);
      }
    }

    /**
     * @return true if any error problems have been collected during validation
     */
    boolean hasErrors() {
      return Problem.hasErrors(problems);
    }

    /**
     * Construct a bookmark using the `id` argument, first by seeing if it's a bookmark, then fallback to assuming it's
     * a URI
     */
    private ResultOrProblems<Bookmark> lookupBookmark(String idOrUri) {
      Project project = context.getProject();

      try {
        return ResultOrProblems.of(project.getBookmarks().get(idOrUri, context.getProblemSink()));
      } catch (NoSuchObjectException | FailedObjectException e) {
        // not a bookmark. Check if it's a file (the resolver will validate it below)
        URI tryThis = BookmarkFactory.uriFromLocation(idOrUri, context.getProject().getRelativeTo()).orElse(null);

        if (tryThis == null) {
          return ResultOrProblems.failed(Problems.caught(e));
        } else {
          return ResultOrProblems.of(Bookmark.fromURI(tryThis, idOrUri));
        }
      }
    }

    /**
     * Construct a bookmark by merging a {@link Bookmark} with any `options` provided in the 2nd argument.
     */
    private ResultOrProblems<Bookmark> mergeBookmarkWithCustomOptions(
        Bookmark originalBookmark,
        Tuple optionsFromArguments
    ) {

      Bookmark bookmark = originalBookmark.clone();

      // convert tuple to a map
      Map<String, Object> optionsMap = optionsFromArguments.toMap();

      /*
       * Format and location need to be treated specially - they are present for all bookmarks and help determine
       * which resolver (similar to a framework with models and functions) will try to convert the parameters in to a
       * a piece of data we will inject in to the platform
       */

      Object format = optionsMap.remove("format");
      if (format != null) {
        bookmark = bookmark.withFormat(format.toString());
      }

      Object location = optionsMap.remove("location");
      if (location != null) {
        ResultOrProblems<URI> uri =
            UriHelper.uriFromLocation(location.toString(), context.getProject().getRelativeTo());

        if (uri.hasErrors()) {
          return ResultOrProblems.failed(uri.getProblems());
        }

        bookmark = bookmark.withLocation(uri.get());
      }

      return ResultOrProblems.of(bookmark.addUnparsed(ParameterSet.normaliseParameterMap(optionsMap)));
    };

    /**
     * Validates that the option argument's type is a struct that only contains binding-friendly values
     */
    boolean validateOptions() {
      FunctionArgument declared = arguments.get("options");
      Argument ast = arguments.getArgument(functionCall, "options").orElse(null);

      // nothing given
      if (ast == null) {
        return true;
      }

      // make sure it's a struct
      Type givenType = givenTypes.get(declared.getIndex());
      Struct optionsType = givenType.find(Struct.class).orElse(null);
      if (optionsType == null) {
        problems.add(TypeProblems.get().mismatch(ast, Struct.class, givenType.getClass()));
        return false;
      }

      // make sure members are suitable for setting as parameter binding inputs
      TypeSet typeSet = context.getProject().getTypeSet();
      boolean valid = true; // check them all, don't stop on first failure
      for (StructMember member: optionsType.getMembers()) {
        Type memberType = member.getType()
            .findAllowNull(RSList.class)
            .map(list -> list.getMemberType())
            .orElse(member.getType());

        if (! (typeSet.isAssignable(memberType, Types.TEXT)
            || typeSet.isAssignable(memberType, Types.BOOLEAN)
            || typeSet.isAssignable(memberType, Types.INTEGER)
            || typeSet.isAssignable(memberType, Types.FLOATING))) {
          problems.add(PROBLEMS.optionValueMustBeTextNumericOrBoolean(member.getKey(), member.getType()));
          valid = false;
        }
      }

      return valid;
    }

    /**
     * @return true if this bookmark function returns a constant value, i.e. does not rely on scope.  Constant
     * bookmark function calls always return the same resolved bookmark value (for each call), whereas dynamic ones will
     * resolve the bookmark anew for each call of the function.
     */
    public boolean isConstant() {
      return constantBookmarkFromIdArg != null && (optionsArg == null || constantOptions != null);
    }

    /**
     * Syntactic sugar/helper for resolving a bookmark
     */
    ResultOrProblems<ResolvedBookmark> resolveBookmark(Bookmark bookmark) {
      return context.getProject().getEngine().getBookmarkResolvers().resolveAndValidate(
          bookmark,
          context.getProject().newExecutionContext(),
          Object.class
      );
    }

    /**
     * Validate that a resolved bookmark's actual return type matches the type argument's type (if given)
     */
    public Problem validateBookmarkType(ResolvedBookmark resolved) {
      if (constantType == null) {
        return null;
      }

      if (! context.getProject().getTypeSet().isAssignable(resolved.getScalarType(), constantType)) {
        // note we use the bookmark toString() here rather than the bookmark itself because
        // the object render (used by the problem renderer) would otherwise just output the
        // bookmark ID. But if the pipeline has switched out the location that isn't likely
        // to be much help to the user.
        String errorContext = resolved.getBookmark().toString();
        return TypeProblems.get().mismatch(errorContext, constantType, resolved.getScalarType());
      }

      return null;
    }
  }

  @Override
  public ResultOrProblems<RiskscapeFunction> realize(RealizationContext context, FunctionCall functionCall,
      List<Type> givenTypes
  ) {

    Instance instance = new Instance(context, functionCall, givenTypes);

    if (!instance.validateArgs()) {
      return ResultOrProblems.failed(instance.problems);
    }

    if (instance.isConstant()) {
      return realizeConstant(instance);
    } else {
      return realizeDynamic(instance);
    }
  }

  private ResultOrProblems<RiskscapeFunction> realizeConstant(Instance instance) {
    return instance.getConstantBookmark()
      .flatMap(bookmark -> instance.resolveBookmark(bookmark))
      .flatMap(resolved -> {
        // for consistency, if a type arg is given for a constant bookmark expression, we check the returned type is
        // assignable to the declared type.  It would make more sense to check that the type is equivalent and support
        // coercion, but we can come back to this if required. (or just ignore it?)
        Problem returnTypeProblem = instance.validateBookmarkType(resolved);
        if (returnTypeProblem != null) {
          return ResultOrProblems.failed(returnTypeProblem);
        }

        // wrap return type with the expression that created it - this allows a nested bookmark expression to extract
        // the original bookmark that was set e.g. (bookmark(bookmark(...)) - this is a thing because the platform will
        // generate expressions that are a `bookmark()` function call as a result of customizing a model input.  The
        // pipeline it gets injected in to might be a parameter that's part of a further bookmark function call, in
        // which case this function has some smarts to deal with this nesting
        Type returnType = WithMetadata.wrap(resolved.getScalarType(), resolved.getBookmark());

        return ResultOrProblems.of(RiskscapeFunction.create(this, instance.givenTypes, returnType,
            (args) -> resolved.getData(Object.class).get()));
      }
    );
  }

  private ResultOrProblems<RiskscapeFunction> realizeDynamic(Instance instance) {
    return ProblemException.catching(() -> {
      Type expectedReturnType = instance.constantType;

      // if the first bookmark arg is constant, we can look this up once and reuse it
      Bookmark constantBookmark;
      if (instance.constantBookmarkFromIdArg != null) {
        constantBookmark = instance.getConstantBookmark().getOrThrow();
      } else {
        constantBookmark = null;
      }

      return RiskscapeFunction.create(this, instance.givenTypes, expectedReturnType, (args) -> {

        // Create a bookmark from just the first argument (or use the memoized one)
        Bookmark originalBookmark;
        if (constantBookmark != null) {
          originalBookmark = constantBookmark;
        } else {
          String id = (String) args.get(0);
          originalBookmark = instance.lookupBookmark(id).orElseThrow(problems -> {
            throw new RiskscapeException(Problems.foundWith(instance.idArg, problems));
          });
        }

        // merge in extra bookmark parameters if given
        Tuple options = (Tuple) args.get(1);
        Bookmark toResolve;
        if (options != null) {
          toResolve = instance.mergeBookmarkWithCustomOptions(originalBookmark, options).orElseThrow(problems -> {
            throw new RiskscapeException(Problems.foundWith(instance.optionsArg, problems));
          });
        } else {
          toResolve = originalBookmark;
        }

        // NB We cache the resolving in the case that the same bookmark ends up being built multiple times - it's
        // assumed that resolving a bookmark can be fairly expensive (and wasteful) so we want to be sure not to do
        // it excessively
        return instance.context.getOrComputeFromCache(LookupBookmark.class, toResolve, Object.class, (key) -> {
          ResolvedBookmark resolved = instance.resolveBookmark(toResolve).orElseThrow(problems -> {
            throw new RiskscapeException(Problems.foundWith(toResolve.toString(), problems));
          });

          // make sure that what was looked up matches declared type
          Problem typeProblem = instance.validateBookmarkType(resolved);
          if (typeProblem != null) {
            throw new RiskscapeException(typeProblem);
          }

          return resolved
              // Hm, do we need to apply the casting version of getData here?  We've already done a type check above
              .getData(expectedReturnType.internalType())
              .orElseThrow(problems -> new RiskscapeException(Problems.foundWith(toResolve.toString())));
        });
      });
    });
  }
}

