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


import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.function.Consumer;

import com.google.common.base.Strings;
import com.google.common.collect.Lists;

import lombok.Getter;
import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.ConstructionCallback;
import nz.org.riskscape.engine.Identified;
import nz.org.riskscape.engine.IdentifiedCollection;
import nz.org.riskscape.engine.IdentifiedException;
import nz.org.riskscape.engine.ObjectAlreadyExistsException;
import nz.org.riskscape.engine.Reference;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.resource.Resource;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.TypeRegistry;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.engine.types.ancestor.AncestorRule;
import nz.org.riskscape.engine.types.ancestor.AncestorType;
import nz.org.riskscape.engine.types.ancestor.AncestorTypeList;
import nz.org.riskscape.engine.types.eqrule.Coercer;
import nz.org.riskscape.engine.types.eqrule.EquivalenceRule;
import nz.org.riskscape.engine.types.varule.Variance;
import nz.org.riskscape.engine.types.varule.VarianceRule;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemSink;


@RequiredArgsConstructor
@Slf4j
public class TypeSet extends IdentifiedCollection.Base<CanonicalType> implements TypeRules {

  /**
   * A jvm-global type set, only use with extreme prejudice
   */
  public static final TypeSet GLOBAL_TYPE_SET = new TypeSet(TypeRegistry.withDefaults());

  @Getter
  private final TypeRegistry typeRegistry;

  public TypeSet() {
    this.typeRegistry = TypeRegistry.withDefaults();
  }


  public void addType(String identifier, Resource resource, ConstructionCallback<Type> builder) {
    super.add(identifier, resource, () -> {
      return builder.create().map(t -> wrap(identifier, t));
    });
  }

  public CanonicalType add(@NonNull String id, @NonNull Type type) throws ObjectAlreadyExistsException {
    CanonicalType ct = wrap(id, type);
    add(ct);
    return ct;
  }

  private CanonicalType wrap(String id, Type type) {
    return new CanonicalType(id, type) {
      @Override
      public TypeSet getTypeSet() {
        return TypeSet.this;
      }
    };
  }

  public LinkedType getLinkedType(@NonNull String id) {
    return new LinkedType(this, id);
  }

  @Deprecated
  public CanonicalType getRequired(@NonNull String id) {
    return getRequired(id, ProblemSink.DEVNULL);
  }

  public CanonicalType getRequired(@NonNull String id, ProblemSink problemSink) {
    try {
      return get(id, problemSink);
    } catch (IdentifiedException e) {
      throw new MissingTypeException(id, e);
    }
  }


  @Override
  public Optional<Coercer> findEquivalenceCoercer(Type sourceType, Type targetType) {
    log.debug("Equivalence search - {} to {}", sourceType, targetType);
    return new WrappedTypeRules(0, this, null).findEquivalenceCoercer(sourceType, targetType);
  }

  @Override
  public Variance testVariance(Type sourceType, Type targetType) {
    log.debug("Variance test - {} to {}", sourceType, targetType);
    return new WrappedTypeRules(0, this, null).testVariance(sourceType, targetType);
  }

  @Override
  public Optional<AncestorType> computeAncestorType(Type type1, Type type2) {
    return new WrappedTypeRules(0, this, null).computeAncestorType(type1, type2);
  }

  /**
   * Equivalent of the old Types#computeAncestor method.  Will only return a type if both types are directly assignable
   * to an ancestor, or ANYTHING if there is no ancestor.
   */
  public Type computeAncestorNoConversion(Type type1, Type type2) {
    Type ancestorType = computeAncestorType(type1, type2)
        .map(ancestor -> ancestor.isConverting() ? null : ancestor)
        .map(AncestorType::getType)
        .orElse(Types.ANYTHING);

    // kludge/workaround: for reasons of hating myself, I made anything not be able to represent null.  To that end,
    // we need to commute the nullability of type1 & type2 to the ancestor. Having written this, I suspect it would make
    // very little practical difference if we allowed anything to hold nulls
    return Nullable.ifTrue(Nullable.is(type1) || Nullable.is(type2), ancestorType);
  }

  /**
   * A list version of computeAncestorType that uses a recursive algorithm to pair-wise 'unify' all the bits in the list
   * to a single type and a conversion function if the values need to change
   */
  public AncestorTypeList computeAncestors(List<Type> allTypes) {
    AncestorTypeList recursed = computeAncestorsRecursive(allTypes);

    // don't bother converting elements if we've failed (no need to convert - anything can hold anything)
    if (recursed.getType() == Types.ANYTHING) {
      if (Nullable.any(allTypes.toArray(new Type[0]))) {
        // need to respect nullable-ness of types, sigh
        return AncestorTypeList.of(Nullable.ANYTHING);
      } else {
        return AncestorTypeList.of(Types.ANYTHING);
      }
    }

    return recursed;
  }
  public AncestorTypeList computeAncestorsRecursive(List<Type> allTypes) {
    List<Type> distinctTypes = allTypes.stream().distinct().toList();

    if (allTypes.size() == 0) {
      // empty list case
      return AncestorTypeList.of(Types.NOTHING);
    }

    if (distinctTypes.size() == 1) {
      // if everything is the same type, we're done
      return AncestorTypeList.of(distinctTypes.get(0));
    } else if (distinctTypes.size() == 2) {
      // if we only have two types then we can use a single result from computeAncestorType and apply any conversion
      // over the whole list
      AncestorType ancestorType =
          computeAncestorType(distinctTypes.get(0), distinctTypes.get(1)).orElse(AncestorType.of(Types.ANYTHING));

      if (ancestorType.isConverting()) {
        return AncestorTypeList.of(
          ancestorType.getType(),
          list -> Lists.transform(list, element -> ancestorType.getConvert().apply(element))
        );
      } else {
        return AncestorTypeList.of(ancestorType.getType());
      }
    } else {
      // if there's more than two distinct types, then we divide and conquer

      // recurse over the tail to get it unified
      List<Type> tail = allTypes.subList(1, allTypes.size());
      AncestorTypeList convertedTail = computeAncestors(tail);

      // now fold the head in to the tail
      Type head = allTypes.get(0);
      AncestorTypeList convertedHead =
          computeAncestors(Arrays.asList(head, convertedTail.getType()));

      return AncestorTypeList.of(convertedHead.getType(), list -> {
        assert list.size() == allTypes.size();

        Object headElement = list.get(0);
        List<Object> tailElements = convertedTail.getConverter().apply(list.subList(1, list.size()));
        List<Object> combined = new ArrayList<>(tailElements.size() + 1);

        combined.add(headElement);
        combined.addAll(tailElements);

        return convertedHead.getConverter().apply(combined);
      });
    }
  }

  public void validate(Consumer<Problem> problemConsumer) {
    for (Reference<CanonicalType> ref : getReferences()) {
      if (ref.getResult().hasProblems()) {
        problemConsumer.accept(
            GeneralProblems.get().failedToValidate(IdentifiedType.class, ref.getId(), ref.getResource())
                .withSeverity(Problem.max(ref.getResult().getProblems()))
                .withChildren(ref.getResult().getProblems())
        );
        // Now that any warnings have gone to project problems we can drain them from the reference
        // to reduce runtime noise
        ref.drainWarnings(p -> {});
      }
    }
  }

  /**
   * Implements TypeRules for this type set and keeps a track of which rule is being tested in a nested way to improve
   * logging readability
   */
  @RequiredArgsConstructor
  private class WrappedTypeRules implements TypeRules {

    private final int indentLevel;

    private final TypeSet wrapped;

    private final Identified rule;

    private String prefix;

    @Override
    public Variance testVariance(Type sourceType, Type targetType) {
      List<VarianceRule> varianceRules = wrapped.getTypeRegistry().getVarianceRules();
      if (rule != null) {
        debug("Testing variance of {} => {}...", sourceType, targetType);
      }

      for (VarianceRule varianceRule : varianceRules) {
        TypeRules next = new WrappedTypeRules(indentLevel + 1, wrapped, varianceRule);

        Variance variance = varianceRule.test(next, sourceType, targetType);
        if (variance != Variance.UNKNOWN) {
          debug("{} variance returned by rule '{}' for {} to {}",
              variance.name(), varianceRule.getId(), sourceType, targetType);
          return variance;
        }
      }

      debug("No rule matched.");
      return Variance.UNKNOWN;
    }

    /**
     * Search through the list of {@link EquivalenceRule}s and find the first one that can safely coerce a value from
     * the source type to the target type.
     * @param sourceType the theoretical l.h.s of an assignment
     * @param targetType the theoretical r.h.s of an assignment
     * @return a {@link Coercer}, or empty if none could be found.
     */
    @Override
    public Optional<Coercer> findEquivalenceCoercer(Type sourceType, Type targetType) {
      List<EquivalenceRule> equivalenceRules = wrapped.getTypeRegistry().getEquivalenceRules();

      if (rule != null) {
        debug("Testing equivalence of {} => {}...", sourceType, targetType);
      }
      if (sourceType == targetType) {
        debug("Exact same types - returning identity coercer for {}", targetType);
        return Optional.of(Coercer.identity(targetType));
      }

      for (EquivalenceRule equivalenceRule : equivalenceRules) {
        TypeRules next = new WrappedTypeRules(indentLevel + 1, wrapped, equivalenceRule);
        // TODO wrap this to detect and report on recursive look ups
        Optional<Coercer> coercer = equivalenceRule.getCoercer(next, sourceType, targetType);
        if (coercer.isPresent()) {
          debug("Equivalence rule '{}' matched for {} => {}", equivalenceRule.getId(), sourceType, targetType);
          return coercer;
        }
      }

      debug("No rule matched");
      return Optional.empty();
    }

    @Override
    public Optional<AncestorType> computeAncestorType(Type type1, Type type2) {
      List<AncestorRule> ancestorRules = wrapped.getTypeRegistry().getAncestorRules();

      if (rule != null) {
        debug("Testing ancestory of {} => {}...", type1, type2);
      }
      if (type1.equals(type2)) {
        debug("Same types - returning identity");
        return Optional.of(AncestorType.of(type1));
      }

      for (AncestorRule ancestorRule : ancestorRules) {
        TypeRules next = new WrappedTypeRules(indentLevel + 1, wrapped, ancestorRule);
        // TODO wrap this to detect and report on recursive look ups
        Optional<AncestorType> ancestor = ancestorRule.computeAncestor(next, type1, type2);
        if (ancestor.isPresent()) {
          debug("Ancestor rule '{}' matched for {} => {}", ancestorRule.getId(), type1, type2);
          return ancestor;
        }
      }

      debug("No rule matched");
      return Optional.empty();
    }

    @Override
    public void debug(String msg, Object... args) {
      if (log.isDebugEnabled()) {
        if (prefix == null) {
          String indent = Strings.repeat("  ", indentLevel);
          if (rule == null) {
            prefix = "";
          } else {
            prefix = indent + " [in " + rule.getId() + "] ";
          }
        }

        log.debug(prefix + msg, args);
      }
    }
  }

  @Override
  public void debug(String msg, Object... args) {
    log.debug(msg, args);
  }


}
