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

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import org.geotools.data.DataUtilities;
import org.geotools.api.data.FeatureSource;
import org.geotools.api.data.SimpleFeatureSource;
import org.geotools.geometry.jts.JTS;
import org.geotools.referencing.CRS;
import org.geotools.referencing.CRS.AxisOrder;
import org.geotools.referencing.crs.DefaultEngineeringCRS;
import org.geotools.util.factory.Hints;
import org.locationtech.jts.geom.CoordinateSequenceFactory;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryFactory;
import org.locationtech.jts.geom.PrecisionModel;
import org.locationtech.jts.geom.impl.CoordinateArraySequenceFactory;
import org.locationtech.jts.operation.valid.IsValidOp;
import org.geotools.api.feature.type.FeatureType;
import org.geotools.api.referencing.FactoryException;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.api.referencing.operation.MathTransform;
import org.geotools.api.referencing.operation.TransformException;

import lombok.Getter;
import lombok.NonNull;
import lombok.Setter;
import nz.org.riskscape.engine.geo.GeometryFixer;
import nz.org.riskscape.engine.geo.GeometryValidation;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.Problems;

/**
 * SRIDSet is the "central authority" in riskscape for managing CoordinateReferenceSystems.  It's responsible for:
 *
 * - maintaining a map of `CoordinateReferenceSystem`s to opaque `srid` integers, dealing with what constitutes an
 *   'equals' CRS for the purposes of RiskScape
 * - Caching a set of maths transformations used for reprojecting geometry between known crs objects, plus a
 *   convenience method for consistently reprojecting geometry
 * - Getting GeometryFactory objects that will return Geometry with the correct SRID so that this SRIDSet can
 *   correctly link it to a CoordinateReferenceSystem
 *
 * TODO the API on this is horrible - it isn't clear about the side effects - rename methods please to
 * make side effects
 */
public class SRIDSet {

  /**
   * Wraps a CRS to change its equality method to ignore metadata - this allows us to use the built in per-key thread
   * safety that comes with a concurrent hash map when seeing if a crs has been assigned an srid or not
   */
  private static class MetaDataFreeCRS {

    private final CoordinateReferenceSystem crs;
    private final int hashCode;

    MetaDataFreeCRS(CoordinateReferenceSystem crs, int hashCode) {
      this.crs = crs;
      this.hashCode = hashCode;
    }

    @Override
    public boolean equals(Object obj) {
      return CRS.equalsIgnoreMetadata(crs, ((MetaDataFreeCRS)obj).crs);
    }

    @Override
    public int hashCode() {
      return hashCode;
    }
  }

  /**
   * Geometry's with this srid have not had an srid assigned by an SridSet.
   */
  public static final int UNSET_SRID = 0;

  /**
   * World coordinate system in lon/lat (x/y).
   */
  public static final CoordinateReferenceSystem EPSG4326_LONLAT = epsgToCrsWithForceXY("EPSG:4326");

  /**
   * World coordinate system in lat/lon (y/x). Caution should be taken with this CRS because it has previously
   * been defined as lon/lat and many other software packages still treat it that way.
   */
  public static final CoordinateReferenceSystem EPSG4326_LATLON = epsgToCrs("EPSG:4326");

  /**
   * CRS for use on mainland New Zealand
   */
  public static final CoordinateReferenceSystem EPSG2193_NZTM = epsgToCrs("EPSG:2193");

  /**
   * A 2D wild card CRS.
   */
  public static final CoordinateReferenceSystem WILDCARD_2D = DefaultEngineeringCRS.GENERIC_2D;

  /**
   * Old grid crs used for NZ before NZTM.  You pretty much always want NZTM, but some test code wants to use this, so
   * here it is.
   */
  public static final CoordinateReferenceSystem EPSG_27200_NZMG = epsgToCrs("EPSG:27200");

  /**
   * @param epsgCode code to get the CRS for
   * @return CRS
   * @throws RiskscapeException if epsgCode does not identify a known CRS
   */
  public static CoordinateReferenceSystem epsgToCrs(String epsgCode) {
    return epsgToCrs(epsgCode, false);
  }

  /**
   * Similar to {@link #epsgToCrs(java.lang.String) } except that a {@link CoordinateReferenceSystem} with
   * an XY axis order will be returned regardless of how the CRS is defined.
   *
   * @param epsgCode code to get the CRS for
   * @return CRS
   * @throws RiskscapeException if epsgCode does not identify a known CRS
   */
  public static CoordinateReferenceSystem epsgToCrsWithForceXY(String epsgCode) {
    return epsgToCrs(epsgCode, true);
  }

  /**
   * @param epsgCode code to get the CRS for
   * @param forceXY when true axis order is forced to XY order despite what the CRS says
   * @return CRS
   * @throws RiskscapeException if epsgCode does not identify a known CRS
   */
  private static CoordinateReferenceSystem epsgToCrs(String epsgCode, boolean forceXY) {
    try {
      return CRS.decode(epsgCode, forceXY);
    } catch (FactoryException e) {
      throw new RiskscapeException(Problems.caught(e));
    }
  }

  // we maintain maps in both directions to allow simple fast lookup in either direction, delegating all locking and
  // sync to the concurrent hashmap

  // The authorative map SRID to CRS. Note that the CRS is the CRS instance that caused the SRID to
  // be assigned.
  protected final Map<Integer, CoordinateReferenceSystem> crsById = new ConcurrentHashMap<>();
  // first level cache which can quickly return a SRID for a CRS without having to call the expensive
  // `crsHashCode` method
  protected final Map<CoordinateReferenceSystem, Integer> sridByCrs = new ConcurrentHashMap<>();
  // second level cach which can return a SRID for CRS that is equivalent to a CRS that is already in the
  // cache. To be equivalent the CRS is equal when metadata is ignored. This means that CRSs that are
  /// effectively the same, will get be assigned the same SRID by RiskScape
  // Checking for this equivalence is quite expensive so we don't want to check it all the time.
  protected final Map<MetaDataFreeCRS, Integer> sridByMetadataFreeCrs = new ConcurrentHashMap<>();

  // this hands out unique srids without any locking
  private final AtomicInteger nextSridCounter = new AtomicInteger(1);

  /**
   * Controls if geometries should be validated after a re-projection and what to do if they have
   * become invalid.
   */
  @Getter @Setter
  private GeometryValidation validationPostReproject = GeometryValidation.OFF;

  private PrecisionModel precisionModel = new PrecisionModel();

  private CoordinateSequenceFactory crsFactory = CoordinateArraySequenceFactory.instance();

  protected final Map<List<Integer>, MathTransform> cachedTransforms = new ConcurrentHashMap<>();

  protected final Map<Integer, GeometryFactory> cachedGeometryFactories = new ConcurrentHashMap<>();

  private final ProblemSink problemSink;

  private final GeometryFixer geometryFixer;

  public SRIDSet() {
    this(ProblemSink.DEVNULL);
  }

  public SRIDSet(ProblemSink problemSink) {
    this(problemSink, GeometryFixer.DEFAULT);
  }

  /**
   * Constructor only intended for use from tests that need to mock the geometryFixer.
   */
  SRIDSet(ProblemSink problemSink, GeometryFixer geometryFixer) {
    this.problemSink = problemSink;
    this.geometryFixer = geometryFixer;
  }

  public CoordinateReferenceSystem get(int id) {
    if (id == UNSET_SRID) {
      throw new IllegalArgumentException("id is 0 - this means no SRID, did the geometry get built from a factory from"
          + "this SRIDSet?");
    }

    CoordinateReferenceSystem found = crsById.get(id);
    if (found == null) {
      throw new UnknownSRIDException(id);
    }

    return found;
  }

  public int get(@NonNull FeatureSource<?, ?> fs) {
    return get(fs.getSchema());
  }

  public int get(@NonNull FeatureType type) {
    return get(type.getCoordinateReferenceSystem());
  }

  public int get(@NonNull CoordinateReferenceSystem crs) {
    // NB: We map twice here - once by crs and then by metadatafreecrs - this is to avoid calling our custom
    // hashcode method for CRSs which is very expensive.  The first time we see a CRS, we will need to compute
    // its identifier-based-hashcode, which is expensive, and then any subsequent lookups of this CRS will be
    // based on the much cheap CRS hashcode method.  See GL601 for more details
    int srid = sridByCrs.computeIfAbsent(crs, newCrs -> {
      return sridByMetadataFreeCrs.computeIfAbsent(new MetaDataFreeCRS(crs, crsHashCode(crs)), newMetaDataFreeCrs -> {
        // no srid has been mapped to this crs, fetch the next free srid ...
        int nextSrid = nextSridCounter.getAndIncrement();
        // ...  update the map in the other direction ...
        crsById.put(nextSrid, crs);
        // ... and return the mapping
        return nextSrid;
      });
    });

    return srid;
  }

  /**
   * Generate a hashCode from the identifing parts of crs. This should return the same hashcode value for
   * crs inputs that are equal when metadata is ignored.
   *
   * This method only exists to allow tests to alter this behaviour.
   * @param crs CRS to generate hash code for
   * @return the hash code
   */
  int crsHashCode(CoordinateReferenceSystem crs) {
    // we need to return the same hashCode for objects that are equal. But crs.hashCode() won't do
    // that if the CRS's are only equal when ignoring the metadata. So we build a hashCode up front
    // from the identifing parts of the CRS
    // - epsgName
    // - axisOrder
    try {
      String epsgName = CRS.lookupIdentifier(crs, false);
      AxisOrder axisOrder = CRS.getAxisOrder(crs);

      return Objects.hash(epsgName, axisOrder);
    } catch (FactoryException e) {
      return crs.hashCode();
    }
  }

 /**
  * Clear all cached state from this SRIDSet
  * See https://gitlab.catalyst.net.nz/riskscape/riskscape/-/issues/389 for some background
  */
  public void clear() {
    sridByMetadataFreeCrs.clear();
    sridByCrs.clear();
    crsById.clear();
    cachedTransforms.clear();
    cachedGeometryFactories.clear();
    nextSridCounter.set(1);
  }

  public CoordinateReferenceSystem get(@NonNull Geometry geometry) {
    return get(geometry.getSRID());
  }

  public void remember(@NonNull SimpleFeatureSource fs) {
    get(fs);
  }

  public org.geotools.api.data.Query queryWithHints(
      @NonNull org.geotools.api.data.Query query,
      @NonNull CoordinateReferenceSystem crs) {

    Hints hints = new Hints();
    hints.put(Hints.JTS_GEOMETRY_FACTORY, getGeometryFactory(crs));
    // We only want to read 2D data. This means we don't get any Z or M values from shapefiles which we
    // don't use anyway. This also prevents issues when these geometries are segmented. When geometries
    // with XYZM coordinates are intersected with geometries that only have XY you can get a mixed bag
    // of coordinate types back from JTS. This mixed bag caused trouble for later geometry operations
    // (refer to GL484)
    hints.put(Hints.FEATURE_2D, true);
    org.geotools.api.data.Query newQuery = new org.geotools.api.data.Query();
    newQuery.setHints(hints);
    return DataUtilities.mixQueries(query, newQuery, null);
  }

  /**
   * Get a {@link MathTransform} that can be used to re-project from sourceCrs to targetCrs.
   *
   * Getting the {@link MathTransform} may be preferable to calling {@link #reproject(Geometry, int) }
   * where the calling code also needs to do an inverse transform.
   *
   * If a datum shift aware transformation cannot be obtained, then a transform that ignores datum shifts
   * will be returned. In this case a warning will be output, but note that the warning will be output once,
   * not every time that that transformation is used.
   *
   * @param sourceSrid the source SRID
   * @param targetSrid the target SRID
   * @return a math transform
   */
  public MathTransform getReprojectionTransform(int sourceSrid, int targetSrid) {
    // a pair of integers will produce a much more efficient cache key than two complex objects like a
    // CoordinateReferenceSystem
    List<Integer> cacheKey = Arrays.asList(sourceSrid, targetSrid);

    return cachedTransforms.computeIfAbsent(cacheKey, (key) -> {
      CoordinateReferenceSystem sourceCrs = get(sourceSrid);
      CoordinateReferenceSystem targetCrs = get(targetSrid);

      try {
        // first we try to find a transform with lenient:false. we want a datum aware transform
        return CRS.findMathTransform(sourceCrs, targetCrs, false);
      } catch (FactoryException e) {
        try {
          // one of the input CRS's may have been missing the bursa-wolf parameters that are required
          // for accurate datum shifts. this can occur when the CRS has come from a *.prj file.
          // we fall back to returning a transform that ignores datum shifts, but we warn the user
          MathTransform transform = CRS.findMathTransform(sourceCrs, targetCrs, true);
          if (!transform.isIdentity()) {
            problemSink.log(GeometryProblems.get().reprojectionIgnoringDatumShift(sourceCrs, targetCrs));
          }
          return transform;
        } catch (FactoryException ee) {
          throw new RuntimeException(e);
        }
      }
    });
  }

  /**
   * Reprojects geom to be in the targetSrid.
   * @param geom
   * @param targetSrid that desired srid
   * @return geom in targetSrid or the input geom if already in targetSrid
   * @throws GeometryReprojectionException if the reprojection is not possible. Normally because the geom
   *         does not wholly fit into the new CRS
   * @throws GeometryInvalidException if the reprojected geometry is not valid (and geometry validation is
   *         set to error)
   */
  public Geometry reproject(Geometry geom, int targetSrid) throws GeometryReprojectionException {
    if (geom.getSRID() == targetSrid || geom.getSRID() == 0) {
      // we skip the reprojection if the geom is already in the right projection and also if it's SRID is zero.
      // Zero is the unknown SRID so it's unknown how to reproject it anyway. Most likely it has come from
      // create_point.
      return geom;
    }

    CoordinateReferenceSystem targetCrs = get(targetSrid);
    CoordinateReferenceSystem sourceCrs = get(geom.getSRID());

    try {

      MathTransform transform = getReprojectionTransform(geom.getSRID(), targetSrid);

      return reproject(geom, transform, sourceCrs, targetCrs);
    } catch (IllegalArgumentException | TransformException ex) {
      throw new GeometryReprojectionException("Geometry could not be reprojected", geom, targetCrs, ex);
    }
  }

  /**
   * Reprojects geom using the supplied {@link MathTransform} and applies standard rules for the fixing
   * and logging of geometry that may become invalid during the reprojection.
   *
   * Note that this method is only intended for use in special cases that would also use
   * {@link #getReprojectionTransform(int, int) }.
   *
   * @param geom      geometry to be reprojected
   * @param transform math transform used to do the reprojection
   * @param sourceCrs the CRS that geom is in. Used for problem logging
   * @param targetCrs the CRS that geom is reprojected to. Used for problem logging
   * @return the reprojected geometry
   * @throws TransformException should the transformation not be possible.
   */
  public Geometry reproject(Geometry geom, MathTransform transform,
    CoordinateReferenceSystem sourceCrs, CoordinateReferenceSystem targetCrs) throws TransformException {

    GeometryFactory targetFactory = getGeometryFactory(targetCrs);

    // make a copy of the reprojected geom using the correct geometry factory to make sure the geometry factory and
    // srid are all set correctly.  Previously we just set the srid and didn't do the copy, which wouldn't set the
    // srid for multi geoms or correctly set the geometry factory.
    // TODO check how efficient (or otherwise) this is and see if there's a way to avoid the deep copy just so that
    // the srid and geometry factory are set correctly
    Geometry reprojected = targetFactory.createGeometry(JTS.transform(geom, transform));

    if (validationPostReproject != GeometryValidation.OFF) {
      IsValidOp isValidOp = new IsValidOp(reprojected);
      if (! isValidOp.isValid()) {
        // we can try to automagically fix the geometry.
        Geometry fixed = geometryFixer.fix(reprojected);
        if (fixed != null) {
          // XXX We don't have any useful context here for the user to point them at what has caused the problem.
          // Having some sort of "evaluation context" would make sense, but it doesn't exist yet and adding it
          // would require a fairly large amount of API surgery to thread it through everywhere, and currently it
          // just isn't worth it for this one thing.  We might, in the future, want some sort of
          // `ExecutionContext#getEvaluationContext` type of thing that could use maybe a thread local or what have
          // you to return a sensible object that described the context in which execution was happening, e.g. what's
          // the function call, what's the pipeline step, what's the source code unit etc
          problemSink.log(GeometryProblems.get()
              .fixedInvalidPostReprojection(sourceCrs, targetCrs, reprojected, fixed));
          return fixed;
        }
        // the reprojected geometry is not valid we should warn the user, but we only do this if the input
        // was valid. It isn't the reprojection's fault if the input was bad to start with. We assume that
        // the user has chossen to ignore bad geometries if it has gotten this far.
        if (geom.isValid()) {
          Problem invalidProblem = GeometryProblems.get()
              .invalidPostReprojection(sourceCrs, targetCrs);

          if (validationPostReproject == GeometryValidation.ERROR) {
            throw new GeometryInvalidException(invalidProblem);
          } else {
            problemSink.log(invalidProblem.withSeverity(Problem.Severity.WARNING));
          }
        }
      }
    }
    return reprojected;
  }

  /**
   * Test if {@link Geometry}s in the input {@link CoordinateReferenceSystem}s would need to be reprojected
   * before geometry operations could be performed accurately.
   * @return true if {@link Geometry} in these CRSs would need to reprojected before doing geometry operations
   */
  public boolean requiresReprojection(CoordinateReferenceSystem crs1, CoordinateReferenceSystem crs2) {
    return ! CRS.equalsIgnoreMetadata(crs1, crs2);
  }

  /**
   * @return a {@link GeometryFactory} that returns {@link Geometry} objects that are assigned an srid that
   * is mapped to the given {@link CoordinateReferenceSystem}.
   */
  public GeometryFactory getGeometryFactory(CoordinateReferenceSystem crs) {
    return cachedGeometryFactories.computeIfAbsent(get(crs), key ->
      new GeometryFactory(precisionModel, key, crsFactory));
  }

}
