/*
 * 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.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import com.google.common.collect.Lists;

import lombok.Getter;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.projection.FlatProjection;
import nz.org.riskscape.engine.projection.FlatProjector;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.RelationType;
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.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ResultOrProblems;

public class UnnestProjection implements FlatProjection {

  private class FlatProjectorImpl implements FlatProjector {

    private class UnnestIterator implements TupleIterator {

      private final Iterator<?>[] values;
      private long index = 0;
      private Tuple yielding;

      UnnestIterator(Tuple tuple) {
        values = new Iterator[sourceMembers.length];

        for (int idx = 0; idx < sourceMembers.length; idx++) {
          Object toUnnest = tuple.fetch(sourceMembers[idx]);

          if (toUnnest instanceof Relation relation) {
            values[idx] = relation.iterator();
          } else {
            Collection<?> coll = (Collection<?>) tuple.fetch(sourceMembers[idx]);
            if (coll == null) {
              values[idx] = Collections.emptyIterator();
            } else {
              values[idx] = coll.iterator();
            }
          }
          if (emitEmpty) {
            // if emit empty we wrap the iterator in one that will always return at least one value
            Iterator backingIterator = values[idx];
            values[idx] = new Iterator() {
              boolean first = true;
              @Override
              public boolean hasNext() {
                return first || backingIterator.hasNext();
              }

              @Override
              public Object next() {
                first = false;
                return backingIterator.hasNext() ? backingIterator.next() : null;
              }
            };
          }
        }

        this.yielding = new Tuple(producedType);
        yielding.setAll(tuple);
      }

      @Override
      public boolean hasNext() {
        // TODO close any relation iterators
        for (int idx = 0; idx < values.length; idx++) {
          if (values[idx].hasNext()) {
            return true;
          }
        }
        return false;
      }

      @Override
      public Tuple next() {
        ++index;
        for (int idx = 0; idx < projectedMembers.length; idx++) {
          yielding.set(projectedMembers[idx], values[idx].hasNext() ? values[idx].next() : null);
        }

        if (indexMember != null) {
          yielding.set(indexMember, index);
        }
        return yielding.clone();
      }
    }

    @Getter
    private final Struct sourceType;
    @Getter
    private final Struct producedType;


    @Getter
    private final StructMember indexMember;
    private final StructMember[] sourceMembers;
    private final StructMember[] projectedMembers;

    FlatProjectorImpl(
        Struct sourceType,
        Struct projectedType,
        List<String> members,
        Optional<String> indexKey
    ) {
      this.sourceType = sourceType;
      this.producedType = projectedType;

      this.sourceMembers = sourceType.getMembers()
          .stream()
          .filter(m -> members.contains(m.getKey()))
          .collect(Collectors.toList())
          .toArray(new StructMember[0]);

      this.projectedMembers = projectedType.getMembers()
          .stream()
          .filter(m -> members.contains(m.getKey()))
          .collect(Collectors.toList())
          .toArray(new StructMember[0]);

      this.indexMember = indexKey.map(key -> projectedType.getMember(key).get()).orElse(null);
    }

    @Override
    public TupleIterator apply(Tuple tuple) {
      return new UnnestIterator(tuple);
    }
  }

  private final List<String> members;
  private final Optional<String> indexKey;
  private final boolean emitEmpty;

  public UnnestProjection(List<String> members, Optional<String> indexKey, boolean emitEmpty) {
    this.members = members;
    this.indexKey = indexKey;
    this.emitEmpty = emitEmpty;
  }

  @Override
  public ResultOrProblems<FlatProjector> getFlatProjector(Struct sourceType) {
    List<Problem> problems = new ArrayList<Problem>();
    StructBuilder builder = new StructBuilder();
    // if there's more than one list, we can't guarantee that the lists are of equal size, and so it's possible we'll
    // get null members in the output
    List<String> membersWip = Lists.newArrayList(this.members);
    boolean nullableMembers = this.members.size() > 1 || emitEmpty;
    for (StructMember sourceMember : sourceType.getMembers()) {
      if (membersWip.remove(sourceMember.getKey())) {
        Optional<RSList> listOpt = sourceMember.getType().findAllowNull(RSList.class);
        RelationType relationType = sourceMember.getType().findAllowNull(RelationType.class).orElse(null);

        if (listOpt.isPresent()) {
          Type listType = listOpt.get().getMemberType();
          if (nullableMembers) {
            listType = Nullable.of(listType);
          }
          builder.add(sourceMember.getKey(), listType);
        } else {
          if (relationType == null) {
            problems.add(Problem.error(
                "Member '%s' is not a list type, was '%s'",
                sourceMember.getKey(),
                sourceMember.getType()
            ));
          } else {
            builder.add(sourceMember.getKey(), Nullable.ifTrue(emitEmpty, relationType.getMemberType()));
          }
        }
      } else {
        builder.add(sourceMember.getKey(), sourceMember.getType());
      }
    }

    if (indexKey.isPresent()) {
      builder.add(indexKey.get(), Types.INTEGER);
    }

    if (!membersWip.isEmpty()) {
      problems.add(Problem.error(
          "Could not find members to unnest %s among %s",
          membersWip,
          sourceType.getMemberKeys()
      ));
    }

    if (Problem.hasErrors(problems)) {
      return ResultOrProblems.failed(problems);
    }

    return builder
      .buildOr()
      .composeProblems("Cannot create index-key '%s', member already present in %s", indexKey.orElse(""), sourceType)
      .map(projectedType  -> {
        return new FlatProjectorImpl(sourceType, projectedType, this.members, this.indexKey);
      });
  }
}
