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

import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Collectors;

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

import hdf.hdf5lib.exceptions.HDF5Exception;
import hdf.hdf5lib.exceptions.HDF5SymbolTableException;
import lombok.Getter;

import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.BindingContext;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.data.relation.RelationBookmarkParams;
import nz.org.riskscape.engine.data.relation.RelationBookmarkResolver;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.relation.BaseRelation;
import nz.org.riskscape.engine.relation.Relation;
import nz.org.riskscape.engine.relation.ZipRelation;
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.hdf5.cursor.H5DatasetCursor;
import nz.org.riskscape.hdf5.relation.H5Relation;
import nz.org.riskscape.hdf5.types.H5CompoundType;
import nz.org.riskscape.hdf5.types.H5Type;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

/**
 * A general purpose resolver to open HDF5 files (or rather datasets within them) in a variety of ways.
 */
public class Hdf5Resolver extends RelationBookmarkResolver<Hdf5Resolver.Parameters> {

  public static final Range<Integer> ALLOWED_READ_SIZE_MB = Range.closed(1, 2048);

  interface LocalProblems extends ProblemFactory {

    /**
     * A file couldn't be loaded. This could be caused by any number of reasons from:
     * - file isn't an HDF5 file
     * - file can't be read be H5 library (maybe file is compressed and lib doesn't support it).
     */
    Problem hdf5FileLoadError(String cause);

    Problem hdf5DatasetDoesNotExist(String dataset);

    /**
     * When a bookmark combines HDF5 datasets, they must all contain the same number of elements
     */
    Problem differentDatasetSizes();
  }

  static final LocalProblems PROBLEMS = Problems.get(LocalProblems.class);

  public enum DatasetMode {
    RELATION(Relation.class);

    private final Class<?> dataType;

    DatasetMode(Class<?> dataType) {
      this.dataType = dataType;
    }

    public Class<?> getDataType() {
      return dataType;
    }
  }

  public static class Parameters extends RelationBookmarkParams {

    @ParameterField
    public DatasetMode mode = DatasetMode.RELATION;

    @ParameterField(minRequired = 1)
    public List<String> dataset;

    @ParameterField
    public Optional<String> includeAttributes = Optional.empty();

    @ParameterField
    public int readSizeMb = 64;

    public H5File h5File;

    @Override
    public Class<?> getDataType() {
      if (mode != null) {
        return mode.getDataType();
      }
      return null;
    }

    H5File getH5File() throws ProblemException {
      if (h5File == null) {
        Path path = getBookmarkedPath();
        if (path == null) {
          // we don't expect to get here anymore because if getBookmarkedPath cannot make a file it
          // should throw a RiskscapeIOException, but just in case we keep this handling rather than
          // risk a NPE on when making the H5File.
          throw new ProblemException(PROBLEMS.hdf5FileLoadError("HDF5 plugin could not load file"));
        }
        h5File = new H5File(path);
      }
      // TODO when to close? on finalize?
      return h5File;
    }

    H5Dataset openDataset(String name) throws ProblemException {
      try {
        return getH5File().openDataset(H5DatasetPath.parse(name));
      } catch (HDF5SymbolTableException e) {
        throw new ProblemException(PROBLEMS.hdf5DatasetDoesNotExist(name));
      } catch (HDF5Exception e) {
        throw new ProblemException(PROBLEMS.hdf5FileLoadError(e.getMessage()));
      }
    }

    List<String> getIncludeAttributeList() {
      return includeAttributes.map(attributesString -> {
        return Arrays.stream(attributesString.split(","))
            .map(param -> param.trim())
            .collect(Collectors.toList());
      }).orElse(Collections.emptyList());
    }

  }

  @Getter
  private final Map<String, String> extensionsToFormats = ImmutableMap.of("hdf5", "hdf5");

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

  @Override
  protected ResultOrProblems build(Parameters parameters) {
    switch (parameters.mode) {
      case RELATION:
        return super.build(parameters);
      default:
        throw new UnsupportedOperationException(parameters.mode.toString());
    }
  }

  private H5Relation createRelation(Parameters params, String datasetName) throws ProblemException {
    H5Dataset dataset = params.openDataset(datasetName);

    Struct producedType = getProducedType(dataset, params);
    Function<H5DatasetCursor, Tuple> bufferToTupleMapper = getBufferToTupleMapper(dataset.getDataType(),
        producedType);

    return new H5Relation(dataset, producedType, bufferToTupleMapper, params.readSizeMb);
  }

  /**
   * Builds a relation for the desired dataset. The returned relation by default will contain all
   * attributes from the target dataset, but only those in {@link Parameters#getIncludeAttributeList() }
   * will be included if that is not empty.
   */
  @Override
  protected ResultOrProblems<Relation> createRawRelationFromBookmark(Parameters params) {
    return ProblemException.catching(() -> {
      if (params.dataset.size() == 1) {
        return createRelation(params, params.dataset.get(0));
      } else {
        // we want to read from multiple datasets and them zip them together into a single result
        List<BaseRelation> relations = new ArrayList<>();
        for (String datasetName : params.dataset) {
          relations.add(createRelation(params, datasetName));
        }
        try {
          return new ZipRelation(relations);
        } catch (RiskscapeException ex) {
          // map the generic error to a more specific HDF5 one
          if (ZipRelation.PROBLEMS.mismatchingRows().equals(ex.getProblem())) {
            throw new ProblemException(PROBLEMS.differentDatasetSizes());
          }
          throw new ProblemException(ex.getProblem());
        }
      }
    });
  }

  /**
   * Returns the {@link Struct} that will be produced from the given dataType and parameters.
   */
  private Struct getProducedType(H5Dataset dataset, Parameters params) throws ProblemException {
    List<String> includeAttributes = params.getIncludeAttributeList();
    StructBuilder builder = Struct.builder();
    List<String> available = new ArrayList<>();
    H5Type dataType = dataset.getDataType();

    if (dataType instanceof H5CompoundType) {
      H5CompoundType compoundDataType = (H5CompoundType) dataType;

      for (H5CompoundMember member : compoundDataType.getMembers()) {
        available.add(member.name);

        // If include-attributes is empty, then the default is to include all. Otherwise we
        // only include the attributes that the bookmark is set to return.
        if (includeAttributes.isEmpty() || includeAttributes.contains(member.name)) {
          builder.add(member.name, member.type.toType());
        }
      }
    } else if (includeAttributes.size() == 1) {
      // user has specified 'include-attributes', so use that as the member name
      available.add(includeAttributes.get(0));
      builder.add(includeAttributes.get(0), dataType.toType());
    } else {
      // just use the last part of the dataset path as the attribute name
      // e.g. /foo/datasets/bar => 'bar' attribute
      available.add(dataset.getDatasetName());
      builder.add(dataset.getDatasetName(), dataType.toType());
    }
    Struct produced = builder.buildOr().getOrThrow();

    // Users could set a bookmark to include attributes that don't actually exist in the data.
    // This should be noticed and be a problem
    for (String attribute : includeAttributes) {
      if (!produced.hasMember(attribute)) {
        throw new ProblemException(GeneralProblems.get().notAnOption(attribute, "include-attributes", available));
      }
    }

    return produced;
  }

  /**
   * Returns a {@link BiFunction} that will produce a {@link Tuple} of producedType when it is passed
   * a {@link ByteBuffer} and offset (into said byte buffer) where the HDF5 element is located.
   */
  private Function<H5DatasetCursor, Tuple> getBufferToTupleMapper(H5Type dataType, Struct producedType) {
    if (!(dataType instanceof H5CompoundType)) {
      return (cursor) -> Tuple.ofValues(producedType, cursor.peek());
    }

    List<H5CompoundMember> membersToFetch = new ArrayList<>();
    H5CompoundType compoundDataType = (H5CompoundType) dataType;
    for (StructMember member: producedType.getMembers()) {
      membersToFetch.add(compoundDataType.findMember(member.getKey()));
    }

    return (cursor) -> {
      Object[] values = membersToFetch.stream()
          .map(member -> cursor.peek(member))
          .toArray();
      return Tuple.ofValues(producedType, values);
    };
  }


  @Override
  protected void validateParameters(Parameters params, BindingContext context) {
    super.validateParameters(params, context);

    if (! ALLOWED_READ_SIZE_MB.contains(params.readSizeMb)) {
      params.problems.add(GeneralProblems.get().valueOutOfRange("read-size-mb", params.readSizeMb,
          ALLOWED_READ_SIZE_MB));
    }
  }

}
