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

import java.awt.Rectangle;
import java.awt.geom.Rectangle2D;
import java.awt.image.Raster;
import java.io.File;
import java.util.concurrent.ExecutionException;

import javax.media.jai.PlanarImage;

import org.geotools.api.coverage.CannotEvaluateException;
import org.geotools.api.coverage.PointOutsideCoverageException;
import org.geotools.api.coverage.grid.GridCoverage;
import org.geotools.api.geometry.MismatchedDimensionException;
import org.geotools.api.geometry.Position;
import org.geotools.api.referencing.operation.MathTransform2D;
import org.geotools.api.referencing.operation.TransformException;
import org.geotools.coverage.grid.GridCoordinates2D;
import org.geotools.coverage.grid.GridCoverage2D;
import org.geotools.coverage.grid.InvalidGridGeometryException;
import org.geotools.coverage.util.CoverageUtilities;
import org.geotools.geometry.Position2D;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;

import it.geosolutions.jaiext.range.NoDataContainer;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.problem.Problems;

/**
 * Extends GridCoverage2D to make use of {@link SparseTIFFImageReader#isEmptyTile(int, int)}.  Using this method
 * allows us to avoid bringing empty tiles in to JAI's tile cache, which uses RAM and processing time.  Instead we can
 * short circuit the usual pixel lookup and return NO_DATA when an empty tile is queried.
 */
@Slf4j
public class SparseTiffCoverage extends GridCoverage2D {

  private final SparseTIFFImageReader imageReader;

  @Getter
  private final double[] noData;

  // profiling shows it's worth memoizing these two fields
  private final MathTransform2D crsToGrid2D;
  private final Rectangle imageBounds;

  // Memoed tile metadata from the image - profiling shows that we spend a lot of time computing these, makes a big
  // difference when doing a lot of sampling if we memo them and call the static version of PlanarImage#XToTileX instead
  private final int tileGridXOffset;
  private final int tileGridYOffset;
  private final int tileWidth;
  private final int tileHeight;

  // local tile cache key
  record XY(int x, int y) {};

  // l1 cache - more often than not we are sampling the same tile over and over
  // again, especially when evaluating all intersections.  Sun's JAI Tile cache is pretty slow compared to a modern
  // implementation like this which uses read/write locks and segmented hash tables to avoid lock conflicts.
  // Note that this seems to have only a marginal impact on run times, despite looking impressive when looking at
  // profiling results.  But then I was only able to run the tests on my laptop
  // which only has two 'fast' cores.  It would be good to do some testing on a 16/32 core server class host and see
  // what impact it has there.
  private final Cache<XY, Raster> cache;

  public static GridCoverage createIfSupported(Project project, GridCoverage coverage, File file) {

    if (coverage instanceof GridCoverage2D grid2d) {
      try {
        NoDataContainer noData = CoverageUtilities.getNoDataProperty(grid2d);

        if (noData == null) {
          log.info("Sparse tiff support not enabled, missing nodata: {}", file);
          return coverage;
        }
      } catch (Exception e) {
        log.warn("Sparse tiff support not enabled, exception: {}", file, e);
        return coverage;
      }

      log.info("Using sparse tiff support for {}", file);
      return new SparseTiffCoverage(project, grid2d.getName(), grid2d);
    }

    log.info("Sparse tiff support not enabled for {}", file);
    return coverage;
  }

  public SparseTiffCoverage(Project project, CharSequence name, GridCoverage2D coverage) {
    super(name, coverage);

    // We need to initialize the reader with knowledge of which image from the TIFF is being sampled (TIFFs can contain
    // many actual images, and this info is not exposed in the GridCoverage2D)
    this.imageReader = SparseTIFFImageReader.initialize(this);
    this.noData = CoverageUtilities.getNoDataProperty(coverage).getAsArray();
    this.crsToGrid2D = getGridGeometry().getCRSToGrid2D();
    this.imageBounds = image.getBounds();

    this.tileGridXOffset = image.getTileGridXOffset();
    this.tileGridYOffset = image.getTileGridYOffset();
    this.tileWidth = image.getTileWidth();
    this.tileHeight = image.getTileHeight();

    int numPipelineThreads = project.getEngine().getPipelineExecutor().getNumThreads();

    this.cache = CacheBuilder.newBuilder()
        // soft values mean they can be garbage collected - should avoid high concurrency or a large number of coverages
        // ruining things
        .softValues()
        // this splits the cache's internal in to this many 'segments', each with their own locks etc
        .concurrencyLevel(numPipelineThreads)
        // keep twice as many tiles as the number of threads to allow a bit of swapping in and out to occur
        .maximumSize(numPipelineThreads * 2)
        .build();
  }

  /**
   * Overrides the evaluate method we use from {@link GridCoverage2D} to do the isEmptyTile short circuit.
   *
   * **Important** - this method returns noData instead of throwing a {@link PointOutsideCoverageException} if the point
   * is out of bounds.  Throwing an exception is expensive and we always swallow these exceptions and turn them in to
   * no-data anyway.
   */
  @Override
  public double[] evaluate(Position coord, double[] dest) {
    Position2D translated = new Position2D();
    try {
      crsToGrid2D.transform(coord, translated);
    } catch (MismatchedDimensionException | InvalidGridGeometryException | TransformException e) {
      throw new CannotEvaluateException("Could not transform coordinate", e);
    }

    final double fx = translated.getX();
    final double fy = translated.getY();
    if (!Double.isNaN(fx) && !Double.isNaN(fy)) {
      final int x = (int) Math.round(fx);
      final int y = (int) Math.round(fy);

      return getPixel(x, y, dest);
    }

    return noData;
  }

  /**
   * Overrides the evaluate method for GridCoordinates2D to do the isEmptyTile short circuit.
   * This method is used by GridTypedCoverage and SparseTIFFImageReader.
   *
   * **Important** - this method returns noData instead of throwing a {@link PointOutsideCoverageException} if the point
   * is out of bounds.  Throwing an exception is expensive and we always swallow these exceptions and turn them in to
   * no-data anyway.
   */
  @Override
  public double[] evaluate(GridCoordinates2D coord, double[] dest) {
    return getPixel(coord.x, coord.y, dest);
  }

  private double[] getPixel(int x, int y, double[] dest) {
    if (imageBounds.contains(x, y)) {
      int tileX = PlanarImage.XToTileX(x, tileGridXOffset, tileWidth);
      int tileY = PlanarImage.YToTileY(y, tileGridYOffset, tileHeight);

      // short circuit - calling image#getTile for an uncached tile is expensive and pointless when it's full of
      // no-data
      if (imageReader.isEmptyTile(tileX, tileY)) {
        return noData;
      }

      Raster tile;
      try {
        tile = cache.get(new XY(tileX, tileY), () -> {
          return image.getTile(tileX, tileY);
        });
      } catch (ExecutionException e) {
        throw new RiskscapeException(Problems.caught(e));
      }

      return tile.getPixel(x, y, dest);

    } else {
      return noData;
    }
  }

  void sampleAPixelForInit() {
    Rectangle2D rectangle2d = getGridGeometry().getGridRange2D().getBounds2D();
      super.evaluate(new GridCoordinates2D((int) rectangle2d.getCenterX(), (int) rectangle2d.getCenterY()),
          new double[2]);
  }
}
