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

import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.function.Function;

import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.coverage.grid.GridCoverage2D;
import org.geotools.coverage.grid.GridEnvelope2D;
import org.geotools.gce.geotiff.GeoTiffWriter;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.locationtech.jts.geom.Geometry;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.GeometryProblems;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.Parameter;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.data.coverage.GridTypedCoverage;
import nz.org.riskscape.engine.geo.GeometryUtils;
import nz.org.riskscape.engine.pipeline.ExecutionContext;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.problem.SeverityLevel;
import nz.org.riskscape.engine.raster.VectorToRaster;
import nz.org.riskscape.engine.resource.CreateHandle;
import nz.org.riskscape.engine.rl.ExpressionRealizer;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.types.Geom;
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.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.PropertyAccess;

public class GeoTiffFormat extends BaseFormat<GeoTiffFormat.Options> {

  public interface LocalProblems extends ProblemFactory {

    /**
     * Tip for when the grid dimensions are invalid
     */
    Problem dimensionsTip();

    /**
     * When features have been ignore because they are outside of the geotiff bounds
     */
    @SeverityLevel(Problem.Severity.WARNING)
    Problem skippedFeaturesOutOfBounds(int count, URI tiffLocation);

    /**
     * When features have not be written to the GeoTIFF because their valus is either null or NaN.
     */
    @SeverityLevel(Problem.Severity.WARNING)
    Problem skippedFeaturesNullOrNan(int count, URI tiffLocation);
  }

  public static final LocalProblems PROBLEMS = Problems.get(LocalProblems.class);

  @RequiredArgsConstructor
  private static class Writer extends RiskscapeWriter {

    private final Function<Tuple, Number> valueExtractor;
    private final Function<Tuple, Geometry> geometryExtractor;
    private final VectorToRaster v2r;
    private final CreateHandle handle;
    private final ProblemSink problemSink;
    private int skippedOutOfBounds = 0;
    private int skippedNullOrNaN = 0;

    @Getter
    private URI storedAt = null;

    @Override
    public void write(Tuple value) {
      VectorToRaster.DrawFeatureResult drawResult = v2r.drawFeature(
          valueExtractor.apply(value),
          geometryExtractor.apply(value)
      );
      if (drawResult == VectorToRaster.DrawFeatureResult.OUT_OF_BOUNDS) {
        ++skippedOutOfBounds;
      } else if (drawResult == VectorToRaster.DrawFeatureResult.SKIPPED_NO_VALUE_OR_GEOMETRY) {
        ++skippedNullOrNaN;
      }
    }

    @Override
    public void close() throws IOException {
      GridCoverage2D coverage = v2r.constructCoverage("layer");

      GeoTiffWriter writer = new GeoTiffWriter(handle.getOutputStream());
      writer.write(coverage, null);
      writer.dispose();
      storedAt = handle.store();

      // warn the user if features were not drawn
      if (skippedOutOfBounds > 0) {
        problemSink.accept(PROBLEMS.skippedFeaturesOutOfBounds(skippedOutOfBounds, storedAt));
      }
      if (skippedNullOrNaN > 0) {
        problemSink.accept(PROBLEMS.skippedFeaturesNullOrNan(skippedNullOrNaN, storedAt));
      }
    }
  }

  public static class Options {

    @ParameterField
    Optional<Double> gridResolution = Optional.empty();

    @ParameterField
    Optional<Expression> bounds =  Optional.empty();

    @ParameterField
    Optional<Expression> value = Optional.empty();

    @ParameterField
    Optional<Expression> geometry = Optional.empty();

    @ParameterField
    VectorToRaster.PixelStrategy pixelStatistic = VectorToRaster.PixelStrategy.MEAN;

    @ParameterField
    Optional<GridTypedCoverage> template = Optional.empty();

    private Struct input;
    private ExpressionRealizer realizer;
    private SRIDSet sridSet;

    private final List<Problem> problems = new ArrayList<>();

    private RealizedExpression validatedValueExpression;
    private RealizedExpression validatedGeometryExpression;
    private Double validatedGridResolution;
    private ReferencedEnvelope validatedBounds;
    private GridTypedCoverage validatedTemplate;
    private CoordinateReferenceSystem validatedTargetCrs;
    private Double validatedScale;

    ResultOrProblems<Options> validate(Struct inputStruct, ExecutionContext context) {
      this.input = inputStruct;
      this.realizer = context.getRealizationContext().getExpressionRealizer();
      this.sridSet = context.getProject().getSridSet();

      this.validatedTemplate = template.orElse(null);

        // sanity-check expressions realize OK
      validateValueExpression();
      validateGeometryExpression();
      validateTargetCrs();
      validateScale();
      validateBounds();

      if (validatedBounds != null && validatedScale != null) {
      // sanity-check dimensions are OK
      VectorToRaster.getDimensions(validatedBounds, validatedScale)
              .map(d -> d, p -> p.withChildren(PROBLEMS.dimensionsTip()))
              .addProblemsTo(problems);
      }

      if (Problem.hasErrors(problems)) {
        return ResultOrProblems.failed(problems);
      } else {
        return ResultOrProblems.of(this, problems);
      }
    }

    private void validateValueExpression() {
      value
          // if it's present, realize it
          .map(val -> realizer.realize(input, val))
          // else try and guess it from the input scope
          .orElseGet(() -> expressionToFirstOfType(false))
          .addProblemsTo(problems)

          // some final validation before memoing realized result
          .ifPresent(expr -> {
            if (!expr.getResultType().isNumeric()) {
              problems.add(
                  TypeProblems.get()
                      .requiresOneOf("value", List.of(Types.FLOATING, Types.INTEGER), expr.getResultType())
              );
            } else {
              this.validatedValueExpression = expr;
            }
        });
    }

    private void validateGeometryExpression() {
      geometry
          // realize it if present
          .map(val -> realizer.realize(input, val))
          // otherwise try and guess it from the input scope
          .orElseGet(() -> expressionToFirstOfType(true))
          .addProblemsTo(problems)

          // some final validation before memoing
          .ifPresent(expr -> {
            if (expr.getResultType().findAllowNull(Geom.class).isEmpty()) {
              problems.add(TypeProblems.get().mismatch("geometry", Types.GEOMETRY, expr.getResultType()));
            } else {
              this.validatedGeometryExpression = expr;
            }
          });
    }

    private void validateTargetCrs() {
      // template's crs takes precedence
      if (validatedTemplate != null) {
        this.validatedTargetCrs = validatedTemplate.getCoordinateReferenceSystem();
        return;
      }

      // otherwise take it from the geometry we are rasterizing
      if (validatedGeometryExpression != null) {
        Type geomType = validatedGeometryExpression.getResultType();
        Referenced referenced = geomType.findAllowNull(Referenced.class).orElse(null);

        if (referenced == null) {
          problems.add(GeometryProblems.get().notReferenced(geomType));
        } else {
          this.validatedTargetCrs = referenced.getCrs();
        }
      }
    }

    private void validateBounds() {
      // see if we got a bounds from the template - we can fallback to this if `bounds` is not set
      ReferencedEnvelope templateBounds = Optional.ofNullable(validatedTemplate)
          .flatMap(GridTypedCoverage::getEnvelope)
          .orElse(null);

      if (bounds.isPresent() && validatedTargetCrs != null) {
        realizer
          .realizeConstant(bounds.get())
          .flatMap(re -> {
            // check our expression yields geometry type
            if (re.getResultType().findAllowNull(Geom.class).isEmpty()) {
              return ResultOrProblems.failed(
                      TypeProblems.get().mismatch("bounds", Types.GEOMETRY, re.getResultType())
              );
            }
            return ResultOrProblems.of((Geometry) re.evaluate(Tuple.EMPTY_TUPLE));
          }).map(boundsVal ->
            // reproject the bounds to match the referenced CRS
            // TODO this isn't going to work - you can not reproject a bounds, they need to be recomputed once
            // reprojected
            sridSet.reproject(boundsVal, sridSet.get(validatedTargetCrs)
          ))
          .addProblemsTo(problems)
          .ifPresent(boundsGeom -> {
            this.validatedBounds = new ReferencedEnvelope(boundsGeom.getEnvelopeInternal(), validatedTargetCrs);
          });

      } else if (templateBounds != null) {
        this.validatedBounds = templateBounds;
      } else {
        problems.add(GeneralProblems.required("bounds", Parameter.class));
      }
    }

    private void validateScale() {
      if (gridResolution.isPresent() && validatedTargetCrs != null) {
        this.validatedScale = 1D / GeometryUtils.toCrsUnits(gridResolution.get(), validatedTargetCrs);
      } else if (validatedTemplate != null) {
        // we infer the grid-resolution from the template GeoTIFF's bounds / pixels height.
        // This will only give us the same grid resolution if the template has square pixels
        // (it is possible to make rectangular pixels though)
        GridEnvelope2D range = validatedTemplate.getCoverage().getGridGeometry().getGridRange2D();
        // TODO what if envelope is not present?
        ReferencedEnvelope templateBounds = validatedTemplate.getEnvelope().get();
        double computedGridResolution = VectorToRaster.getHeightCrsUnits(templateBounds) / range.height;
        // note that this value is already in the correct CRS units
        this.validatedScale = 1D / computedGridResolution;
      } else {
        problems.add(GeneralProblems.required("grid-resolution", Parameter.class));
      }
    }

    ResultOrProblems<Writer> buildWriter(CreateHandle handle, ProblemSink problemSink) {
      // we should be validated by this point already
      if (Problem.hasErrors(problems)) {
        throw new AssertionError("Validation should have failed");
      }

      VectorToRaster v2r = new VectorToRaster();
      return ProblemException.catching(() -> {
        v2r.initialize(validatedBounds, validatedScale, pixelStatistic);

        // make sure we reproject to match the target CRS for the output
        int srid = sridSet.get(validatedTargetCrs);
        return new Writer(
            (t) -> (Number) validatedValueExpression.evaluate(t),
            (t) -> {
              // NB: reproject() isn't null-safe
              Geometry geom = (Geometry) validatedGeometryExpression.evaluate(t);
              if (geom == null) {
                return geom;
              }
              return sridSet.reproject(geom, srid);
            },
            v2r,
            handle,
            problemSink
        );
      });
    }

    /**
     * Builds an expression to access the first found attribute of struct of the desired type.
     * <p>
     * Desired type is determined by isGeometry (Types.GEOMETRY when set, any numeric type otherwise).
     */
    private ResultOrProblems<RealizedExpression> expressionToFirstOfType(boolean isGeometry) {
      List<String> members = findSegmentsToFirstMemberOfType(input, isGeometry);
      if (members.isEmpty()) {
        return ResultOrProblems.failed(TypeProblems.get()
                .structMustHaveMemberType(isGeometry ? Types.GEOMETRY : Types.FLOATING, input));
      }
      return realizer.realize(input, PropertyAccess.of(members));
    }

    private List<String> findSegmentsToFirstMemberOfType(Struct struct, boolean isGeometry) {
      for (StructMember member : struct.getMembers()) {
        boolean hasType = isGeometry
                ? member.getType().findAllowNull(Geom.class).isPresent()
                : member.getType().isNumeric();
        if (hasType) {
          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 = findSegmentsToFirstMemberOfType(structMember.get(), isGeometry);
          if (!nested.isEmpty()) {
            nested.add(0, member.getKey());
            return nested;
          }
        }
      }
      return Collections.emptyList();
    }
  }

  public GeoTiffFormat() {
    super("geotiff", "tif", "image/tiff");
  }

  @Override
  protected ResultOrProblems<? extends RiskscapeWriter> newWriter(ExecutionContext context, Struct input,
      CreateHandle handle, Options bound) {
    return bound.buildWriter(handle, context.getProblemSink());
  }

  @Override
  protected ResultOrProblems<Options> validateWriterOptions(ExecutionContext context, Struct input, CreateHandle handle,
      Options bound) {
    return bound.validate(input, context);
  }

  // test only constructor - we want to bypass binding - NB I think I'll get rid of this in the next commit or so
  Optional<WriterConstructor> getWriterConstructor(Options options) {
    return Optional.of((context, type, handle, formatOptions) -> {
      return options.buildWriter(handle, context.getProblemSink());
    });
  }
}
