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

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import com.google.common.collect.Range;

import lombok.Data;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.defaults.lookup.MapLookupTable;
import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.FunctionArgument;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.rl.agg.Accumulator;
import nz.org.riskscape.engine.rl.agg.AggregationFunction;
import nz.org.riskscape.engine.rl.agg.RealizedAggregateExpression;
import nz.org.riskscape.engine.types.LambdaType;
import nz.org.riskscape.engine.types.LookupTableType;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.RSList;
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.TypeProblems;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.engine.util.FunctionCallOptions;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.FunctionCall;

public class ToLookupTable implements AggregationFunction {

  @Data
  public static class Options {
    @ParameterField
    private boolean unique = false;
  }

  public interface LocalProblems extends ProblemFactory {
    Problem keyNotUnique(Object key, List<Object> values);
  }

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

  @RequiredArgsConstructor
  public static class LookupTableAccumulator implements Accumulator {

    private final RealizedExpression keyExpression;
    private final RealizedExpression valueExpression;
    private final boolean unique;

    @SuppressWarnings("rawtypes")
    private final HashMap values = new HashMap<>();

    @Override
    public Accumulator combine(Accumulator rhs) {
      LookupTableAccumulator ltRhs = (LookupTableAccumulator) rhs;

      // TODO do something clever with a combined map view of the many individual maps?
      HashMap<Object, Object> rhsValues = ltRhs.values;
      for (Map.Entry entry : rhsValues.entrySet()) {
        if (unique) {
          Object replaced = values.put(entry.getKey(), entry.getValue());
          if (replaced != null) {
            throw new RiskscapeException(
                PROBLEMS.keyNotUnique(entry.getKey(), Arrays.asList(entry.getValue(), replaced)));
          }
        } else {
          values.merge(entry.getKey(), entry.getValue(), (old, andNew) -> {
            List<Object> oldList = (List<Object>) old;
            List<Object> newList = (List<Object>) andNew;

            oldList.addAll(newList);
            return oldList;
          });
        }
      }

      return this;
    }

    @Override
    public void accumulate(Object input) {
      Object key = keyExpression.evaluate(input);
      Object value  = valueExpression.evaluate(input);

      if (unique) {
        @SuppressWarnings("unchecked")
        Object replaced = values.put(key, value);
        if (replaced != null) {
          throw new RiskscapeException(PROBLEMS.keyNotUnique(key, Arrays.asList(value, replaced)));
        }
      } else {
        @SuppressWarnings("unchecked")
        HashMap<Object, List<Object>> listValues = values;
        List<Object> keyedValues = listValues.computeIfAbsent(key, k -> new ArrayList<>());
        keyedValues.add(value);
      }

    }

    @Override
    public Object process() {
      return new MapLookupTable(values, keyExpression.getResultType(), valueExpression.getResultType());
    }

    @Override
    public boolean isEmpty() {
      return values.isEmpty();
    }
  }

  private class FunctionImpl implements RiskscapeFunction, RealizableFunction {

    private final RSList expectedType = RSList.create(Struct.of("key", Types.ANYTHING, "value", Types.ANYTHING));

    @Getter
    private final ArgumentList arguments = ArgumentList.fromArray(
        new FunctionArgument("list", expectedType),
        new FunctionArgument("options", Struct.of("unique", Types.BOOLEAN))
    );

    @Override
    public Type getReturnType() {
      return LookupTableType.WILD;
    }

    @Override
    public List<Type> getArgumentTypes() {
      return arguments.getArgumentTypes();
    }

    @Override
    public Object call(List<Object> args) {
      throw new UnsupportedOperationException();
    }

    @Override
    public boolean isDoTypeAdaptation() {
      // nope - otherwise we won't see the real types of key and value
      return false;
    }

    @Override
    public ResultOrProblems<RiskscapeFunction> realize(RealizationContext context, FunctionCall functionCall,
        List<Type> argumentTypes) {

      if (argumentTypes.size() < 1) {
        return ResultOrProblems.failed(ArgsProblems.get().wrongNumber(Range.closed(1, 2), arguments.size()));
      }

      return ProblemException.catching(() -> {
        ProblemException listMismatch = new ProblemException(
            TypeProblems.get().mismatch(functionCall.getArguments().get(0), expectedType, argumentTypes.get(0))
        );

        // we need to track if the list input is nullable because arguments.getRequiredAs strips nulls
        // behind the scenes.
        boolean nullable = Nullable.is(argumentTypes.get(0));
        RSList list = arguments.getRequiredAs(argumentTypes, 0, RSList.class).orElseThrow(() -> listMismatch);

        Struct memberType = list.getMemberType()
          .findAllowNull(Struct.class)
          .orElseThrow(() -> listMismatch);

        Options options = FunctionCallOptions.bindOptions(Options.class, context, arguments, functionCall).getOrThrow();

        StructMember keyMember = memberType.getMember("key").orElseThrow(() -> listMismatch);
        StructMember valueMember = memberType.getMember("value").orElseThrow(() -> listMismatch);

        Type valueType;
        if (options.unique) {
          valueType = valueMember.getType();
        } else {
          valueType = RSList.create(valueMember.getType());
        }

        LookupTableType returnType = LookupTableType.create(keyMember.getType(), valueType);
        return RiskscapeFunction.create(this, argumentTypes, Nullable.ifTrue(nullable, returnType), args -> {

          @SuppressWarnings("rawtypes")
          HashMap values = new HashMap<>();

          @SuppressWarnings({ "rawtypes" })
          List<?> toReduce = (List<?>) args.get(0);

          if (toReduce == null) {
            return null;
          }

          for (Object object : toReduce) {

            // we support null elements - they get skipped
            if (object == null) {
              continue;
            }

            Tuple tuple = (Tuple) object;

            Object key = tuple.fetch(keyMember);
            Object value = tuple.fetch(valueMember);

            if (options.unique) {
              @SuppressWarnings("unchecked")
              Object replaced = values.put(key, value);
              if (replaced != null) {
                throw new RiskscapeException(PROBLEMS.keyNotUnique(key, Arrays.asList(value, replaced)));
              }
            } else {
              @SuppressWarnings("unchecked")
              HashMap<Object, List<Object>> listValues = values;
              List<Object> keyedValues = listValues.computeIfAbsent(key, k -> new ArrayList<>());
              keyedValues.add(value);
            }
          }

          return new MapLookupTable(values, returnType.getKeyType(), valueType);
        });
      });
    }

    @Override
    public Optional<AggregationFunction> getAggregationFunction() {
      return Optional.of(ToLookupTable.this);
    }
  };

  private final FunctionImpl function = new FunctionImpl();

  // TODO have some methods on ArgumentList for doing sublist with appropriate cloning
  private final ArgumentList arguments = ArgumentList.create(
      "keyExpression", new LambdaType("listElement"),
      "valueExpression", new LambdaType("listElement"),
      "options", Struct.of("unique", Types.BOOLEAN)
  );
  @Override
  public ResultOrProblems<RealizedAggregateExpression> realize(
      RealizationContext context,
      Type inputType,
      FunctionCall fc
  ) {

    return ProblemException.catching(() -> {
      // check arity
      if (fc.getArguments().size() < 2) {
        throw new ProblemException(ArgsProblems.get().wrongNumber(2, fc.getArguments().size()));
      }

      Options options = FunctionCallOptions.bindOptions(Options.class, context, arguments, fc).getOrThrow();

      RealizedExpression keyExpression = context.getExpressionRealizer()
          .realize(inputType, fc.getArguments().get(0).getExpression())
          .getOrThrow(Problems.foundWith(fc.getArguments().get(0)));

      RealizedExpression valueExpression = context.getExpressionRealizer()
          .realize(inputType, fc.getArguments().get(1).getExpression())
          .getOrThrow(Problems.foundWith(fc.getArguments().get(1)));

      Type lookupValueType =
          options.isUnique() ? valueExpression.getResultType() : RSList.create(valueExpression.getResultType());

      Type lookupTableType = new LookupTableType(keyExpression.getResultType(), lookupValueType);

      return RealizedAggregateExpression.create(inputType, lookupTableType, fc,
          () -> new LookupTableAccumulator(keyExpression, valueExpression, options.isUnique()));
    });
  }

  public RiskscapeFunction asRiskscapeFunction() {
    return function;
  }

}
