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

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.geotools.api.data.Query;
import org.geotools.api.data.SimpleFeatureSource;
import org.geotools.factory.CommonFactoryFinder;
import org.geotools.filter.text.cql2.CQLException;
import org.geotools.filter.text.ecql.ECQL;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.api.style.FeatureTypeStyle;
import org.locationtech.jts.geom.Geometry;
import org.geotools.api.feature.simple.SimpleFeature;
import org.geotools.api.feature.simple.SimpleFeatureType;
import org.geotools.api.feature.type.AttributeDescriptor;
import org.geotools.api.feature.type.GeometryDescriptor;
import org.geotools.api.filter.Filter;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.data.relation.FeatureSourceBookmarkResolver;
import nz.org.riskscape.engine.relation.RenameFilter.Outcome;
import nz.org.riskscape.engine.restriction.ExpressionRestriction;
import nz.org.riskscape.engine.restriction.Restriction;
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.StructBuilder;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.engine.types.Text;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;

/**
 * Provides a stream of {@link Tuple}s based on a {@link SimpleFeatureSource}
 */
@Slf4j
public class FeatureSourceRelation extends BaseRelation {

  /**
   * Name of tuple member that identifies the value of {@link SimpleFeature#getID()}
   */
  public static final String FID_MEMBER_NAME = "_fid";

  /**
   * Special instance of Text type used to identify feature id - makes it easier to treat it specially
   */
  public static final Type FID_TYPE = new Text();

  @NonNull @Getter
  private final SRIDSet sridSet;

  @NonNull @Getter
  private final SimpleFeatureSource featureSource;

  @NonNull @Getter
  private final Filter filter;

  protected final CoordinateReferenceSystem crs;

  private Optional<Long> size = null;

  /**
   * Infer a struct from the given FeatureSource
   * TODO should we be passing in the CRS here?  We have the potential to scrape a different CRS for our referenced
   * value compared to the one that gets passed in to the spatial metadata
   * @param fs a feature source to build a {@link Struct} from
   * @param includeFid if true, a special `_fid` member is added to the type which holds the value of
   * {@link SimpleFeature#getID()}.
   */
  public static Struct fromFeatureSource(SimpleFeatureSource fs, boolean includeFid) {
    SimpleFeatureType type = fs.getSchema();

    ReferencedEnvelope bounds = null;
    try {
      bounds = fs.getBounds();
      if (bounds == null) {
        // FeatureSource#getBounds says that it may return null if calculating the bounds it judged
        // too costly.
        log.warn("Feature source returned a null bounds. This is because it would be too costly to "
            + "calculate the actual bounds");
      }
    } catch (IOException e) {
      // Likewise if getting the bounds results in IOException we log and carry on.
      log.warn("Bounds could not be calculated", e);
    }
    return fromFeatureType(type, bounds, includeFid);
  }

    /**
   * Infer a struct from the given SimpleFeatureType
   * @param type a feature type to build a {@link Struct} from
   * @param bounds the geometric bounds if known (else null)
   * @param includeFid if true, a special `_fid` member is added to the type which holds the value of
   * {@link SimpleFeature#getID()}.
   */
  public static Struct fromFeatureType(SimpleFeatureType type, ReferencedEnvelope bounds, boolean includeFid) {
    StructBuilder structBuilder = Struct.builder();

    if (includeFid) {
      structBuilder.add(FID_MEMBER_NAME, FID_TYPE);
    }

    GeometryDescriptor geomDescriptor = type.getGeometryDescriptor();
    for (AttributeDescriptor attributeDescriptor : type.getAttributeDescriptors()) {
      String key = attributeDescriptor.getLocalName();
      Class binding = attributeDescriptor.getType().getBinding();
      if (binding == FeatureTypeStyle.class) {
        // Ignore styling
        continue;
      }
      Type riskscapeType = Types.fromJavaType(binding);
      if (attributeDescriptor.equals(geomDescriptor)) {
        // feature sources can use the most accurate geometry type as the feature source is responsible
        // for ensuring all geometries are of the same type. So we use Geom.fromJavaType()
        Type geometryType = Geom.fromJavaType(attributeDescriptor.getType().getBinding().asSubclass(Geometry.class));

        if (bounds != null && bounds.getCoordinateReferenceSystem() != null) {
          riskscapeType = Referenced.ofNullable(geometryType, bounds);
        }

      }
      structBuilder.add(key, riskscapeType);
    }
    return structBuilder.build();
  }

  /**
   * @return the StructMember in the given struct the corresponds to the feature's geometry
   */
  public static StructMember geometryStructMember(Struct struct, SimpleFeatureType schema) {
    GeometryDescriptor descriptor = schema.getGeometryDescriptor();
    return descriptor == null ? null : struct.getEntry(descriptor.getLocalName());
  }

  public static SpatialMetadata spatialMetadataFromSchema(
      Struct struct,
      SimpleFeatureType schema,
      CoordinateReferenceSystem crs
  ) {
    StructMember geometryMember = geometryStructMember(struct, schema);

    if (geometryMember != null) {
      // NB we used to fetch the crs from the featuresource, but that can't parse the wide range of prj formats that
      // the world has thrown at us, so now we are given a crs to use
      return new SpatialMetadata(crs, geometryMember);
    } else {
      return null;
    }
  }

  public FeatureSourceRelation(SimpleFeatureSource fs, SRIDSet sridSet, CoordinateReferenceSystem crs) {
    this(fromFeatureSource(fs, false), fs, sridSet, Filter.INCLUDE, crs);
  }

  public FeatureSourceRelation(@NonNull Struct type,
      @NonNull SimpleFeatureSource fs, @NonNull SRIDSet sridSet, CoordinateReferenceSystem crs) {
    this(type, fs, sridSet, Filter.INCLUDE, crs);
  }

  public FeatureSourceRelation(
      @NonNull Struct type,
      @NonNull SimpleFeatureSource fs,
      @NonNull SRIDSet sridSet,
      @NonNull Filter filter,
      @NonNull CoordinateReferenceSystem crs) {

    super(type, null, spatialMetadataFromSchema(type, fs.getSchema(), crs));
    this.sridSet = sridSet;
    this.featureSource = fs;
    this.filter = filter;
    this.crs = crs;

    getSpatialMetadata().ifPresent(sm -> sridSet.get(sm.getCrs()));
  }

  protected FeatureSourceRelation(
      @NonNull BaseRelation.Fields clonedFields,
      @NonNull SimpleFeatureSource fs, @NonNull SRIDSet sridSet, @NonNull Filter filter,
      @NonNull CoordinateReferenceSystem crs) {
    super(clonedFields);
    this.sridSet = sridSet;
    this.featureSource = fs;
    this.filter = filter;
    this.crs = crs;

    getSpatialMetadata().ifPresent(sm -> sridSet.get(sm.getCrs()));
  }

  @Override
  public FeatureSourceTupleIterator rawIterator() {
    Query query = new Query();
    query.setFilter(filter);

    return FeatureSourceBookmarkResolver.LOCK_DEFEATER.call("build-fs-from-query",
        () -> FeatureSourceTupleIterator.fromQuery(featureSource, query, sridSet, getRawType(), crs));
  }

  @Override
  public String getSourceInformation() {
    return featureSource.getDataStore() == null ? "unknown" : featureSource.getDataStore().toString();
  }

  @Override
  protected BaseRelation cloneWithRestrictionIfSupported(Restriction restriction) {
    // only works if we can get an expression to convert
    if (!(restriction instanceof ExpressionRestriction)) {
      return null;
    }
    ExpressionRestriction exprRestriction = (ExpressionRestriction) restriction;
    Filter newFilter;
    try {
      newFilter = ECQL.toFilter(exprRestriction.getExpression().toSource());
    } catch (CQLException e) {
      //We cannot clone if the expression doesn't parse as ECQL - simple expressions should be fine as riskscape
      // expressions are semi descended from ECQL
      return null;
    }

    // so begins the process of determining whether we can build a feature source compatible feature
    // from the one we've been given
    // 1 - we build a list of all the transformations from all the projections in our transformation list in
    // the order that they've been applied, so that we end up with:
    //   T0      a->a'        b->b'   c->c'
    //   T1      a'->a''      b->b''  c->c''
    // where the un-prime attributes are from the raw type, and the prime-prime attributes are in the projected type
    List<Map<StructMember, StructMember>> transformations = new ArrayList<>(transformers.length);

    for (Transformer transformer : this.transformers) {
      if (transformer.projectionFunction != null) {
        Map<List<StructMember>,List<StructMember>> targetToSource = transformer.projectionFunction.getDirectMapping();
        Map<StructMember, StructMember> sourceToTarget = new HashMap<>();

        for (Map.Entry<List<StructMember>, List<StructMember>> entry : targetToSource.entrySet()) {
          if (entry.getValue().size() == 1 && entry.getKey().size() == 1) {
            // crucially here, we invert the map, so instead of getting target to source, we got source to target
            // this means if the map is not bi-directional, that we arbitrarily select the last mapped value to be
            // the one we use - it doesn't really matter for our purposes which route back we select (TODO add a test
            // case to prove it?)
            sourceToTarget.put(entry.getValue().get(0), entry.getKey().get(0));
          }
        }
        transformations.add(sourceToTarget);
      }
    }

    // Now we build a working copy map of mappings, that is used to track an attributes name back from the relation's
    // projected type, aka getType, back to the the source type, aka getRawType()
    // To begin with, the map looks like:
    // T    a->a     b->b      c->c
    Map<StructMember, StructMember> mappedAttributes = this.getRawType().getMembers()
        .stream()
        .collect(Collectors.toMap(Function.identity(), Function.identity()));


    // With each iteration of the loop, we adjust the value set to point to the projected members of the current
    // projected type.  Values without a mapping are set to null
    for (Map.Entry<StructMember, StructMember> mapped : mappedAttributes.entrySet()) {
      for (Map<StructMember, StructMember> sourceToTarget: transformations) {
        StructMember targetMember = sourceToTarget.get(mapped.getValue());
        if (targetMember != null) {
          mapped.setValue(targetMember);
        } else {
          mapped.setValue(null);
          break;
        }
      }
    }

    // using our attribute map, we attempt to rebuild the filter
    RenameFilter converter = new RenameFilter(mappedAttributes);
    RenameFilter.Result renamed = converter.convert(newFilter);

    // must have converted unmapped properties - can't rebuild the filter, give up
    if (renamed.getOutcome() == Outcome.NONE) {
      return null;
    } else {
      log.debug(
          "Filter {} converted to {} for push up in to feature source {}",
          newFilter,
          renamed.getFilter(),
          getSourceInformation());
    }

    Filter combinedFilter = newFilter == Filter.INCLUDE ?  renamed.getFilter()
        : CommonFactoryFinder.getFilterFactory().and(renamed.getFilter(), this.filter);
    BaseRelation converted = new FeatureSourceRelation(
        new BaseRelation.Fields(), featureSource, sridSet, combinedFilter, crs);
    if (renamed.getOutcome() == Outcome.PARTIAL) {
      //If renamed has non error problems we assume it cannot apply the whole filter, in which case we clone with
      //original restriction so all of the filtering is done.
      converted = converted.cloneWithTransformer(new Transformer(getType(), restriction));
    }
    return converted;
  }

  @Override
  protected BaseRelation clone(Fields fields) {
    return new FeatureSourceRelation(fields, featureSource, sridSet, filter, crs);
  }

  @Override
  public Optional<Long> size() {
    if (size == null) {
      try {
        Query query = new Query();
        query.setFilter(filter);
        long featureCount = featureSource.getCount(query);
        if (featureCount >= 0) {
          size = Optional.of(featureCount);
        } else {
          // FeatureSource.getCount() can return -1 if the operation is "too costly"
          size = Optional.empty();
        }
      } catch (IOException e) {
        // something went wrong... as determining the size is optional, just fall back
        // to returning "unknown". If the data source is bad, we'll handle that error
        // when we try to read the tuples properly
        size = Optional.empty();
      }
    }
    return size;
  }
}
