/*
 * 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.wizard.model2.input;

import static nz.org.riskscape.wizard.AskerHints.*;
import static nz.org.riskscape.wizard.model2.DslHelper.*;

import java.util.ArrayList;
import java.util.Formatter;
import java.util.List;
import java.util.stream.Collectors;

import org.locationtech.jts.geom.Lineal;

import nz.org.riskscape.engine.data.ResolvedBookmark;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.query.TupleUtils.FindOption;
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.rl.TokenTypes;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.PropertyAccess;
import nz.org.riskscape.rl.ast.StructDeclaration;
import nz.org.riskscape.wizard.ExpressionHelper;
import nz.org.riskscape.wizard.bld.PipelineChange;
import nz.org.riskscape.wizard.bld.change.AppendJoinChange;
import nz.org.riskscape.wizard.bld.dsl.PipelineChangeInput;
import nz.org.riskscape.wizard.bld.loc.ChangeLocation;
import nz.org.riskscape.wizard.bld.loc.EndOfBranchWith;
import nz.org.riskscape.wizard.survey2.DefaultQuestionSet2;

public class Geoprocessing {

  public static final String GEOPROCESS_QUESTION = "geoprocess";

  public void addQuestions(DefaultQuestionSet2 questionSet, String layerAttribute) {

    addFilterQuestions(questionSet, layerAttribute);
    addEnlargeQuestions(questionSet, layerAttribute);
    addCutQuestions(questionSet, layerAttribute);

    // let the user save/see the raw results after geoprocessing has been done
    questionSet.addQuestion("save-raw-results", Boolean.class)
      .customizeQuestion(q -> q.optionalOne())
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .then((input, response) -> {
        if (response.equals(Boolean.TRUE)) {
          String stepToHangOff = layerAttribute + "_geoprocessed";
          return PipelineChange.many(
              PipelineChange.chainStep("select({*}) as %s", stepToHangOff),
              PipelineChange.newChain("%s -> save('%s-geoprocessed')",
                  stepToHangOff, layerAttribute)
              );
        } else {
          return PipelineChange.noChange();
        }
      });

    // if we did any geoprocessing at all, strip off any work-in-progress attributes
    // we added to the tuple
    questionSet.addHiddenQuestion("strip-wip-attrs")
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .then(input -> PipelineChange.chainStep("select({ %s })", layerAttribute)
    );
  }

  private void addFilterQuestions(DefaultQuestionSet2 questionSet, String layerAttribute) {
    // somewhat unnecessary step, but useful for stats as geoprocessing is generally slow
    questionSet.addHiddenQuestion("geoprocess-start")
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .then(input ->
      PipelineChange.chainStep("select({*}) as \"geoprocess_%s\"", layerAttribute)
    );

    questionSet.addQuestion("filter", FilterType.class)
      .customizeQuestion(q -> q.optionalOne())
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .then(r -> PipelineChange.noChange());

    questionSet.addQuestion("filter-expression", Expression.class)
      .askWhenResponseIs("filter", FilterType.BY_EXPRESSION)
      .customizeQuestion(q -> q.withAnnotations(
          ExpressionHelper.TAG_EXPRESSION_TYPE, ExpressionHelper.TAG_EXPRESSION_TYPE_BOOLEAN
        ).withAskAs(FilterByAttribute.ASK_AS)
      )
      .then((input, filterExpression) -> PipelineChange.chainStep(
          "filter(%s)",
          filterExpression.toSource()
      ));

    questionSet.addQuestion("filter-bounds", ResolvedBookmark.class)
      .askWhenResponseIs("filter", FilterType.BY_LAYER_BOUNDS)
      .then((input, response) ->
        PipelineChange.chainStep("filter(intersects(%s, bounds(bookmark(%s))))",
            layerAttribute, TokenTypes.quoteText(response.getId()))
      );

    // we split the filter-layer PipelineChange into two parts
    String filterLayerJoinStep = layerAttribute + "_filter_layer_join";
    String filterLayerInputStep = "filter_layer_" + layerAttribute;
    String filterLayerCoverage = layerAttribute + "_filter_coverage";

    // first, ask the user what layer to use. We can then pull it in immediately
    questionSet.addQuestion("filter-layer", ResolvedBookmark.class)
      .askWhenResponseIs("filter", FilterType.BY_FEATURE_IN_LAYER)
      .then((input, bookmark) -> PipelineChange.newChain(
            "input(relation: %s) as %s",
            TokenTypes.quoteText(bookmark.getId()),
            filterLayerInputStep
      ));

    // now build the filter condition for the feature(s) we're interested in,
    // and then turn it into a coverage we can sample
    questionSet.addQuestion("filter-layer-condition", Expression.class)
      .customizeQuestion(q -> q
          .atLocation(ChangeLocation.atStep(filterLayerInputStep))
          .withAnnotations(
              ExpressionHelper.TAG_EXPRESSION_TYPE, ExpressionHelper.TAG_EXPRESSION_TYPE_BOOLEAN,
              ExpressionHelper.TAG_LAYER_QUESTION_NAME, "filter-layer"
          )
      )
      .dependsOn("filter-layer")
      .then((input, response) -> PipelineChange.chainStep(""
            + "filter(%s) -> "
            + "group({to_coverage(*) as %s})",
            response.toSource(),
            filterLayerCoverage
      ));

    // the second part then joins the filter-layer back to the main branch. We then
    // filter out any features in our main dataset that don't intersect the filter-layer.
    // The last select removes the filter-layer coverage once we're done with it
    questionSet.addHiddenQuestion("filter-layer-join")
      .dependsOn("filter-layer-condition")
      .then((input, response) ->
        answer -> AppendJoinChange.builder(answer)
          .append(""
              + "join(on: true).lhs as %s -> "
              + "filter(length(sample_intersections(%s, coverage: %s)) > 0) -> "
              + "select({%s})",
              filterLayerJoinStep,
              layerAttribute,
              filterLayerCoverage,
              layerAttribute
          )
          .joins(EndOfBranchWith.stepNamed(filterLayerInputStep), "%s.rhs", filterLayerJoinStep)
          .build()
      );
  }

  private void addEnlargeQuestions(DefaultQuestionSet2 questionSet, String layerAttribute) {
    // TODO change this from boolean to enum: BY_DISTANCE or BY_ATTRIBUTE, and then
    // populate backingData with available attributes in BY_ATTRIBUTE case
    questionSet.addQuestion("enlarge", Boolean.class)
      .customizeQuestion(q -> q.optionalOne())
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .thenNoChange();

    questionSet.addQuestion("enlarge-by", EnlargeBy.class)
      .askWhenResponseIs("enlarge", Boolean.TRUE)
      .then((input, enlargeBy) -> {
          Type geomType = getGeometryAttribute(input, layerAttribute).getType();
          boolean lineal = Lineal.class.isAssignableFrom(geomType.internalType());

          return PipelineChange.chainStep("enlarge(distance: %s%s, mode: '%s', remove-overlaps: %b)",
              enlargeBy.distance.toSource(),
              // distance is divided by two for lineal (line) inputs. This is a convenience for users who
              // are likely to find it more natural to state the buffered line width.
              lineal ? " / 2" : "",
              enlargeBy.mode.name(), enlargeBy.removeOverlaps);
    });
  }

  private void addCutQuestions(DefaultQuestionSet2 questionSet, String layerAttribute) {
    questionSet.addQuestion("cut", CutType.class)
      .customizeQuestion(q -> q.optionalOne())
      .askWhenResponseIs(GEOPROCESS_QUESTION, Boolean.TRUE)
      .thenNoChange();

    questionSet.addQuestion("cut-distance", Double.class)
      .askWhenResponseIs("cut", CutType.BY_DISTANCE)
      .then((input, distance) -> {
        String cutGeom = layerAttribute + "_cut_geom";
        String geomAttr = getGeometryAttributeName(input, layerAttribute);

        // split up writing the steps, so it's a bit easier to keep track of the format args
        StringBuilder sb = new StringBuilder();
        try (Formatter dsl = new Formatter(sb)) {
          dsl.format("select({*, segment(%s, %f) as %s}) -> ", layerAttribute, distance, cutGeom);
          dsl.format("unnest(%s, 'segmentID') -> ", cutGeom);
          dsl.format("select({ ");
          dsl.format(" { %s.*, segmentID, %s as %s } as %s,",
              layerAttribute, cutGeom, geomAttr, layerAttribute);
          dsl.format(" measure(%s) / measure(%s.%s) as scale_factor",
              cutGeom, layerAttribute, geomAttr);
          dsl.format(" })");
          return PipelineChange.chainStep(sb.toString());
        }
      });

    questionSet.addQuestion("cut-by-layer", ResolvedBookmark.class)
      .askWhenResponseIs("cut", CutType.BY_LAYER)
      .thenNoChange();

    questionSet.addQuestion("cut-by-layer-attributes", StructDeclaration.class)
      .askWhenResponseIs("cut", CutType.BY_LAYER)
      .customizeQuestion(q -> q.optionalMany()
          // look for attributes in the cut-by-layer relation, not the current pipeline
          .withAnnotations(ExpressionHelper.TAG_LAYER_QUESTION_NAME, "cut-by-layer"))
      .thenNoChange();

    questionSet.addQuestion("cut-by-layer-include", CuttingIncludes.class)
      .askWhenResponseIs("cut", CutType.BY_LAYER)
      .then((input, include) -> {
        ResolvedBookmark bookmark = input.getBuildState().getResponse(
            questionSet.getId(),
            "cut-by-layer",
            ResolvedBookmark.class).get();
        List<StructDeclaration> membersToCopy = input.getBuildState().getAnswersTo(
            questionSet.getId(),
            "cut-by-layer-attributes",
            StructDeclaration.class);
        String copyAttrs = getCutLayerAttributesToCopy(membersToCopy);

        // if the user wants to include all parts of the original exposure-layer geometry,
        // then we need to return the geometry difference as well as the intersecting geometry
        boolean returnDifference = include == CuttingIncludes.ALL_PARTS;
        String geomAttr = getGeometryAttributeName(input, layerAttribute);
        String geomAttrFull = layerAttribute + "." + geomAttr;

        StringBuilder sb = new StringBuilder();
        try (Formatter dsl = new Formatter(sb)) {
          // save the original geom and cut by layer
          dsl.format("select({ orig_geom: %s,", geomAttrFull);
          dsl.format(" %s: layer_intersection(%s, bookmark('%s'), {%s}, %b) ",
              layerAttribute, layerAttribute, bookmark.getId(), copyAttrs, returnDifference);
          dsl.format("}) -> ");
          // unroll the result, add in the segmentID, and calculate the ratio it was cut by
          dsl.format("unnest(%s, 'segmentID') -> ", layerAttribute);
          dsl.format("select({ ");
          dsl.format(" { %s.*, segmentID } as %s,", layerAttribute, layerAttribute);
          dsl.format(" measure(%s) / measure(orig_geom) as scale_factor", geomAttrFull);
          dsl.format("})");
          return PipelineChange.chainStep(sb.toString());
        }
    });

    // optionally automatically scale certain attributes when the geometry changes in size
    questionSet.addQuestion("scale-attributes", PropertyAccess.class)
      .dependsOn("cut")
      .customizeQuestion(q -> q.optionalMany().withAnnotations(PICK_ATTRIBUTE, NUMERIC))
      .thenForEach((input, responses) -> {
        // scale the specified attributes by the amount they have changed in size
        String scaledAttrs = responses.stream()
            .map(attr -> String.format("scale_factor * %s as %s",
                attr.toSource(), attr.getLastIdentifier().getValue()))
            .collect(Collectors.joining(", "));

        return PipelineChange.chainStep("select({ { %s.*, %s } as %s })",
            layerAttribute, scaledAttrs, layerAttribute);
      });

    // manually recompute certain attributes when the geometry changes in size
    questionSet.addQuestion("recompute", RecomputeAttribute.class)
      .dependsOn("cut")
      .customizeQuestion(q -> q.optionalMany().withAnnotations(PICK_ATTRIBUTE, NUMERIC))
      .thenForEach((input, responses) -> {
        String recomputedAttrs = responses.stream()
            .map(r -> String.format("%s as %s", r.expression.toSource(), r.attribute.getLastIdentifier().getValue()))
            .collect(Collectors.joining(", "));

        return PipelineChange.chainStep("select({ { %s.*, %s } as %s })",
            layerAttribute, recomputedAttrs, layerAttribute);
      });
  }

  private String getCutLayerAttributesToCopy(List<StructDeclaration> answers) {
    // the answers come in as regular "foo as bar" StructDeclaration expressions
    List<StructDeclaration.Member> membersToCopy = new ArrayList<>();
    for (StructDeclaration expr : answers) {
      membersToCopy.addAll(expr.getMembers());
    }

    // the layer intersections is a bit of a special case. Because the cut-layer attributes
    // aren't actually in pipeline scope yet, they need to be specified as text values
    return membersToCopy.stream()
        .map(attr -> String.format("%s: '%s'",
            attr.getAttributeName(),
            attr.getExpression().toSource()))
        .collect(Collectors.joining(", "));
  }

  /**
   * Returns the name of the first geometry member that's found within a struct.
   */
  private StructMember getGeometryAttribute(PipelineChangeInput input, String layerAttribute) {
    Struct datasetStruct = getAttributeType(input.getStruct(), layerAttribute, Struct.class).get();

    return TupleUtils.findGeometryMember(datasetStruct, FindOption.ANY_REQUIRED);
  }

  private String getGeometryAttributeName(PipelineChangeInput input, String layerAttribute) {
    StructMember geom = getGeometryAttribute(input, layerAttribute);
    return TokenTypes.quoteIdent(geom.getKey());
  }
}
