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

import static nz.org.riskscape.rl.TokenTypes.*;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import org.locationtech.jts.geom.Geometry;

import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.Timer;
import com.google.common.collect.Lists;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.dsl.Token;
import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.HasMeter;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.BindingContext;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.expr.StructMemberAccessExpression;
import nz.org.riskscape.engine.geo.GeometryFamily;
import nz.org.riskscape.engine.geo.GeometryUtils;
import nz.org.riskscape.engine.geo.IntersectionIndex;
import nz.org.riskscape.engine.geo.OverlayOperations;
import nz.org.riskscape.engine.pipeline.Collector;
import nz.org.riskscape.engine.pipeline.RealizationInput;
import nz.org.riskscape.engine.pipeline.Realized;
import nz.org.riskscape.engine.pipeline.RealizedStep;
import nz.org.riskscape.engine.projection.Projector;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.query.TupleUtils.FindOption;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.rl.ExpressionRealizer;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.types.Geom;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Referenced;
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.util.Pair;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.problem.StandardCodes;
import nz.org.riskscape.rl.ExpressionParser;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.PropertyAccess;

/**
 * Step to enlarge (buffer) a geometry member.
 *
 * Mainly intended for use with lines that represent road centre lines.
 *
 * May optionally remove overlaps from the enlarged geometries. For each overlap (between two geometries)
 * the whole overlap will be removed from one of those geometries. As such overlap removal is not likely
 * to be completely fair. But for the intended use of enlarging road centrelines this should be sufficient.
 *
 * This functionality has been added as a step to work around not having a means to replace struct
 * members in the expression language. If that had been available enlarging (without overlap removal) could
 * be done in a function or an aggregation function when removing overlaps. Using expression language functions
 * is always preferred to added special purpose steps.
 */
public class EnlargeStep extends BaseStep<EnlargeStep.Params> {

  public static final String PROCESS_METRIC_NAME = "process";

  @RequiredArgsConstructor
  public enum EnlargeMode {
    ROUND("{vertex: 'round', cap: 'round'}"),
    BEVELED_SQUARE("{vertex: 'bevel', cap: 'square'}"),
    BEVELED_FLAT("{vertex: 'bevel', cap: 'flat'}"),
    MITRED_SQUARE("{vertex: 'mitre', cap: 'square'}"),
    MITRED_FLAT("{vertex: 'mitre', cap: 'flat'}");

    public final String mode;
  }

  public static class Params {

    @ParameterField
    Expression distance;

    @ParameterField
    Optional<PropertyAccess> geomExpression;

    @ParameterField
    Optional<EnlargeMode> mode;

    @ParameterField
    Optional<Boolean> removeOverlaps;

    @Input
    public RealizedStep input;

    public RealizationInput rInput;

  }

  /**
   * A {@link Projector} that will enlarge a geometry by applying a buffer expression.
   */
  @RequiredArgsConstructor
  static class EnlargeGeometryProjector implements Projector {
    @Getter
    private final Struct sourceType;
    @Getter
    private final Struct producedType;

    private final StructMemberAccessExpression geomAccessor;
    private final RealizedExpression bufferExpression;

    @Override
    public Tuple apply(Tuple t) {
      Object original = geomAccessor.evaluate(t);
      Object buffered = bufferExpression.evaluate(t);

      // We make a clone of the tuple so don't update the originals content
      Tuple clone = t.clone();
      if (original != null) {
        // now set the buffered geom into the clone.
        // we only do this if the original is not null because if original is null then there is
        // no need to set as buffered must also be null. But we do still set a null value because this
        // could occur if the distance was null.
        geomAccessor.setValue(clone, buffered);
      }

      // finish up by coercing clone to the expected type. this is required as the geometry type
      // is changed to the generic Types.GEOMETRY
      @SuppressWarnings("unchecked")
      Tuple projected = (Tuple)producedType.coerce(clone);
      return projected;
    }
  }

  /**
   * Accumulator that will collect a list of {@link Tuples} containing the enlarged geometries.
   *
   * The actual enlarging is deferred to the enlargingProjector.
   */
  @RequiredArgsConstructor
  static class AccumInstance {

    private final Projector enlargingProjector;

    private final List<Tuple> items = new ArrayList<>();

    public void accumulate(Tuple tuple) {
      Tuple enlarged = enlargingProjector.apply(tuple);
      items.add(enlarged);
    }

    public boolean isEmpty() {
      return items.isEmpty();
    }

    public AccumInstance combine(AccumInstance rhs) {
      if (rhs.isEmpty()) {
        return this;
      }
      if (isEmpty()) {
        return rhs;
      }
      AccumInstance combined = new AccumInstance(enlargingProjector);
      combined.items.addAll(this.items);
      combined.items.addAll(rhs.items);
      return combined;
    }
  }

  static class EnlargingCollector implements Collector<AccumInstance>, HasMeter {

    @Getter
    private final Class<AccumInstance> accumulatorClass = AccumInstance.class;

    @Getter
    private final Struct sourceType;
    @Getter
    private final Struct producedType;

    private final Projector enlargingProjector;

    /**
     * A geometry accessor that is relative to {@link #producedType}.
     */
    private final StructMemberAccessExpression geomAccessor;

    private final SRIDSet sridSet;

    private final MetricRegistry metrics = new MetricRegistry();
    private final Timer processTimer;
    private final List<String> progressMetricNames = Lists.newArrayList(PROCESS_METRIC_NAME);


    EnlargingCollector(Struct sourceType, Struct targetType, Projector enlargingProjector,
        StructMemberAccessExpression geomAccessor, SRIDSet sridSet) {
      this.sourceType = sourceType;
      this.producedType = targetType;
      this.enlargingProjector = enlargingProjector;
      this.geomAccessor = geomAccessor;
      this.sridSet = sridSet;

      this.processTimer = metrics.timer(PROCESS_METRIC_NAME);
    }

    @Override
    public AccumInstance newAccumulator() {
      return new AccumInstance(enlargingProjector);
    }

    @Override
    public void accumulate(AccumInstance accumulator, Tuple tuple) {
      accumulator.accumulate(tuple);
    }

    @Override
    public AccumInstance combine(AccumInstance lhs, AccumInstance rhs) {
      return lhs.combine(rhs);
    }

    @Override
    public TupleIterator process(AccumInstance accumulator) {
      OverlayOperations overlayOps = OverlayOperations.get();

      List<Tuple> tuples = accumulator.items;

      // All the tuples are loaded into the index with their enlarged geometries. As this makes it
      // easier to find the intersections later on.
      // Then we will start removing any overlaps which of course will update the geometry in the tuples
      // that are stored in the index. This is why the index is used to get potential intersections
      // but the geometries are checked again to get the current intersection.
      IntersectionIndex index = IntersectionIndex.withDefaultOptions(geomAccessor, sridSet);
      tuples.stream().forEach(t -> index.insert(t));
      index.build();

      for (Tuple indexed: tuples) {
        long start = System.nanoTime();
        Geometry indexedGeom = geomAccessor.evaluate(indexed, Geometry.class);
        GeometryFamily indexedGeomFamily = GeometryFamily.from(indexedGeom);
        boolean indexedGeomUpdated = false;

        // to make things a little more fairer, we alternate from which side the overlap is removed.
        // this is likely to be fairer than removing all of the overlaps from indexedGeom
        boolean removeOverlapFromLHS = false;
        for (Pair<Geometry, Tuple> potentialIntersection: index.findIntersections(indexedGeom)) {
          if (indexed.equals(potentialIntersection.getRight())) {
            // we expect to find a self intersection, but we don't process it.
            continue;
          }

          Geometry intersectingGeom = geomAccessor.evaluate(potentialIntersection.getRight(), Geometry.class);

          // We re-calculate the intersection because the geometries in the tuples may have already
          // been changed by this method. But the intersection index is only dealing with the geometry
          // as they were at insertion.
          Geometry actualIntersection = GeometryUtils.removeNonFamilyMembers(
              overlayOps.intersection(indexedGeom, intersectingGeom),
              indexedGeomFamily);

          if (actualIntersection.isEmpty()) {
            // actualIntersection could be empty because overlaps have already been removed from this
            // pair of potentially intersecting geometries.
            continue;
          } else {
            // toggle the remove flag to keep alternating where overlap is removed.
            removeOverlapFromLHS = !removeOverlapFromLHS;

            Geometry differenceFrom = removeOverlapFromLHS ? indexedGeom : intersectingGeom;
            Geometry difference = GeometryUtils.removeNonFamilyMembers(
                overlayOps.difference(differenceFrom, actualIntersection),
                indexedGeomFamily);
            if (difference.isEmpty()) {
              // We don't want to remove overlaps to the point that the geometry is empty. That can
              // cause flow on affects with sampling etc.
              continue;
            }
            if (removeOverlapFromLHS) {
              // if we are on evens we update the indexedGeom
              indexedGeomUpdated = true;
              indexedGeom = difference;
            } else {
              // other wise we update the other side. This makes things a little fairer
              geomAccessor.setValue(potentialIntersection.getRight(), difference);
            }
          }
        }

        if (indexedGeomUpdated) {
          geomAccessor.setValue(indexed, indexedGeom);
        }
        processTimer.update(System.nanoTime() - start, TimeUnit.NANOSECONDS);
      }
      return TupleIterator.wrapped(tuples.iterator(), Optional.empty());
    }

    @Override
    public MetricRegistry getRegistry() {
      return metrics;
    }

    @Override
    public List<String> getProgressMetricNames() {
      // This is a real method (rather than a @Getter) to ensure that it overrides HasMeter.getProgressMetricNames()
      return progressMetricNames;
    }
  }

  public EnlargeStep(Engine engine) {
    super(engine);
  }

  @Override
  public ResultOrProblems<? extends Realized> realize(Params parameters) {
    Struct inputType = parameters.input.getProduces();
    RealizationContext context = parameters.rInput.getExecutionContext().getRealizationContext();
    ExpressionRealizer realizer = context.getExpressionRealizer();
    BindingContext bindingContext = parameters.rInput.getBindingContext();

    return ProblemException.catching(() -> {
      // We realize the distance Expression to ensure it is good.
      RealizedExpression distanceExpression = realizer.realize(inputType, parameters.distance)
          .getOrThrow();
      if (! distanceExpression.getResultType().isNumeric()) {
        throw new ProblemException(TypeProblems.get().mismatch(parameters.distance,
            Types.FLOATING, distanceExpression.getResultType()));
      }

      List<String> geomPath;
      if (parameters.geomExpression.isPresent()) {
        geomPath = getMembersFromExpression(parameters.geomExpression.get(), inputType, realizer);
      } else {
        // If the user hasn't said what geometry to process we use the first one we can find.
        geomPath = findFirstGeometry(inputType);
        if (geomPath.isEmpty()) {
          throw new ProblemException(new Problem(Problem.Severity.ERROR, StandardCodes.GEOMETRY_REQUIRED, inputType));
        }
      }

      // Now create the buffer expression that we need.
      String bufferExpression = String.format("buffer(%s, %s, options: %s)",
          geomPath.stream()
          .map(i -> quoteIdent(i))    // path elements may need to be quoted in an expression
          .collect(Collectors.joining(".")),
          parameters.distance.toSource(),
          parameters.mode.orElse(EnlargeMode.ROUND).mode
      );

      // NB if there is a problem with the expression, it's not the user's fault - we have a coding error, so throw an
      // exception for failed parsing or realizing(don't return a problem)
      RealizedExpression rBufferExpr =
          realizer.realize(inputType, ExpressionParser.INSTANCE.parse(bufferExpression)).getOrThrow();

      StructMemberAccessExpression geomAccessor = StructMemberAccessExpression.build(inputType, geomPath)
          .getOrThrow();

      // the target type could be a little different as point and lines become polygons (maybe multipolygons)
      Struct targetType = getTargetType(inputType, geomPath).findAllowNull(Struct.class).get();

      Projector enlargingProjector =  new EnlargeGeometryProjector(inputType, targetType, geomAccessor, rBufferExpr);

      if (! parameters.removeOverlaps.orElse(false)) {
        // We are not removing overlaps to the projector can be returned.
        return enlargingProjector;
      }

      // We need a geom accessor to work with the targetType. This for the EnlargingCollector to use
      // when removing overlaps.
      StructMemberAccessExpression targetGeomAccessor = StructMemberAccessExpression.build(targetType, geomPath)
          .getOrThrow();
      return new EnlargingCollector(inputType, targetType, enlargingProjector, targetGeomAccessor,
          context.getProject().getSridSet());
    });
  }

  private Type getTargetType(Struct inputType, List<String> geomPath) {
    StructMember member = inputType.getEntry(geomPath.get(0));
    // We have to take care not to drop any nullables
    boolean nullableMember = Nullable.is(member.getType());
    Optional<Struct> isStruct = member.getType().findAllowNull(Struct.class);
    if (isStruct.isPresent()) {
      Type amended = getTargetType(isStruct.get(), geomPath.subList(1, geomPath.size()));
      return inputType.replace(member.getKey(), Nullable.ifTrue(nullableMember, amended));
    }
    // if it's not a struct member then it must be the geometry member.
    // enlarging will create polygons (or multi-polygons due to removing overlaps), so we need to
    // use the generic geometry type
    Optional<Referenced> referencedGeom = member.getType().findAllowNull(Referenced.class);
    Type amendedGeomType = referencedGeom.map(ref -> Referenced.of(Types.GEOMETRY, ref.getCrs()))
        .orElse(Types.GEOMETRY);
    return inputType.replace(member.getKey(), Nullable.ifTrue(nullableMember, amendedGeomType));
  }

  private List<String> getMembersFromExpression(PropertyAccess geomExpression, Struct inputType,
      ExpressionRealizer realizer)
      throws ProblemException {
    if (geomExpression.getReceiver().isPresent()) {
      // Not allowed.
      // TODO error
      throw new ProblemException();
    }
    RealizedExpression expr = realizer.realize(inputType, geomExpression).getOrThrow();


    List<String> members = new ArrayList<>();
    Struct toAccess = inputType;
    for (int i = 0; i < geomExpression.getIdentifiers().size(); i++) {
      Token identifier = geomExpression.getIdentifiers().get(i);
      StructMember member = toAccess.getEntry(identifier.getValue());
      members.add(member.getKey());

      Optional<Struct> nestedStruct = member.getType().findAllowNull(Struct.class);
      if (nestedStruct.isPresent()) {
        toAccess = nestedStruct.get();

      }
    }
    Optional<Struct> isStruct = expr.getResultType().find(Struct.class);
    if (isStruct.isPresent()) {
      // if the expression returns a struct then we check if it has a geometry member
      Struct struct = isStruct.get();
        StructMember geometry = TupleUtils.findGeometryMember(struct, FindOption.OPTIONAL);
        if (geometry == null) {
          throw new ProblemException(Problem.error(StandardCodes.GEOMETRY_REQUIRED, struct));
        }
        members.add(geometry.getKey());
    } else if (! expr.getResultType().find(Geom.class).isPresent()) {
      throw new ProblemException(TypeProblems.get().mismatch(geomExpression, Types.GEOMETRY, expr.getResultType()));
    }
    return members;
  }

  /**
   * Traverse the struct and return a property access to the first geometry member that is found.
   * @param struct
   * @return
   */
  List<String> findFirstGeometry(Struct struct) {

    for (StructMember member : struct.getMembers()) {
      Optional<Geom> geomType = member.getType().findAllowNull(Geom.class);
      if (geomType.isPresent()) {
        List<String> found = new ArrayList<>();
        found.add(member.getKey());
        return found;
      }
      Optional<Struct> structMember = member.getType().findAllowNull(Struct.class);
      if (structMember.isPresent()) {
        List<String> nested = findFirstGeometry(structMember.get());
        if (! nested.isEmpty()) {
          nested.add(0, member.getKey());
          return nested;
        }
      }
    }

    return Collections.emptyList();
  }

}
