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

import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;

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

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.pipeline.Realized;
import nz.org.riskscape.engine.projection.FlatProjector;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.types.CoercionException;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.engine.types.Text;
import nz.org.riskscape.engine.types.Type;

/**
 * Step that will transform flat input data into a more nested form.
 *
 * Input data would normally be expected from a source such as CSV which is inherently all text so
 * coercion is used to attempt to convert. Of course this can cause errors at runtime when fields do
 * not contain the expected values.
 */
public class SegmentStep extends BaseStep<SegmentStep.Parameters> {

  public static class Parameters extends ProjectionStep.ProjectionParameters {

    /**
     * The {@link Type} that input relation should be segmented to. This type should be a
     * {@link Struct} with members that are {@link Struct}s containing only simple types. Eg no
     * deeper nesting.
     */
    @ParameterField
    Type segmented;

    /**
     * When true and coercion errors are silently ignored and the relation item that generated them
     * is skipped.
     */
    @ParameterField
    boolean skipInvalid;

    /**
     * Set this to true to allow rows in the source data to having missing attributes
     */
    @ParameterField
    private boolean allowMissingAttributes;
  }

  @RequiredArgsConstructor
  private static class SegmentFlatProjector implements FlatProjector {

    @Getter
    private final Struct sourceType;

    @Getter
    private final Struct producedType;

    /**
     * Function that given a {@link Tuple} of the sourceType will return a {@link Tuple} of the
     * projected type.
     */
    private final Function<Tuple, Tuple> tupleMapper;

    private final boolean skipInvalid;

    @Override
    public TupleIterator apply(Tuple source) {
      try {
        Tuple result = tupleMapper.apply(source);
        return TupleIterator.singleton(result);
      } catch (Exception e) {
        if (skipInvalid) {
          return TupleIterator.EMPTY;
        }
        throw e;
      }
    }

  }

  public SegmentStep(Engine engine) {
    super(engine);
  }

  @Override
  public ResultOrProblems<? extends Realized> realize(Parameters parameters) {
    Struct sourceType = parameters.input.getProduces();

    Set<String> seen = Sets.newHashSet();
    List<Problem> problems = Lists.newArrayList();

    //a list of functions that can produce each segment of produced type
    List<Function<Tuple, Tuple>> segmentMappers = Lists.newArrayList();
    Struct projected = parameters.segmented.asStruct();
    for (StructMember segment : projected.getMembers()) {
      boolean targetNullable = Nullable.is(segment.getType());
      Struct segmentType = segment.getType().asStruct();

      //A list of functions that given the source tuple will return the value to set each attribute
      //in the produced segment. The returned value would normally come from the source tuple.
      List<Function<Tuple, Object>> attributeMappers = Lists.newArrayList();
      for (StructMember member : segmentType.getMembers()) {
        boolean nullable = Nullable.is(member.getType());

        if (sourceType.hasMember(member.getKey())) {
          if (sourceType.getEntry(member.getKey()).getType().getUnwrappedType() instanceof Struct) {
            problems.add(Problem.error("Source attribute '%s' is a struct", member.getKey()));
          }
        } else if (! (nullable || parameters.allowMissingAttributes)) {
          problems.add(Problem.error("Segment '%s' requires attribute '%s' which does not exist in input %s",
              segment.getKey(), member.getKey(), sourceType));
        }
        if (!seen.add(member.getKey())) {
          problems.add(Problem.error("Segment '%s' contains attribute '%s' which is used in other segments",
              segment.getKey(), member.getKey()));
        }
        if (member.getType().getUnwrappedType() instanceof Struct) {
          problems.add(Problem.error("Segment '%s' contains nested structs in '%s'",
              segment.getKey(), member.getKey()));
        }
        attributeMappers.add(attributeMapper(sourceType, member));
      }
      segmentMappers.add(segmentBuilder(segmentType, attributeMappers, targetNullable));
    }

    if (!problems.isEmpty()) {
      return ResultOrProblems.failed(problems);
    }
    return ResultOrProblems.of(new SegmentFlatProjector(sourceType, projected,
        resultBuilder(projected, segmentMappers), parameters.skipInvalid));
  }

  private Function<Tuple, Object> attributeMapper(Struct sourceType, StructMember targetMember) {
    Type targetType = targetMember.getType();
    Optional<StructMember> sourceMember = sourceType.getMember(targetMember.getKey());
    if (!sourceMember.isPresent()) {
      //If the sourceType does not have an attribute of the same key, we will not find it. So return null.
      return (t) -> null;
    }
    boolean targetNullable = Nullable.is(targetType);
    boolean targetText = targetType instanceof Text;
    return (t) -> {
      Object value = t.fetch(sourceMember.get());
      if ("".equals(value) && targetNullable && !targetText) {
        return null;
      }
      try {
        return targetType.coerce(value);
      } catch (CoercionException e) {
        throw new CoercionException(e.getValue(), e.getType(),
            "Value '%s' from source attribute %s could not be converted to type %s - %s",
            value, sourceMember.get().getKey(), targetType, e.getMessage()
        );
      }
    };
  }

  /**
   * Get a {@link Function} that will produce a {@link Tuple} of segmentType give an input
   * {@link Tuple}.
   *
   * @param segmentType type of the {@link Tuple} to be returned
   * @param attributeMappers functions to obtain value to set to each attribute of segmentType
   * @return function to return segment tuple
   */
  private Function<Tuple, Tuple> segmentBuilder(Struct segmentType, List<Function<Tuple, Object>> attributeMappers,
      boolean segmentNullable) {
    return (source) -> {
      try {
        Object[] values = new Object[attributeMappers.size()];
        for (int i = 0; i < attributeMappers.size(); i++) {
          values[i] = attributeMappers.get(i).apply(source);
        }
        return Tuple.ofValues(segmentType, values);
      } catch (CoercionException e) {
        if (segmentNullable) {
          return null;
        }
        throw e;
      }
    };
  }

  /**
   * Get a {@link Function} that will produce a {@link Tuple} of resultType give an input
   * {@link Tuple}.
   *
   * @param resultType type of the {@link Tuple} to be returned
   * @param segmentMappers functions to obtain value to set to each attribute of segmentType
   * @return function to return result tuple
   */
  private Function<Tuple, Tuple> resultBuilder(Struct resultType, List<Function<Tuple, Tuple>> segmentMappers) {
    return (source) -> {
      Object[] values = new Object[segmentMappers.size()];
      for (int i = 0; i < segmentMappers.size(); i++) {
        values[i] = segmentMappers.get(i).apply(source);
      }
      return Tuple.ofValues(resultType, values);
    };
  }

}
