/*
 * 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.List;
import java.util.function.Function;

import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryFactory;

import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.GeometryProblems;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.coverage.TypedCoverage;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.BaseRealizableFunction;
import nz.org.riskscape.engine.function.RiskscapeFunction;
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.types.CoverageType;
import nz.org.riskscape.engine.types.Nullable;
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.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.FunctionCall;

/**
 * Gets the bounding box (BBOX) as a {@link Geometry} from a {@link Referenced}.
 */
public class Bounds extends BaseRealizableFunction {

  /**
   * @return a geometry that covers that extent of the referenced envelope.
   */
  public static Geometry toBoundsGeom(ReferencedEnvelope env, SRIDSet sridSet) {
    return toBoundsGeom(env, env.getCoordinateReferenceSystem(), sridSet);
  }

  /**
   * @return a geometry that covers that extent of the envelope.
   */
  public static Geometry toBoundsGeom(Envelope env, CoordinateReferenceSystem crs, SRIDSet sridSet) {
    GeometryFactory gf = sridSet.getGeometryFactory(crs);
    return gf.toGeometry(env);
  }

  public Bounds() {
    super(
      ArgumentList.create("relation", RelationType.WILD),
      Types.GEOMETRY
    );
  }

  @Override
  public ResultOrProblems<RiskscapeFunction> realize(RealizationContext context, FunctionCall functionCall,
      List<Type> givenTypes) {

    return ProblemException.catching(() -> {
      getArguments().validateArity(givenTypes);

      SRIDSet sridSet = context.getProject().getSridSet();

      Object constant = arguments.get("relation").evaluateConstant(context, functionCall, Object.class).orElse(null);

      Type boundsArgType;
      // we introduce an argMapper because if we use a bookmark expression we will need to use that
      // to get the coverage (because we can't use the return type to get the bounds)
      Function<Object, Object> argMapper;
      if (constant != null && constant instanceof String stringConstant) {
        // bounds has been given a constant string argument. we treat this as a shortcut for
        // bounds(bookmark(value))
        // 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();

        boundsArgType = bookmarkExpr.getResultType();
        Object bookmarkResult = bookmarkExpr.evaluate(Tuple.EMPTY_TUPLE);

        // note, the argMapper is only used in the coverage case. relations will always use the constant value
        // if it is set.
        argMapper = (arg) -> bookmarkResult;
        constant = bookmarkResult;
      } else {
        boundsArgType = givenTypes.get(0);
        argMapper = Function.identity();
      }

      if (boundsArgType.find(CoverageType.class).isPresent()) {
        // coverages have no Referenced type, so we return the envelope of the TypedCoverage args
        return RiskscapeFunction.create(this, givenTypes, Types.POLYGON, args -> {
          TypedCoverage coverage = (TypedCoverage) argMapper.apply(args.get(0));
          return coverage.getEnvelope()
            .map(bounds -> toBoundsGeom(bounds, sridSet))
            .orElse(null);
        });

      } else if (!boundsArgType.find(RelationType.class).isPresent()) {
        // Only RelationType and CoverageType arguments are supported
        throw new ProblemException(ArgsProblems.mismatch(arguments.get(0), boundsArgType));
      }

      // check we weren't given a non-spatial relation, e.g. CSV without any geometry
      Referenced referencedType = boundsArgType.find(Referenced.class)
          .orElseThrow(() -> new ProblemException(GeometryProblems.get().notReferenced(boundsArgType)));
      Type returnType = referencedType.wrapNullable(Types.POLYGON);

      if (constant != null) {
        if (referencedType.getBounds() == null) {
          // the relation itself is constant here, so we only need to calculate the bounds once
          Geometry bounds = calculateBounds((Relation)constant, referencedType.getCrs(), sridSet);
          return RiskscapeFunction.create(this, givenTypes, Nullable.of(returnType), args -> bounds);
        } else {
          // with a constant relation args, we can just pull the bounds out of the Referenced type
          Geometry constantBounds = toBoundsGeom(referencedType.getBounds(), sridSet);
          return RiskscapeFunction.create(this, givenTypes, returnType, args -> constantBounds);
        }
      }

      return RiskscapeFunction.create(this, givenTypes, Nullable.of(returnType), args -> {
        Relation relation = (Relation) args.get(0);
        return calculateBounds(relation, referencedType.getCrs(), context.getProject().getSridSet());
      });
    });
  }

  private static Geometry calculateBounds(Relation relation, CoordinateReferenceSystem crs, SRIDSet sridSet) {
      // Relations from CSV will not have  bounds calculated until they are required.
      // TODO consider if we should log a problem is bounds is not calculated
      return relation.calculateBounds()
          .map(envelope -> toBoundsGeom(envelope, crs, sridSet))
          .orElse(null);
  }
}
