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


import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
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.engine.types.Type;
import nz.org.riskscape.engine.types.Value;

/**
 * Offers quick, efficient, Map-like storage (populating the map is 100% quicker than a {@link HashMap} where the keys
 * are defined according to a {@link Struct}.  Random access to values is slightly slower (~30%) than with a standard
 * {@link HashMap}.  Does not guarantee type correctness in the same way that the {@link Value} class does
 *
 * Tuples are used to store user data as it flows through the pipeline. A tuple is a set of attributes
 * (e.g. a row of input data), but tuples can also be comprised of other tuples (e.g. the 'exposure' tuple contains
 * all the exposure-layer attributes). When a tuple is passed to a user's Python function, it gets converted into
 * a Python dictionary.
 */
@Slf4j
public class Tuple {

  /**
   * Options for error suppression/checking when using the {@link Tuple#coerce(Struct, Map)} method
   * @author nickg
   *
   */
  public enum CoerceOptions {
    /**
     * Any extra keys in the map will be ignored and not included in the resulting tuple
     */
    SURPLUS_IGNORED,
    /**
     * Suppresses an error for when {@link Tuple#hasNulls()} returns true after building the tuple from the map
     */
    MISSING_IGNORED
  }

  /**
   * 'Helpful' constant for an empty tuple - built from {@link Struct#EMPTY_STRUCT}
   */
  public static final Tuple EMPTY_TUPLE = new Tuple(Struct.EMPTY_STRUCT);

  /**
   * Convenience constructor for building an empty tuple of this type
   */
  public static Tuple of(Struct type) {
    return new Tuple(type);
  }

  /**
   * Convenience constructor for building a tuple with a single key/value
   */
  public static Tuple of(Struct type, String key, Object value) {
    return new Tuple(type).set(key, value);
  }

  /**
   * Convenience constructor for building a tuple with two keys/values
   */
  public static Tuple of(Struct type, String key1, Object value1, String key2, Object value2) {
    return new Tuple(type).set(key1, value1).set(key2, value2);
  }

  /**
   * Convenience constructor for building a tuple from an ordered (wrt to the struct) set of values
   */
  public static Tuple ofValues(Struct type, Object... values) {
    final int size = type.getMembers().size();
    if (values.length == size) {
      Object[] dest = new Object[size];
      // TODO can we just take ownership?  for varargs, it'll be implicitly creating a new array anyway
      System.arraycopy(values, 0, dest, 0, size);
      return new Tuple(type, dest);
    } else {
      Tuple tuple = new Tuple(type);
      for (int i = 0; i < values.length; i++) {
        tuple.set(type.getMembers().get(i), values[i]);
      }
      return tuple;
    }

  }

  public static Tuple coerce(Struct type, Map<?, ?> rawValues) {
    return coerce(type, rawValues, EnumSet.noneOf(CoerceOptions.class));
  }

  /**
   * Build a tuple from the given map, applying the type's coercion to the result.
   *
   * @param type The type to coerce the map to
   * @param rawValues A map of uncoerced values
   * @param options Options for error handling
   * @return a fresh tuple
   */
  public static Tuple coerce(Struct type, Map<?, ?> rawValues, EnumSet<CoerceOptions> options) {
    Tuple tuple = new Tuple(type);

    rawValues.entrySet().stream().forEach(entry -> {
      StructMember structMember = type.getEntry((String) entry.getKey());
      if (structMember == null) {
        if (!options.contains(CoerceOptions.SURPLUS_IGNORED)) {
          throw new NoSuchMemberException(type, (String) entry.getKey());
        }
      } else {
        tuple.set(structMember, structMember.getType().coerce(entry.getValue()));
      }
    });

    if (!options.contains(CoerceOptions.MISSING_IGNORED) && tuple.hasRequiredNulls()) {
      List<String> missing = new ArrayList<>();
      for (StructMember member : tuple.struct.getMembers()) {
        if (tuple.fetch(member) == null && !member.getType().isNullable()) {
          missing.add(member.getKey());
        }
      }

      throw new CoercionException(
          rawValues,
          type,
          "Supplied map '%s' is missing required keys - %s",
          rawValues,
          missing);
    }

    return tuple;
  }

  /**
   * The struct defines the shape of the data - the attribute keys and what RiskScape {@link Type} they are.
   */
  @Getter
  private Struct struct;

  /**
   * The entries store the actual data values. The entries are indexed by the struct member, i.e. `entries[0]` is
   * the corresponding value for the first struct member.
   */
  private Object[] entries;

  public Tuple(Struct struct) {
    this.struct = struct;
    this.entries = new Object[struct.size()];
  }

  private Tuple(Struct struct, Object[] entries) {
    this.struct = struct;
    this.entries = entries;
  }

  public Tuple(Struct struct, Map<String,String> values) {
    this(struct);
    values.entrySet().forEach(entry -> {
      set(entry.getKey(), entry.getValue());
    });
  }

  /**
   * Set a value by its key
   * @return this to allow tuple.set("foo", bar).set("baz", foo)...
   * @throws NoSuchMemberException if no member named key in underlying struct
   */
  public Tuple set(@NonNull String key, Object value) {
    entries[find(key).getIndex()] = value;
    return this;
  }

  /**
   * Set a value using a {@link StructMember} as a key
   * @return this to allow tuple.set("foo", bar).set("baz", foo)...
   * @throws IllegalArgumentException if key doesn't belong to struct this tuple is of
   */
  public Tuple set(@NonNull StructMember key, Object value) {
    assert checkOwner(key);
    entries[key.getIndex()] = value;
    return this;
  }

  /**
   * Set a value using its struct index
   */
  public Tuple set(int index, Object value) {
    // TODO assert correct type?
    entries[index] = value;
    return this;
  }

  /**
   * Populate this tuple with the given values
   */
  public void setAll(List<Object> values) {
    int index = 0;
    for (Object object : values) {
      // TODO assert correct type
      entries[index++] = object;
    }
  }

  /**
   * Copies (appends) the values from the given source tuple into this tuple.
   * Useful when merging or combining tuples together
   *
   * @param source     Tuple containing the values to copy
   * @param destOffset The struct index in this tuple to start copying into
   */
  public void setAll(int destOffset, Tuple source) {
    int index = destOffset;
    for (Object object : source.entries) {
      // TODO assert correct type
      entries[index++] = object;
    }
  }

  /**
   * Copies the values from the given source tuple into this tuple.
   */
  public void setAll(Tuple source) {
    setAll(0, source);
  }

  /**
   * Clears a value from the tuple
   */
  public void remove(@NonNull String key) {
    entries[find(key).getIndex()] = null;
  }

  /**
   * Clears a value from the tuple
   */
  public void remove(@NonNull StructMember key) {
    assert checkOwner(key);
    entries[key.getIndex()] = null;
  }

  /**
   * Return a value from the tuple, with an implicit cast to the receiving type
   */
  @SuppressWarnings("unchecked")
  public <T> T fetch(String key) {
    return (T) entries[find(key).getIndex()];
  }

  /**
   * Return a value from the tuple, with an implicit cast to the receiving type
   * @param entry the member of the struct we're fetching.  Slightly more efficient than lookup by key
   */
  @SuppressWarnings("unchecked")
  public <T> T fetch(StructMember entry) {
    assert checkOwner(entry);
    return (T) entries[entry.getIndex()];
  }

  @SuppressWarnings("unchecked")
  public <T> T fetch(int index) {
    return (T) entries[index];
  }
  /**
   * TBD is this actually necessary?
   */
  public Tuple fetchChild(String key) {
    StructMember entry = find(key);

    if (Nullable.unwrap(entry.getType()).find(Struct.class).isPresent()) {
      return (Tuple) entries[entry.getIndex()];
    } else {
      throw new IllegalArgumentException(key + " is not a struct");
    }
  }

  /**
   * Shortcut for accessing the type of a member
   * @param key
   * @return Type of member
   */
  public Type getType(String key) {
    return find(key).getType();
  }

  private StructMember find(String key) {
    StructMember entry = struct.getEntry(key);

    if (entry == null) {
      throw new NoSuchMemberException(struct, key);
    }

    return entry;
  }

  /**
   * Perform a tuple-aware clone of this tuple, performing a shallow copy of the entry array, unless the entry is
   * itself a tuple, which will be recursively cloned
   */
  @Override
  public Tuple clone() {
    Tuple clone = new Tuple(struct, Arrays.copyOf(entries, entries.length));

    for (int i = 0; i < clone.entries.length; i++) {
      if (clone.entries[i] instanceof Tuple) {
        Tuple child = (Tuple) clone.entries[i];
        clone.entries[i] = child.clone();
      }
    }

    return clone;
  }

  // declare to silence checkstyle and make it clear equals has a friend
  @Override
  public int hashCode() {
    return Objects.hash(entries);
  }

  @Override
  public boolean equals(Object rhs) {
    if (rhs instanceof Tuple) {
      Tuple rhsTuple = (Tuple) rhs;

      if (!rhsTuple.struct.equals(struct)) {
        return false;
      }

      for (int i = 0; i < this.entries.length; i++) {
        if (!Objects.equals(rhsTuple.entries[i], this.entries[i])) {
          return false;
        }
      }

      return true;
    } else if (rhs instanceof Map) {
      return super.equals(rhs);
    } else {
      return false;
    }
  }

  /**
   * Checks that the given key belongs to this, or an equal, struct, throwing an assertion error if it is not.
   * This method should only be called with the assert keyword, so that it's restricted to when `-enableassertions` is
   * on (which is true of our unit tests) to avoid slowing down model runs with this check - it's not particularly
   * expensive, but this is a *very* hot code path.
   *
   * Note that this check used to run all the time, but did a much cheaper identity check.  This gave rise to a lot of
   * false negatives - the structs were equal, but not the same.  We introduced some struct normalization to make it
   * easier to avoid these types of errors, but that was papering over this issue and required extra code/thought.  This
   * seems like a better compromise - make this expensive (in total) operation only run when assertions are enabled.
   */
  private boolean checkOwner(StructMember key) {
    // NB this check is only going to run in test-mode - the reason being we don't want to slow down model
    // runs with a ton of equality checks, instead we'll do this more expensive and permissive
    if (!key.getOwner().equals(this.struct)) {
      log.warn("Detected mismatching struct access - key is from {}, but our struct is {}", key.getOwner(), struct);
      return false;
    }
    return true;
  }

  public int size() {
    return this.entries.length;
  }

  /**
   * @return if any of this tuples entries are null.
   */
  public boolean hasNulls() {
    for (int i = 0; i < this.entries.length; i++) {
      if (this.entries[i] == null) {
        return true;
      }
    }
    return false;
  }

  /**
   * @return if any of this tuples non-nullable entries are null.
   */
  public boolean hasRequiredNulls() {
    List<StructMember> members = this.struct.getMembers();
    for (int i = 0; i < this.entries.length; i++) {
      if (this.entries[i] == null && !members.get(i).getType().isNullable()) {
        return true;
      }
    }
    return false;
  }

  /**
   * Estimate the number of bytes used to store this {@link Tuple} when serialized.
   * @return the number of bytes.
   */
  public int estimateSize() {
    int counter = 0;

    List<StructMember> members = this.struct.getMembers();
    for (int i = 0; i < members.size(); i++) {
      Object entry = this.entries[i];

      if (entry == null) {
        continue;
      }

      StructMember member = members.get(i);
      counter = counter + member.getType().estimateSize(entry);
    }

    return counter;
  }

  @Override
  public String toString() {
    StringBuilder bldr = new StringBuilder("{");

    for (StructMember member : struct.getMembers()) {
      if (bldr.length() != 1) {
        bldr.append(", ");
      }

      bldr.append(member.getKey()).append("=").append(member.getType().toString((this.fetch(member))));
    }

    return bldr.append("}").toString();
  }

  /**
   * @return a read-only view on to the underlying values that back this tuple
   */
  public List<Object> getValues() {
    return Collections.unmodifiableList(Arrays.asList(this.entries));
  }

  public Map<String, Object> toMap() {
    Map<String, Object> newMap = new HashMap<>(entries.length);

    for (StructMember member : struct.getMembers()) {
      Object value = fetch(member);

      if (value != null) {
        if (value instanceof Tuple tuple) {
          value = tuple.toMap();
        }
        newMap.put(member.getKey(), value);
      }
    }

    return newMap;
  }

  public Object[] toArray() {
    Object[] clone = new Object[entries.length];
    System.arraycopy(entries, 0, clone, 0, entries.length);
    return clone;
  }

}
