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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;

import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.referencing.crs.DefaultEngineeringCRS;
import org.locationtech.jts.geom.Geometry;

import lombok.EqualsAndHashCode;
import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.NoSuchMemberException;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.coverage.TypedCoverage;
import nz.org.riskscape.engine.data.coverage.IndexedTypedCoverage;
import nz.org.riskscape.engine.data.coverage.NearestNeighbourCoverage;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.BaseRealizableFunction;
import nz.org.riskscape.engine.function.ExpensiveResource;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.function.UntypedFunction;
import nz.org.riskscape.engine.geo.IntersectionIndex;
import nz.org.riskscape.engine.geo.IntersectionIndex.Options;
import nz.org.riskscape.engine.geo.NearestNeighbourIndex;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.relation.Relation;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.rl.agg.Accumulator;
import nz.org.riskscape.engine.rl.agg.AggregationFunction;
import nz.org.riskscape.engine.rl.agg.RealizedAggregateExpression;
import nz.org.riskscape.engine.types.CoercionException;
import nz.org.riskscape.engine.types.CoverageType;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.RelationType;
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.WithinSet;
import nz.org.riskscape.engine.util.SegmentedList;
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;

public class ToTypedCoverage extends BaseRealizableFunction implements AggregationFunction {

  /*
   * map keys to avoid typo errors
   */
  private static final String OPT_INDEX_CUT = "intersection_cut";
  private static final String OPT_INDEX_CUT_POINTS = "intersection_cut_points";
  private static final String OPT_INDEX_CUT_RATIO = "intersection_cut_ratio";
  private static final String OPT_INDEX_CUT_RATIO_POINTS = "intersection_cut_ratio_points";
  private static final String OPT_INDEX_CUT_SIZE = "intersection_cut_size";

  /**
   * A {@link Struct} defining all options that may be passed to this function.
   */
  public static final Struct ALLOWED_OPTIONS = Struct.builder()
      .add("index",
        new WithinSet(
          Types.TEXT,
          "intersection",
          "nearest_neighbour"
        )
      )
      .add("nearest_neighbour_max_distance", Types.FLOATING)
      .add(OPT_INDEX_CUT, Types.BOOLEAN)
      .add(OPT_INDEX_CUT_POINTS, Types.INTEGER)
      .add(OPT_INDEX_CUT_RATIO, Types.FLOATING)
      .add(OPT_INDEX_CUT_RATIO_POINTS, Types.INTEGER)
      .add(OPT_INDEX_CUT_SIZE, Types.INTEGER)
      .build();


  // function pointer style interface for building the two types of coverage that are supported
  private interface CoverageBuilder extends Function<Iterator<Tuple>, TypedCoverage> {}

  @EqualsAndHashCode
  private static class CoverageProducingAccumulator implements Accumulator {

    private final CoverageBuilder coverageBuilder;
    /**
     * {@link RealizedExpression} to obtain the value (tuple) to be accumulated
     */
    private final RealizedExpression valueExpression;

    private final StructMember geometryMember;
    private final SRIDSet sridSet;

    final List<Tuple> accumulated;

    CoverageProducingAccumulator(CoverageBuilder coverageBuilder, RealizedExpression valueExpression,
        StructMember geometryMember, SRIDSet sridSet) {
      this.coverageBuilder = coverageBuilder;
      this.valueExpression = valueExpression;
      this.geometryMember = geometryMember;
      this.sridSet = sridSet;

      // Coverages are often made from spatial inputs that don't correctly report the nullability of
      // the input attributes. We make everything nullable here to force SegmentedList to use a null
      // safe version.
      this.accumulated = SegmentedList.forType(Nullable.of(valueExpression.getResultType()));
    }

    @Override
    public Accumulator combine(Accumulator rhs) {
      if (rhs.isEmpty()) {
        return this;
      }
      if (isEmpty()) {
        return rhs;
      }
      CoverageProducingAccumulator combined = new CoverageProducingAccumulator(coverageBuilder, valueExpression,
          geometryMember, sridSet);
      combined.accumulated.addAll(this.accumulated);
      combined.accumulated.addAll(((CoverageProducingAccumulator)rhs).accumulated);
      return combined;
    }

    @Override
    public void accumulate(Object input) {
      Tuple toAccumulate = (Tuple) valueExpression.evaluate(input);
      if (toAccumulate == null) {
        return;
      }
      accumulated.add(toAccumulate);
    }

    @Override
    public Object process() {
      return coverageBuilder.apply(accumulated.iterator());
    }

    @Override
    public boolean isEmpty() {
      return accumulated.isEmpty();
    }
  }

  /**
   * Various bits of realization state that are common between the aggregate and scalar forms of this function
   */
  private static class RealizationState {
    SRIDSet sridSet; // used for reprojection etc
    Map<String, Object> options; // 2nd optional arg with index options
    Struct memberType; // the member type of the resulting coverage
    boolean nullable; // was the original scalar input nullable - irrelevant for aggregate functions
    StructMember geometryMember; // the member of memberType that holds the geometry we index

    /**
     * @return the CRS linked to the geometry member
     */
    Optional<CoordinateReferenceSystem> sniffCrsFromType() {
      return geometryMember.getType().findAllowNull(Referenced.class).
          map(Referenced::getCrs);
    }

    /**
     * Gets the first CRS found in a tuple from the given iterator, or return a wildcard/no-crs if the iterator is
     * empty.
     */
    CoordinateReferenceSystem getFromTuples(Iterator<Tuple> tuples, Consumer<Tuple> consumer) {
      // use the first tuple to detect the crs
      if (tuples.hasNext()) {
        Tuple firstTuple = tuples.next();
        consumer.accept(firstTuple);

        // NB this throws if the srid is 0 - this is correct.
        return sridSet.get((Geometry) firstTuple.fetch(geometryMember));
      }

      // no tuples - fall back to wildcard/none
      return DefaultEngineeringCRS.CARTESIAN_2D;
    }

    /**
     * Gets a CRS to use for the coverage we are building, first deferring to the type, then fall back to using the
     * first tuple in the given iterator.  If the iterator is empty, a wildcard engineering CRS is returned.  If the
     * iterator contains geometries in mixed CRS', then it's up to each index type how they handle it.  If at all.
     */
    CoordinateReferenceSystem nominateCrs(Iterator<Tuple> tuples, Consumer<Tuple> consumer) {
      return sniffCrsFromType()
      .orElseGet(() -> getFromTuples(tuples, consumer));
    }
  }

  public ToTypedCoverage() {
    super(
      ArgumentList.create("values", RelationType.WILD, "options", Types.ANYTHING),
      CoverageType.WILD
    );
  }

  @Override
  public RiskscapeFunction asFunction() {
    return AggregationFunction.addAggregationTo(this, super.asFunction());
  }

  // scalar version
  @Override
  public ResultOrProblems<RiskscapeFunction> realize(RealizationContext context, FunctionCall fc,
      List<Type> argumentTypes) {
    return ProblemException.catching(() -> {
      RealizationState rState = newRealizationState(context, fc);
      Type inputType = argumentTypes.get(0);
      rState.nullable = Nullable.is(inputType);

      if (inputType.findAllowNull(CoverageType.class).isPresent()) {
        // short circuit if we already have a coverage, then just return that
        return RiskscapeFunction.create(RiskscapeFunction.BUILT_IN, argumentTypes,
            inputType, args -> args.get(0));
      }

      // if values is constant, we can build a single coverage for all evaluations and save lots of cpu cycles
      Object constant = arguments.get("values").evaluateConstant(context, fc, Object.class).orElse(null);
      if (constant != null && constant instanceof String stringConstant) {
        // we can only do this for constant string args otherwise bookmark requires a type argument as well.
        RealizedExpression bookmarkExpr = context.getExpressionRealizer()
            .realize(Struct.EMPTY_STRUCT, "bookmark('" + stringConstant + "')")
            // the constant bookmark expression will fail to realize if stringConstant is not a bookmark ID or file.
            // any problems from this will be wrapped in a found with bookmark() which may be a little unexpected
            // by users but shouldn't be too alarming either.
            .getOrThrow();

        constant = bookmarkExpr.evaluate(Tuple.EMPTY_TUPLE);
      }

      if (constant != null) {
        // if the values are constant then we will try to extract the actual type from the constant value.
        // only relations and lists of tuples are supported.
        if (constant instanceof Relation) {
          // get the member type from the relation we're ultimately going to get our tuples from - if we use the
          // given argument types version, it's possibly a different instance of the same struct, which leads to the
          // 'not my struct' error when we go to access the relation's tuples later
          rState.memberType = ((Relation)constant).getType();
        } else if (constant instanceof List) {
          // if the constant is a list then we can get the member type from the first tuple. to_coverage will only work
          // with a non-empty list of *tuples*, so we don't need to worry about supporting the other
          // cases here, instead we check for a null memberType below and give the user an error about an unsupported
          // type for values
          List<?> constantList = (List<?>)constant;

          // NB we could probably support an empty list and just return an empty coverage, which might be nice from
          // the perspective of 'completeness', e.g we have empty structs, empty list type etc.
          if (! constantList.isEmpty()) {
            Object firstItem = constantList.get(0);
            if (firstItem instanceof Tuple) {
              rState.memberType = ((Tuple)firstItem).getStruct();
            }
          }
        }
      } else {
        Type listMemberType = inputType.findAllowNull(RSList.class)
            .map(RSList::getMemberType)
            .orElse(null);

        if (listMemberType != null) {
          rState.memberType = listMemberType.find(Struct.class).orElse(null);
        } else {
          RelationType relationType = inputType.findAllowNull(RelationType.class).orElse(null);
          if (relationType != null) {
            rState.memberType = relationType.getMemberType();
          }
        }
      }

      if (rState.memberType == null) {
        // if member type is still null it's because arg zero is not a supported type
        throw new ProblemException(ArgsProblems.mismatch(arguments.get(0), inputType,
              Arrays.asList(RSList.create(Struct.EMPTY_STRUCT), RelationType.WILD)));
      }

      validateMemberType(rState);

      CoverageBuilder coverageBuilder = coverageBuilder(rState);

      UntypedFunction function;
      if (constant != null) {
        function = constantFunction(context, argumentTypes, constant, coverageBuilder);
      } else {
        function = function(coverageBuilder);
      }
      return RiskscapeFunction.create(RiskscapeFunction.BUILT_IN, argumentTypes,
          Nullable.ifTrue(rState.nullable, new CoverageType(rState.memberType)), function);
    });
  }

  // aggregate version
  @Override
  public ResultOrProblems<RealizedAggregateExpression> realize(RealizationContext context, Type inputType,
      FunctionCall fc) {
    return ProblemException.catching(() -> {
      RealizationState rState = newRealizationState(context, fc);

      RealizedExpression itemExpression = context.getExpressionRealizer()
          .realize(inputType, fc.getArguments().get(0).getExpression())
          .getOrThrow(Problems.foundWith(fc.getArguments().get(0)));

      rState.memberType = itemExpression.getResultType().findAllowNull(Struct.class)
        .orElseThrow(() -> new ProblemException(
            TypeProblems.get().notStruct(fc.getArguments().get(0).getExpression(),
            itemExpression.getResultType()))
        );

      validateMemberType(rState);

      CoverageBuilder coverageBuilder = coverageBuilder(rState);

      // NB in other cases, the coverage type gets wrapped with a referenced type from the geometry member (where
      // applicable), but the function doesn't bother.  Not sure whether that's important or not, as the crs of a
      // coverage tends not to be important for realization (compared to, say, geometry operations that require geometry
      // operands to be in the same CRS before they are allowed)
      return RealizedAggregateExpression.create(inputType, new CoverageType(rState.memberType), fc,
          () -> new CoverageProducingAccumulator(coverageBuilder, itemExpression, rState.geometryMember,
              rState.sridSet));
    });
  }


  /**
   * Construct an RealizationState object for the given context and function call, setting sridSet and options
   */
  private RealizationState newRealizationState(RealizationContext context, FunctionCall fc) throws ProblemException {
    if (fc.getArguments().size() < 1 || fc.getArguments().size() > arguments.size()) {
      throw new ProblemException(ArgsProblems.get().wrongNumberRange(1, arguments.size(), fc.getArguments().size()));
    }

    RealizationState input = new RealizationState();
    input.sridSet = context.getProject().getSridSet();
    input.options = parseAndValidateOptions(arguments.getArgument(fc, "options"), context);
    return input;
  }

  /**
   * Validate the member type set in details, setting geometryMember and crs
   */
  private void validateMemberType(RealizationState details) throws ProblemException {
    details.geometryMember = TupleUtils.findGeometryMember(details.memberType, TupleUtils.FindOption.OPTIONAL);
    if (details.geometryMember == null) {
      throw new ProblemException(Problems.foundWith(arguments.get(0),
          TypeProblems.get().structMustHaveMemberType(Types.GEOMETRY, details.memberType)
      ));
    }
  }

  /**
   * @return a RiskscapeFunction that returns the same constant value every time it's called based on the tuples held
   * within `constant` (which must be either a Relation or a List of tuples).  Wraps construction in an
   * ExpensiveResource to do it lazily and provide some feedback if it's big
   */
  private UntypedFunction constantFunction(
      RealizationContext context,
      List<Type> argumentTypes,
      Object constant,
      CoverageBuilder builder
  ) {

    // defer constructing the coverage until it's needed
    ExpensiveResource<Object> lazilyBuilt =
      new ExpensiveResource<Object>(context.getProblemSink(), "build-index-", () ->
        function(builder).call(Arrays.asList(constant))
      );

    return args -> lazilyBuilt.get();
  }

  /**
   * @return an UntypedFunction that uses the given builder to construct a TypedCoverage from whatever input is given
   * from the first argument
   */
  @SuppressWarnings("unchecked")
  private UntypedFunction function(CoverageBuilder builder) {
    return args -> {
      Object values = args.get(0);
      if (values == null) {
        return values;
      }
      Iterator<Tuple> tupleIterator;
      if (values instanceof List) {
        tupleIterator = ((List<Tuple>) values).iterator();
      } else {
        tupleIterator = ((Relation) values).iterator();
      }
      return builder.apply(tupleIterator);
    };
  }

  private CoverageBuilder coverageBuilder(RealizationState rState) {
    // TODO add some warnings on unused options
    if ("nearest_neighbour".equals(rState.options.get("index"))) {
      // NB we might be OK to assumbe a double here, but this is basically the same as the cast with an unboxing anyway
      double maxDistance = ((Number) rState.options.get("nearest_neighbour_max_distance")).doubleValue();
      return getNearestNeighbourBuilder(maxDistance, rState);
    } else {
      IntersectionIndex.Options options = IntersectionIndex.defaultOptions();
      // we default cut before adding to empty as this lets intersection index choose based on the geometry type
      // that is first put into the index. (the index only allows polygons to be cut)
      options.setCutBeforeAdding(Optional.ofNullable((Boolean) rState.options.getOrDefault(OPT_INDEX_CUT, null)));
      options.setCutRatio((double) rState.options.getOrDefault(OPT_INDEX_CUT_RATIO, options.getCutRatio()));
      options.setCutPoints((long) rState.options.getOrDefault(OPT_INDEX_CUT_POINTS, options.getCutPoints()));
      options.setCutRatioPoints((long)
          rState.options.getOrDefault(OPT_INDEX_CUT_RATIO_POINTS, options.getCutRatioPoints())
      );
      options.setCutSizeMapUnits((long) rState.options.getOrDefault(OPT_INDEX_CUT_SIZE, options.getCutSizeMapUnits()));
      // Intersection index is the default index type.
      return getIntersectionCoverageBuilder(rState, options);
    }
  }

  private CoverageBuilder getIntersectionCoverageBuilder(RealizationState rState, Options options) {
    return (tuples) -> {
      IntersectionIndex index = new IntersectionIndex(rState.geometryMember, rState.sridSet, options);

      CoordinateReferenceSystem crs = rState.nominateCrs(tuples, index::insert);

      // populate the index
      tuples.forEachRemaining(t -> index.insert(t));
      index.build();

      return new IndexedTypedCoverage(rState.memberType, rState.sridSet, crs) {
        @Override
        protected StructMember getGeomMember() {
          return rState.geometryMember;
        }

        @Override
        protected IntersectionIndex getIndex() {
          return index;
        }
      };
    };
  }

  private CoverageBuilder getNearestNeighbourBuilder(double maxDistanceMetres, RealizationState rState) {
    return (tuples) -> {
      // we can't add the tuple to the index yet,  we need to work out the crs first - stash it here for later
      List<Tuple> sniffedTuple = new ArrayList<>(1);
      CoordinateReferenceSystem crs = rState.sniffCrsFromType()
          .orElseGet(() -> rState.getFromTuples(tuples, sniffedTuple::add));

      NearestNeighbourIndex index =
          NearestNeighbourIndex.metricMaxDistance(rState.geometryMember, rState.sridSet, crs, maxDistanceMetres);

      // if we sniffed a tuple, pop it in the index
      sniffedTuple.forEach(index::insert);
      tuples.forEachRemaining(t -> index.insert(t));

      return new NearestNeighbourCoverage(() -> index, rState.memberType, crs, rState.sridSet);
    };
  }

  private Map<String, Object> parseAndValidateOptions(Optional<Argument> optionsArg, RealizationContext context)
      throws ProblemException {
    if (! optionsArg.isPresent()) {
      return Collections.emptyMap();
    }
    Tuple options = optionsArg.get()
        .evaluateConstant(context, nz.org.riskscape.engine.Tuple.class, Struct.EMPTY_STRUCT)
        .getOrThrow();

    try {
      options = Tuple.coerce(ALLOWED_OPTIONS, options.toMap(), EnumSet.of(Tuple.CoerceOptions.MISSING_IGNORED));
    } catch (CoercionException | NoSuchMemberException e) {
      throw new ProblemException(Problems.foundWith(arguments.get(1), Problems.caught(e)));
    }
    if ("nearest_neighbour".equals(options.fetch("index")) && options.fetch("nearest_neighbour_max_distance") == null) {
      throw new ProblemException(Problems.foundWith(arguments.get(1),
          GeneralProblems.get().required("nearest_neighbour_max_distance")));
    }
    return options.toMap();
  }

}
