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

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import lombok.ToString;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.Tuple.CoerceOptions;
import nz.org.riskscape.engine.data.InputDataProblems;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.relation.UnexpectedNullValueException;
import nz.org.riskscape.engine.typexp.AST;
import nz.org.riskscape.engine.typexp.AST.ComplexType;
import nz.org.riskscape.engine.typexp.AST.Symbol;
import nz.org.riskscape.engine.typexp.ComplexTypeConstructor;
import nz.org.riskscape.engine.typexp.TypeArgumentException;
import nz.org.riskscape.engine.typexp.TypeBuilder;
import nz.org.riskscape.engine.util.Pair;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.ExpressionProblems;

/**
 * Permissive struct-like type that allows a c-like struct (like a java class with only fields) to be declared
 * with members being other riskscape types.
 *
 * Struct should be constructed using the builder. eg:
 *
 * ```java
 * Struct.of("name", Types.TEXT).and("value", Types.INTEGER).build();
 * ```
 */
public class Struct implements Type {

  public static final Class<?> INTERNAL_TYPE = Tuple.class;

  public static final TypeInformation TYPE_INFORMATION = new TypeInformation(
      "struct",
      Struct.class,
      Tuple.class,
      (ComplexTypeConstructor) (TypeBuilder typeBuilder, ComplexType type) -> {
        List<AST> args = type.args();

        if (args.size() != 1) {
          throw new TypeArgumentException(type, "one argument expected");
        }

        AST.Dictionary dict = typeBuilder.expectAST(AST.Dictionary.class, args.remove(0),
            (ast) -> new TypeArgumentException(type, "argument must be a dictionary"));

        StructBuilder builder = new StructBuilder();
        Set<Entry<Symbol,AST>> entrySet = dict.values.entrySet();

        for (Entry<Symbol, AST> entry : entrySet) {
          Type memberType = typeBuilder.expectType(entry.getValue(),
              (ast) -> new TypeArgumentException(type, "dictionary values must be types"));

          builder.add(entry.getKey().ident(), memberType);
        }

        return builder.build();
      }
  );

  @Data
  @RequiredArgsConstructor
  @EqualsAndHashCode(exclude="owner")
  @ToString(exclude="owner")
  public static final class StructMember {
    @NonNull
    final java.lang.String key;
    @NonNull
    final Type type;
    final int index;
    Struct owner;

    @Override
    public StructMember clone() {
      return new StructMember(key, type, index);
    }
  }

  /**
   * Builder class for constructing struct types with a fluid api
   */
  public static class StructBuilder {
    private final List<StructMember> pairs;

    public StructBuilder() {
      pairs = new ArrayList<>();
    }

    public StructBuilder(int expectedSize) {
      pairs = new ArrayList<>(expectedSize);
    }

    /**
     * @deprecated use add
     */
    @Deprecated
    public StructBuilder and(java.lang.String key, Type type) {
      return add(key, type);
    }

    public StructBuilder add(String key, Type type) {
      this.pairs.add(new StructMember(key, type, pairs.size()));
      return this;
    }

    public StructBuilder addAll(Struct type) {
      for (StructMember member : type.members) {
        this.pairs.add(new StructMember(member.key, member.type, pairs.size()));
      }
      return this;
    }

    /**
     * Variation of build that returns a ResultOrProblems that can be used in user error reporting, rather than
     * throwing an exception (on key duplicates)
     * @return ResultOrProblems
     */
    public ResultOrProblems<Struct> buildOr() {
      List<StructMember> memberList = new ArrayList<>(this.pairs.size());
      for (StructMember member : this.pairs) {
        memberList.add(member.clone());
      }
      try {
        return ResultOrProblems.of(new Struct(memberList));
      } catch (DuplicateKeysException ex) {
        return ResultOrProblems.error(ex);
      }
    }

    /**
     * Finish defining a struct
     * @return the struct that has been defined.
     */
    public Struct build() {
      List<StructMember> members = new ArrayList<>(this.pairs.size());
      for (StructMember member : this.pairs) {
        members.add(member.clone());
      }
      return new Struct(members);
    }

    public boolean isEmpty() {
      return this.pairs.isEmpty();
    }

    public int size() {
      return this.pairs.size();
    }
  }

  public static final Struct EMPTY_STRUCT = new StructBuilder().build();


  /**
   * Create an empty struct
   */
  public static Struct of() {
    return EMPTY_STRUCT;
  }

  /**
   * Define a struct, starting with the given key and type
   * @return a builder to use to define the struct.
   */
  public static Struct of(String key, Type type) {
    return new StructBuilder().add(key, type).build();
  }

  public static Struct of(String k0, Type t0, String k1, Type t1) {
    return new StructBuilder().add(k0, t0).add(k1, t1).build();
  }

  public static Struct of(String k0, Type t0, String k1, Type t1, String k2, Type t2) {
    return new StructBuilder().add(k0, t0).add(k1, t1).add(k2, t2).build();
  }

  public static Struct of(String k0, Type t0, String k1, Type t1, String k2, Type t2, String k3, Type t3) {
    return new StructBuilder().add(k0, t0).add(k1, t1).add(k2, t2).add(k3, t3).build();
  }

  public static Struct of(String k0, Type t0, String k1, Type t1, String k2, Type t2, String k3, Type t3, String k4,
      Type t4) {
    return new StructBuilder().add(k0, t0).add(k1, t1).add(k2, t2).add(k3, t3).add(k4, t4).build();
  }

  public static StructBuilder builder() {
    return new StructBuilder();
  }

  /**
   * @return a list of member keys in DFS order, where children have their key appended to their parent's with a dot
   * between them, so they are in the style a user would enter in an expression, e.g. `foo.bar`
   */
  public static List<String> flattenMemberKeys(Struct struct) {
    List<String> names = new LinkedList<>();
    flattenMemberKeys(names, "", struct);
    return names;
  }

  private static void flattenMemberKeys(List<String> appendTo, String prefix, Struct struct) {
    for (StructMember member : struct.getMembers()) {
      Struct unwrapped = member.getType().findAllowNull(Struct.class).orElse(null);
      String name = prefix + member.getKey();
      if (unwrapped != null) {
        flattenMemberKeys(appendTo, name + ".", unwrapped);
      } else {
        appendTo.add(name);
      }
    }
  }

  private final List<StructMember> members;
  private final Map<String, StructMember> membersMap;

  /**
   * Private constructor - use the static constructors, as this one will reuse whatever list it has been given, which
   * could lead to bugs if used improperly
   */
  private Struct(List<StructMember> members) {
    this.members = members;
    this.membersMap = new HashMap<>(members.size());

    Set<StructMember> collisions = new HashSet<>();
    for (StructMember member : members) {
      StructMember replaced = membersMap.put(member.key, member);
      if (replaced != null) {
        collisions.add(member);
        collisions.add(replaced);
      }

      // safe to assert this, we don't have any APIs that should allow it
      assert member.owner == null: "Members can not be reused, create a new one";

      member.owner = this;
    }

    if (!collisions.isEmpty()) {
      throw new DuplicateKeysException(collisions);
    }
  }

  /**
   * @return a {@link List} of the names/keys of all the members in this struct, sorted alphabetically.
   */
  public List<String> getMemberKeys() {
    List<String> keys = new ArrayList<>(membersMap.size());
    keys.addAll(membersMap.keySet());
    Collections.sort(keys);

    return keys;
  }

  public List<StructMember> getMembers() {
    return Collections.unmodifiableList(members);
  }

  @Override
  public Object coerce(Object value) {
    if (value instanceof Tuple && ((Tuple) value).getStruct() == this) {
      // source is already of this type - don't bother coercing
      return value;
    }

    if (value instanceof Map) {
      Map sourceTuple = (Map) value;
      return Tuple.coerce(this, sourceTuple, EnumSet.of(CoerceOptions.SURPLUS_IGNORED));
    } else if (value instanceof Tuple tuple) {
      return Tuple.coerce(this, tuple.toMap(), EnumSet.of(CoerceOptions.SURPLUS_IGNORED));
    } else if (members.size() == 1) {
      StructMember m = members.get(0);
      Map<String,Object> map = new HashMap<>();
      map.put(m.key, m.type.coerce(value));
      return Tuple.coerce(this, map, EnumSet.of(CoerceOptions.MISSING_IGNORED));
    } else {
      throw new CoercionException(value, this);
    }
  }

  @Override
  public Class<?> internalType() {
    return INTERNAL_TYPE;
  }

  /**
   * Compatability shim to cope with API change to Struct.of/Struct.and
   * @deprecated
   */
  @Deprecated
  public Struct build() {
    return this;
  }

  @Override
  public String toString() {
    return members.stream()
          .map(se -> String.format("%s=>%s", se.key, se.type))
          .collect(Collectors.joining(", ", "{", "}"));
  }

  /**
   * @return a StructMember identified by the given key, or Optional.empty() if non exists
   */
  public Optional<StructMember> getMember(String key) {
    return Optional.ofNullable(membersMap.get(key));
  }

  /**
   * null-returning variation of getMember
   * TODO consider making this package level or change it's name or something
   */
  public StructMember getEntry(String key) {
    return membersMap.get(key);
  }

  public Struct parent(String asKey) {
    return Struct.of(asKey, this).build();
  }

  /**
   * Returns this with the extra member added on
   * TODO using some kind of optimized copying immutable map might cut down on garbage collection etc
   * (if guava isn't already doing that)
   */
  public Struct add(String asKey, Type otherType) {
    List<StructMember> newMembers = new ArrayList<>(this.size() + 1);
    for (StructMember member : this.members) {
      newMembers.add(member.clone());
    }
    newMembers.add(new StructMember(asKey, otherType, this.size()));

    return new Struct(newMembers);
  }

  /**
   * Returns a ResultOrProblems wrapping a new struct without the given members
   * Adds non-present keys as NoSuchStructMember problems
   */
  public ResultOrProblems<Struct> remove(List<String> toRemove) {

    ArrayList<String> stillToRemove = new ArrayList<>(toRemove);

    List<StructMember> newMembers = new ArrayList<>(this.size());
    int count = 0;

    for (StructMember member: this.members) {
      if (!stillToRemove.contains(member.key)) {
        newMembers.add(new StructMember(member.key, member.type, count));
        count++;
      } else {
        stillToRemove.remove(member.key);
      }
    }

    if (!stillToRemove.isEmpty()) {
      return ResultOrProblems.of(new Struct(newMembers),
          stillToRemove
            .stream()
            .map(key -> ExpressionProblems.get().noSuchStructMember(key, this.getMemberKeys()))
            .toList()
      );
    }

    return ResultOrProblems.of(new Struct(newMembers));
  }

  /**
   * @return a new Struct with the given member replaced with one of a new type.
   */
  public Struct replace(String key, Type newType) {
    List<StructMember> newMembers = new ArrayList<>(this.size());
    boolean replaced = false;
    int count = 0;
    for (StructMember member : this.members) {
      Type type;
      if (member.getKey().equals(key)) {
        type = newType;
        replaced = true;
      } else {
        type = member.getType();
      }
      newMembers.add(new StructMember(member.getKey(), type, count));
      count++;
    }

    if (!replaced) {
      throw new IllegalArgumentException(String.format(
          "This struct does not contain a member with key %s, has %s",
          key,
          this.members));
    }

    return new Struct(newMembers);
  }

  /**
   * Synonym of {@link #add(String, Type)} to support change in builder API
   */
  public Struct and(Struct struct) {
    return builder().addAll(this).addAll(struct).build();

  }

  /**
   * Synonym of {@link #add(String, Type)} to support change in builder API
   */
  public Struct and(String newKey, Type newType) {
    return add(newKey, newType);
  }

  @Override
  public Struct asStruct() {
    return this;
  }

  public int size() {
    return this.members.size();
  }

  public boolean contains(StructMember entry) {
    return this.members.contains(entry);
  }

  /**
   * @return true if the two structs have the same ordered list of member types - e.g. their tuples could be combined
   * in a union
   */
  public boolean isUnion(Struct rhs) {
    return this.members.stream().map(StructMember::getType).collect(Collectors.toList()).equals(
        rhs.members.stream().map(StructMember::getType).collect(Collectors.toList()));
  }

  /**
   * @return true if the two structs have the same members, considering only the keys and the types.
   */
  public boolean isEquivalent(Struct rhs) {
    return this.members.stream().map(se -> Arrays.asList(se.key, se.type)).collect(Collectors.toSet()).equals(
        rhs.members.stream().map(se -> Arrays.asList(se.key, se.type)).collect(Collectors.toSet()));
  }

  /**
   * @return true if this struct is a superset, that is contains at least the same members as rhs.
   */
  public boolean isSupersetOf(Struct rhs) {
    throw new RuntimeException("not yet implemented");
  }

  public boolean hasMember(String key) {
    return membersMap.containsKey(key);
  }

  @Override
  public int hashCode() {
    return this.members.hashCode();
  }

  @Override
  public boolean equals(Object rhs) {
    if (rhs instanceof Struct) {
      Struct rhsStruct = (Struct) rhs;

      return this.members.equals(rhsStruct.members);
    } else {
      return false;
    }
  }

  @Override
  public int estimateSize(Object entry) {
    if (entry instanceof Tuple) {
      Tuple tuple = (Tuple) entry;
      return tuple.estimateSize();
    }

    return 0;
  }

  @Override
  public void toBytes(DataOutputStream os, Object toWrite) throws IOException {
    Tuple value = (Tuple) toWrite;

    // see Tuple#checkOwner for why this check is an assert
    assert value.getStruct().equals(this);

    for (StructMember member : members) {
      Object valueToWrite = value.fetch(member);

      // sanity-check for unexpected null values - they may be present in the input
      // data, but we won't be able to write them safely
      if (valueToWrite == null && !member.type.isNullable()) {
        throw new UnexpectedNullValueException(
            InputDataProblems.get().unexpectedNullValue(member.key, TupleUtils.getContentSummary(value)));
      }
      member.type.toBytes(os, valueToWrite);
    }
  }

  @Override
  public Object fromBytes(DataInputStream in) throws IOException {
    Tuple tuple = new Tuple(this);

    for (StructMember member : members) {
      Object read = member.type.fromBytes(in);
      tuple.set(member, read);
    }

    return tuple;
  }

  public Struct addOrReplace(String key, Type type) {
    if (this.hasMember(key)) {
      return replace(key, type);
    } else {
      return add(key, type);
    }
  }

  @Override
  public <T, U> U visit(TypeVisitor<T, U> tv, T data) {
    // function needs an explicit type signature or the compiler complains
    final Function<StructMember, Pair<Type, ?>> toPair = m -> Pair.of(m.getType(), m.getKey());
    return tv.compoundType(this, members.stream().map(toPair).toList(), data);
  }
}

