/*
 * 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 static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import org.geotools.api.data.Query;
import org.geotools.data.simple.SimpleFeatureCollection;
import org.geotools.data.simple.SimpleFeatureIterator;
import org.geotools.api.data.SimpleFeatureSource;
import org.geotools.feature.collection.SimpleFeatureIteratorImpl;
import org.geotools.filter.text.ecql.ECQL;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.crs.DefaultEngineeringCRS;
import org.hamcrest.collection.IsIterableContainingInOrder;
import org.junit.Before;
import org.junit.Test;
import org.locationtech.jts.geom.Geometry;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
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.AttributeType;
import org.geotools.api.feature.type.GeometryDescriptor;
import org.geotools.api.feature.type.GeometryType;
import org.geotools.api.filter.Filter;
import org.geotools.api.filter.IncludeFilter;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;

import nz.org.riskscape.engine.CrsHelper;
import nz.org.riskscape.engine.ProjectTest;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.filter.FilterFactory;
import nz.org.riskscape.engine.function.StringFunctions;
import nz.org.riskscape.engine.projection.TypeProjection;
import nz.org.riskscape.engine.restriction.ExpressionRestriction;
import nz.org.riskscape.engine.rl.DefaultOperators;
import nz.org.riskscape.engine.types.Enumeration;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.rl.ExpressionParser;

public class FeatureSourceRelationTest extends ProjectTest implements CrsHelper {

  FilterFactory ff;

  SRIDSet sridSet;
  @Mock
  SimpleFeatureSource fs;
  @Mock
  SimpleFeatureType type;
  @Mock
  SimpleFeatureCollection collection;

  SimpleFeatureIterator iterator;

  List<AttributeDescriptor> descriptors = new ArrayList<>();

  ArgumentCaptor<org.geotools.api.data.Query> queryCaptor;

  private Collection<SimpleFeature> features = new ArrayList<>();

  CoordinateReferenceSystem crs = DefaultEngineeringCRS.GENERIC_2D;

  FeatureSourceRelation relation;

  public FeatureSourceRelationTest() {
    project.getFunctionSet().insertFirst(DefaultOperators.INSTANCE);
    project.getFunctionSet().addAll(StringFunctions.FUNCTIONS);

    ff = engine.getFilterFactory();
    sridSet = project.getSridSet();
  }

  public FeatureSourceRelation getRelation() {
    if (relation == null) {
      relation = new FeatureSourceRelation(fs, sridSet, crs);
    }

    return relation;
  }

  @Before
  public void setupMocks() throws IOException {
    MockitoAnnotations.initMocks(this);
    queryCaptor = ArgumentCaptor.forClass(org.geotools.api.data.Query.class);
    when(fs.getSchema()).thenReturn(type);
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);
    when(fs.getFeatures(queryCaptor.capture())).thenReturn(collection);
    when(type.getAttributeDescriptors()).thenReturn(descriptors);
    when(type.getCoordinateReferenceSystem()).thenReturn(nzMapGrid());

    // replicate underlying SimpleFeatureSource behaviour - counting all features
    // returns the size, but more specific queries return -1 as it's too costly
    // when(fs.getCount(argThat(new IsQueryAll()))).thenReturn(collection.size());
    when(fs.getCount(any())).thenAnswer(i -> {
      boolean isQueryAll = ((Query) i.getArgument(0)).getFilter() instanceof IncludeFilter;
      return isQueryAll ? features.size() : -1;
    });
  }

  @Test
  public void canCreateAnEmptyRelation() throws IOException {
    List<Tuple> elements = getRelation().stream().collect(Collectors.toList());

    assertTrue(relation.getType().getMembers().isEmpty());
    assertTrue(elements.isEmpty());
  }

  @Test
  public void canInferTheTypeFromTheFeatureSourceSchema() {
    descriptors.add(descriptor("id", Long.class));
    descriptors.add(descriptor("desc", String.class));
    descriptors.add(descriptor("the_geom", Geometry.class));

    Struct inferred = getRelation().getType();
    Struct expected = Struct.of("id", Types.INTEGER).and("desc", Types.TEXT)
        .and("the_geom", Types.GEOMETRY);

    assertTrue(inferred.isEquivalent(expected));
  }

  @SuppressWarnings("unchecked")
  @Test
  public void canBuildSpatialMetadataIgnoringTheFeatureSourceSchema() throws Exception {
    GeometryDescriptor geometryDescriptor = mock(GeometryDescriptor.class);

    @SuppressWarnings("rawtypes")
    Class clazz = Geometry.class;
    GeometryType attributeType = mock(GeometryType.class);
    when(attributeType.getBinding()).thenReturn(clazz);
    when(geometryDescriptor.getType()).thenReturn(attributeType);

    when(geometryDescriptor.getLocalName()).thenReturn("the_geom");
    when(geometryDescriptor.getType()).thenReturn(attributeType);
    when(geometryDescriptor.getCoordinateReferenceSystem()).thenReturn(nzMapGrid());
    descriptors.add(geometryDescriptor);
    when(type.getGeometryDescriptor()).thenReturn(geometryDescriptor);

    ReferencedEnvelope bounds = new ReferencedEnvelope(nzMapGrid());
    bounds.expandToInclude(20, 20); // make sure we don't inadvertently match against a mistakenly empty envelope
    when(fs.getBounds()).thenReturn(bounds);

    Struct inferred = getRelation().getType();

    Referenced geometryType = inferred.getEntry("the_geom").getType().find(Referenced.class).orElse(null);
    assertEquals(nzMapGrid(), geometryType.getCrs());
    assertEquals(bounds, geometryType.getBounds());

    assertEquals(Types.GEOMETRY, geometryType.getUnwrappedType());
    SpatialMetadata spatialMetadata = getRelation().getSpatialMetadata().get();
    assertNotEquals(nzMapGrid(), spatialMetadata.getCrs());
    assertSame(inferred.getMembers().get(0), spatialMetadata.getGeometryStructMember());
  }

  @Test
  public void willReturnASetOfTuplesConvertedFromFeatures() {
    descriptors.add(descriptor("desc", String.class));
    descriptors.add(descriptor("value", Integer.class));

    features.add(feature(ImmutableMap.of("desc", "foo", "value", 10.0D)));
    features.add(feature(ImmutableMap.of("desc", "bar", "value", 12)));
    features.add(feature(ImmutableMap.of("desc", "baz", "value", 14.0D)));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct struct = getRelation().getType();
    assertThat(struct, is(Struct.of("desc", Types.TEXT, "value", Types.INTEGER)));
    List<Tuple> expected = Lists.newArrayList(
      Tuple.ofValues(struct, "foo", 10L),
      Tuple.ofValues(struct, "bar", 12L),
      Tuple.ofValues(struct, "baz", 14L)
    );

    List<Tuple> elements = getRelation().stream().collect(Collectors.toList());
    assertEquals(expected, elements);
    assertEquals(Filter.INCLUDE, queryCaptor.getValue().getFilter());
  }

  @Test
  public void willNotSetANullValue() {
    descriptors.add(descriptor("desc", String.class));

    features.add(feature(ImmutableMap.of()));
    features.add(feature(ImmutableMap.of("desc", "bar")));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct struct = getRelation().getType();
    List<Tuple> expected = Lists.newArrayList(
      new Tuple(struct),
      new Tuple(struct).set("desc",  "bar")
    );

    List<Tuple> elements = getRelation().stream().collect(Collectors.toList());
    assertEquals(expected, elements);
    assertEquals(Filter.INCLUDE, queryCaptor.getValue().getFilter());
  }

  @Test
  public void canApplyAFilterToATypeProjectedRelation() throws Exception {
    descriptors.add(descriptor("desc", String.class));

    features.add(feature(ImmutableMap.of("desc", "foo")));
    features.add(feature(ImmutableMap.of("desc", "barr")));
    features.add(feature(ImmutableMap.of("desc", "bazzz")));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct projectedType = Struct.of("desc-mapped", Types.TEXT, "desc-mangled", Types.INTEGER);

    FeatureSourceRelation fsRelation = (FeatureSourceRelation) getRelation()
        .project(new TypeProjection(projectedType, ImmutableMap.of(
            "desc-mapped", expressionParser.parse("desc"),
            "desc-mangled", expressionParser.parse("str_length(desc)")
          ), realizationContext))
        .get();

    List<Tuple> expected = Lists.newArrayList(
      Tuple.ofValues(projectedType, "foo", 3L),
      Tuple.ofValues(projectedType, "barr", 4L),
      Tuple.ofValues(projectedType, "bazzz", 5L));

    List<Tuple> elements = fsRelation.stream().collect(Collectors.toList());
    assertEquals(expected, elements);
    assertEquals(Filter.INCLUDE, queryCaptor.getValue().getFilter());
    assertThat(fsRelation.size().get(), is(3L));

    FeatureSourceRelation restricted = (FeatureSourceRelation) fsRelation
        .restrict(restriction("\"desc-mapped\" = 'foo'")).get();
    assertEquals("desc = 'foo' AND INCLUDE", ECQL.toCQL(restricted.getFilter()));
    // don't report a misleading size when a filter is applied
    assertThat(restricted.size(), is(Optional.empty()));
  }

  @Test
  public void canApplyAFilterToAMangledAndProjectedRelationFallingBackToBruteForceFiltering() throws Exception {
    descriptors.add(descriptor("desc", String.class));

    features.add(feature(ImmutableMap.of("desc", "foo")));
    features.add(feature(ImmutableMap.of("desc", "barr")));
    features.add(feature(ImmutableMap.of("desc", "bazzz")));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct projectedType = Struct.of("desc-mapped", Types.TEXT, "desc-mangled", Types.INTEGER);

    FeatureSourceRelation fsRelation = (FeatureSourceRelation) getRelation()
        .project(new TypeProjection(projectedType, ImmutableMap.of(
            "desc-mapped", expressionParser.parse("desc"),
            "desc-mangled", expressionParser.parse("str_length(desc)")
          ), realizationContext))
        .get();

    FeatureSourceRelation restricted = (FeatureSourceRelation) fsRelation
        .restrict(restriction("\"desc-mangled\" = 4")).get();

    List<Tuple> expected = Lists.newArrayList(
        Tuple.ofValues(projectedType, "barr", 4L));

    assertEquals(Filter.INCLUDE, restricted.getFilter());
    List<Tuple> elements = restricted.stream().collect(Collectors.toList());
    assertEquals(expected, elements);

  }

  @Test
  public void canApplyAFilterToEnumAndTextTypes() throws Exception {
    //Checking that parts ifBounded the restriction are not lost. Enumerations currently cannot be upstreamed.
    descriptors.add(descriptor("desc", String.class));
    descriptors.add(descriptor("type", Integer.class));
    features.add(feature(ImmutableMap.of("desc", "house1", "type", 1)));
    features.add(feature(ImmutableMap.of("desc", "house2", "type", 1)));
    features.add(feature(ImmutableMap.of("desc", "shop", "type", 2)));
    features.add(feature(ImmutableMap.of("desc", "office", "type", 3)));
    features.add(feature(ImmutableMap.of("desc", "factory", "type", 3)));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct projectedType = Struct.of("desc", Types.TEXT, "type", Enumeration.oneBased("timber", "brick", "steel"));

    FeatureSourceRelation fsRelation = (FeatureSourceRelation) getRelation()
        .project(new TypeProjection(projectedType, ImmutableMap.of(
            "desc", expressionParser.parse("desc"),
            "type", expressionParser.parse("type")
          ), realizationContext))
        .get();

    FeatureSourceRelation restricted = (FeatureSourceRelation) fsRelation
        .restrict(restriction("\"type\" = 'steel' and desc = 'office'")).get();
    assertThat(restricted.stream().collect(Collectors.toList()), IsIterableContainingInOrder.contains(
        Tuple.of(projectedType, "desc", "office", "type", "steel")));
  }

  @Test
  public void canApplyFilterToSubtreesThatAreSafeToDoSo() throws Exception {
    descriptors.add(descriptor("desc", String.class));

    features.add(feature(ImmutableMap.of("desc", "foo")));
    features.add(feature(ImmutableMap.of("desc", "barr")));
    features.add(feature(ImmutableMap.of("desc", "bazzz")));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct projectedType = Struct.of("desc-mapped", Types.TEXT, "desc-mangled", Types.INTEGER);

    FeatureSourceRelation fsRelation = (FeatureSourceRelation) getRelation()
        .project(new TypeProjection(projectedType, ImmutableMap.of(
            "desc-mapped", expressionParser.parse("desc"),
            "desc-mangled", expressionParser.parse("str_length(desc)")
          ), realizationContext))
        .get();

    FeatureSourceRelation restricted = (FeatureSourceRelation) fsRelation
        .restrict(restriction("\"desc-mapped\" = 'foo' and \"desc-mangled\" = 3")).get();

    //Can lift subtrees when and'd
    assertEquals("desc = 'foo' AND INCLUDE", ECQL.toCQL(restricted.getFilter()));
  }

  @Test
  public void canApplyFilterToSubtreesThatAreSafeToDoSo_WhichMayBeNone() throws Exception {
    descriptors.add(descriptor("desc", String.class));

    features.add(feature(ImmutableMap.of("desc", "foo")));
    features.add(feature(ImmutableMap.of("desc", "barr")));
    features.add(feature(ImmutableMap.of("desc", "bazzz")));
    iterator = new SimpleFeatureIteratorImpl(features);
    when(collection.features()).thenReturn(iterator);

    Struct projectedType = Struct.of("desc-mapped", Types.TEXT, "desc-mangled", Types.INTEGER);

    FeatureSourceRelation fsRelation = (FeatureSourceRelation) getRelation()
        .project(new TypeProjection(projectedType, ImmutableMap.of(
            "desc-mapped", expressionParser.parse("desc"),
            "desc-mangled", expressionParser.parse("str_length(desc)")
          ), realizationContext))
        .get();

    FeatureSourceRelation restricted = (FeatureSourceRelation) fsRelation
        .restrict(restriction("\"desc-mapped\" = 'foo' or \"desc-mangled\" = 3")).get();

    //Can't restrict when or'd
    assertEquals(Filter.INCLUDE, restricted.getFilter());
  }


  private SimpleFeature feature(ImmutableMap<String, ? extends Object> values) {
    SimpleFeature feature = mock(SimpleFeature.class);
    values.entrySet().forEach(entry -> {
      when(feature.getAttribute(entry.getKey())).thenReturn(entry.getValue());
    });
    when(feature.getFeatureType()).thenReturn(mock(SimpleFeatureType.class));
    return feature;
  }

  @SuppressWarnings({ "unchecked", "rawtypes" })
  private AttributeDescriptor descriptor(String localName, Class class1) {
    AttributeDescriptor descriptor = mock(AttributeDescriptor.class);
    when(descriptor.getLocalName()).thenReturn(localName);
    AttributeType attributeType = mock(AttributeType.class);
    when(attributeType.getBinding()).thenReturn(class1);
    when(descriptor.getType()).thenReturn(attributeType);

    return descriptor;
  }

  private ExpressionRestriction restriction(String expression) {
    return new ExpressionRestriction(ExpressionParser.parseString(expression), expressionRealizer);
  }

}
