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

import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeIOException;
import nz.org.riskscape.engine.bind.BindingContext;
import nz.org.riskscape.engine.bind.BoundParameters;
import nz.org.riskscape.engine.bind.JavaParameterSet;
import nz.org.riskscape.engine.bind.ParamProblems;
import nz.org.riskscape.engine.bind.Parameter;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.bind.ParameterSet;
import nz.org.riskscape.engine.bind.ParameterTemplate;
import nz.org.riskscape.engine.ini.IniFileParameterBuilder;
import nz.org.riskscape.engine.model.BaseModel2;
import nz.org.riskscape.engine.model.Model;
import nz.org.riskscape.engine.model.ModelFramework;
import nz.org.riskscape.engine.model.ModelParameter;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.resource.Resource;
import nz.org.riskscape.engine.resource.ResourceLoadingException;
import nz.org.riskscape.engine.resource.StringResource;
import nz.org.riskscape.engine.util.Pair;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problem.Severity;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ExpressionParser;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.ParameterToken;
import nz.org.riskscape.rl.ast.PipelineDeclaration;
import nz.org.riskscape.util.ListUtils;

/**
 * Takes pipelines in plain-text DSL form and slots them into the model
 * framework, so that they can be run the same way as any other model.
 */
@RequiredArgsConstructor
public class ParameterizedPipelineModelFramework extends BaseModelFramework {

  private final Engine engine;

  public static class Params {
    @ParameterField
    public Optional<String> pipeline = Optional.empty();

    @ParameterField
    public Optional<URI> location = Optional.empty();
  }

  @Override
  public String getId() {
    return "pipeline";
  }

  @Getter
  private final JavaParameterSet<Params> buildParameterSet = JavaParameterSet.fromBindingClass(Params.class);

  @Override
  public ResultOrProblems<Model> build(Project project, BoundParameters buildWith) {
    Params params = buildParameterSet.bindToObject(buildWith).getBoundToObject();

    Resource pipelineSource;

    if (params.pipeline.isPresent() && params.location.isPresent()) {
      // only one of these params should be specified
      return ResultOrProblems.failed(
          ParamProblems.get().mutuallyExclusive("pipeline", "location"));
    }

    if (params.pipeline.isPresent()) {
      // TODO replace project output base location with specific ini location once we add new ini parser config model
      // stuff
      pipelineSource = new StringResource(project.getOutputBaseLocation(), params.pipeline.get());
    } else if (params.location.isPresent()) {
      try {
        pipelineSource = engine.getResourceFactory().load(params.location.get());
      } catch (RiskscapeIOException | ResourceLoadingException ex) {
        return ResultOrProblems.failed(
            Problems.foundWith(Parameter.class, "location", Problems.caught(ex)));
      }
    } else {
      return ResultOrProblems.failed(ParamProblems.oneOfTheseRequired("pipeline", "location"));
    }

    return parsePipelineDecl(pipelineSource)
        .flatMap(pipelineAst -> build(buildWith, buildWith.getContext(), pipelineAst, pipelineSource));
  }

  /**
   * Returns the parameters that the user has defined for a pipeline model
   */
  private List<PipelineModelParameter> getPipelineParameters(PipelineDeclaration parsed, Map<String, List<?>> unbound,
      BindingContext context, ProblemSink problemSink) {
    IniFileParameterBuilder parameterBuilder = new IniFileParameterBuilder(context.getProject());
    List<Problem> problems = new ArrayList<>();

    // now we need to extract the required parameters from the pipeline declaration
    List<ParameterToken> tokens = parsed.getUniqueParameters();
    List<PipelineModelParameter> collected = new ArrayList<>();

    for (ParameterToken parameterToken : tokens) {
      String paramName = parameterToken.getToken().getValue();

      // the user can optionally specify template config that specifies things like help text
      // for the parameter, what options it can accept, etc. Collect any such template config
      Map<String, List<String>> templateConfig = new HashMap<>();
      for (String key : IniFileParameterBuilder.KEYS) {
        List<?> values = unbound.get(definitionKey(paramName, key));
        if (values != null) {
          templateConfig.put(key, Lists.transform(values, v -> v.toString()));
        }
      }

      // If the user hasn't set a template, and one exists with the same name as the parameter, use that one
      boolean noSpecifiedTemplate = !templateConfig.containsKey(IniFileParameterBuilder.TEMPLATE);
      if (noSpecifiedTemplate && context.getProject().getParameterTemplates().getOr(paramName).isPresent()) {
        templateConfig.put(IniFileParameterBuilder.TEMPLATE, List.of(paramName));
      }

      ParameterTemplate template = ProblemException.catching(() ->
          parameterBuilder.build(templateConfig)
        )
        // downgrade any template errors to warnings. The model can still run without a valid
        // parameter template, it just might not constrain parameter values like the user intended
        .map(t -> t, prob -> prob.withSeverity(Severity.WARNING))
        .addProblemsTo(problems)
        .orElse(ParameterTemplate.EMPTY);

      PipelineModelParameter parameter = PipelineModelParameter.create(paramName, template);

      // apply the default value, as defined in the project.ini file
      List<?> values = unbound.get(definitionKey(paramName));
      if (values != null && values.size() > 0) {
        parameter = parameter.withDefaultValue(context, values.get(0));
        if (values.size() > 1) {
          problems.add(ParamProblems.get().wrongNumberGiven(paramName, "1", values.size())
              .withSeverity(Severity.WARNING));
        }
      } else if (template.getDefaultValue().isPresent()) {
        parameter = parameter.withDefaultValue(context, template.getDefaultValue().get());
      }
      collected.add(parameter);
    }

    problems.forEach(problemSink);

    return collected;
  }

  /**
   * @return a parameter key as defined in the project.ini file, e.g. param.foo or param.foo.bar
   */
  private String definitionKey(String paramName, String... subKeys) {
    return String.join(".", ListUtils.concat(Arrays.asList("param", paramName), Arrays.asList(subKeys)));
  }

  private ResultOrProblems<PipelineDeclaration> parsePipelineDecl(Resource pipelineSource) {
    return ExpressionParser.parseParameterizedPipeline(pipelineSource)
        // KLUDGE: We need to check for cycles before we do parameter substitution, as that is screwing up the line
        // locations of the substituted tokens.  The 'proper' fix would be to build a new string during
        // parameterization and re-parse the model (or somehow preserve the locations...)
       .flatMap(decl -> decl.checkValid(decl.getStepNameFunction(new DefaultStepNamingPolicy())));
  }

  private ResultOrProblems<Model> build(
      BoundParameters builtWith,
      BindingContext context,
      PipelineDeclaration pipeline,
      Resource pipelineSource
  ) {
    List<Problem> problems = new ArrayList<>();
    List<PipelineModelParameter> parameters = getPipelineParameters(pipeline, builtWith.getExtraneous(), context,
        p -> problems.add(p));

    // sanity-check the user's parameter names for duplicates
    for (ModelParameter param : parameters) {
      if (builtWith.getBoundTo().contains(param.getName())) {
        problems.add(GeneralProblems.get().nameAlreadyUsedBy(param.getName(), ModelFramework.class, Parameter.class));
      }
    }
    ParameterSet parameterSet = new ParameterSet(Lists.transform(parameters, mp -> mp.getParameter()));
    Map<String, List<?>> unbound = getUnboundMap(builtWith.getExtraneous(), parameters);

    return buildWithIncompleteParameters(parameterSet.bind(context, unbound).withMoreProblems(problems),
        bound -> new PipelineModelImpl(builtWith, bound, pipeline, pipelineSource, parameters)
    );
  }

  /**
   * Removes any project.ini file config (e.g. `param.foo = bar`) from the extraneous framework parameters.
   * @return the unbound parameter values that will be needed to build a BoundParameters for the model.
   */
  private Map<String, List<?>> getUnboundMap(Map<String, List<?>> extraneous, List<PipelineModelParameter> parameters) {
    // we make a new unbound map from the framework's extraneous parameters. This map will be what is
    // turned into the model's BoundParameters.
    Map<String, List<?>> unbound = Maps.newHashMap(extraneous);

    // at this point we expect the unbound map will contain:
    // - `param.PARAM_NAME` entries from the model definition. These will have been used as
    //                      the parameters' default values. We remove these, so they don't appear as
    //                      as extraneous entries in the model's BoundParameters.
    // - `param.PARAM_NAME.KEY` entries that describe any ParameterTemplate config applied to the ModelParameter.
    //                      We can remove these entries now
    // - `PARAM_NAME` entries that have come from the users --param command line options. We leave these alone
    for (ModelParameter param : parameters) {
      unbound.remove(definitionKey(param.getName()));
      IniFileParameterBuilder.KEYS.forEach(key -> unbound.remove(definitionKey(param.getName(), key)));
    }
    return unbound;
  }

  /**
   * A runnable instance of a pipeline model (which just basically wraps up a
   * {@link Pipeline} that was parsed from DSL earlier).
   */
  private class PipelineModelImpl extends BaseModel2 {

    /**
     * The {@link PipelineDeclaration} that this model was created with. The declaration may contain
     * parameters that need to be substituted to get a pipeline that can be realized.
     */
    @Getter
    private final PipelineDeclaration parameterizedPipeline;

    /**
     * The {@link PipelineDeclaration} with any parameters substituted with their actual values.
     */
    @Getter
    private PipelineDeclaration pipeline;

    PipelineModelImpl(BoundParameters typeParameters, BoundParameters bound, PipelineDeclaration parameterizedPipeline,
        Resource pipelineSource, List<PipelineModelParameter> parameters) {
      super(ParameterizedPipelineModelFramework.this, typeParameters, bound, parameters);
      this.parameterizedPipeline = parameterizedPipeline;
    }

    @Override
    public ResultOrProblems<RealizedPipeline> realize(ExecutionContext context) {
      return getBoundParameters().flatMap(boundParameters -> {
        // swap out parameters
        Map<String, Expression> toReplace = parameterizedPipeline.getUniqueParameters().stream()
          .map(token -> token.getValue())
          .map(paramName -> Pair.of(paramName, getExpression(paramName)))
          .collect(Pair.toMap());

        // sanity-check the expressions we're inserting don't have any errors in them
        List<Problem> paramProblems = new ArrayList<>();
        toReplace.keySet().forEach(paramName -> paramProblems.addAll(validateParameter(paramName)));
        if (Problem.hasErrors(paramProblems)) {
          return ResultOrProblems.failed(paramProblems);
        }

        // realize shouldn't be getting called with an incomplete set of parameters - the parameters in the pipeline ast
        // have already been slurped in to a parameter set and are all `required` - realize is only called if the
        // parameters are valid.
        pipeline = parameterizedPipeline.replaceParameters(toReplace).get();

        RealizedPipeline realized = context.realize(pipeline);

        // log any warnings associated with the pipeline
        realized = realized.drainWarnings(context.getProblemSink());

        return ResultOrProblems.of(realized);
      });
    }

    private Expression getExpression(String parameterName) {
      return getModelParameter(parameterName)
          .map(mp -> ((PipelineModelParameter) mp).getExpression(getBoundParameters())).orElse(null);
    }

    private List<Problem> validateParameter(String parameterName) {
      return getModelParameter(parameterName)
          .map(mp -> ((PipelineModelParameter) mp).validateExpression(getBoundParameters()))
          .orElse(Collections.emptyList());
    }
  }
}
