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

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.rl.ExpressionRealizer;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.TokenTypes;
import nz.org.riskscape.rl.ast.BaseExpressionConverter;
import nz.org.riskscape.rl.ast.BinaryOperation;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.ExpressionConverter;
import nz.org.riskscape.rl.ast.FunctionCall;
import nz.org.riskscape.rl.ast.FunctionCall.Argument;
import nz.org.riskscape.rl.ast.PropertyAccess;

/**
 * {@link JoinIndexer} implementation that builds an in-memory hash index of tuples when the
 * {@link JoinCondition} is an equality filter. Eg 'lhs.thing = rhs.thing'
 */
@Slf4j
public class HashIndexer extends RealizedExpressionJoinIndexer {

  /**
   * The default capacity that indexes will be created for.
   *
   * The HashIndexer contains a map of right hand side values to items.
   *
   * For best performance the index should be as large as the expected number of unique right hand
   * side values.
   */
  public static final int DEFAULT_INITIAL_INDEX_SIZE = 100000;

  public static final Constructor CONSTRUCTOR = (join, expressionRealizer, initialIndexSize) -> {
    return new HashIndexer(join, expressionRealizer, initialIndexSize.orElse(DEFAULT_INITIAL_INDEX_SIZE));
  };

  /**
   * An {@link ExpressionConverter} that will adjust {@link FunctionCall} and {@link BinaryOperation}
   * expressions so that they can be evaluated on a {@link Tuple} of {@link #joinTestType}.
   *
   * This is achieved by modifying the children of any {@link FunctionCall} and {@link BinaryOperation}
   * expressions such that:
   * {@link #lhsExpression} -> lhs
   * {@link #rhsExpression} -> rhs
   */
  private BaseExpressionConverter converter = new BaseExpressionConverter() {
    @Override
    public Expression visit(FunctionCall expression, Object data) {
      List<FunctionCall.Argument> convertedArgs = Lists.newArrayListWithCapacity(expression.getArguments().size());
      for (Argument arg : expression.getArguments()) {
        String asRL = arg.getExpression().toSource();
        if (asRL.equals(lhsExpression.getExpression().toSource())) {
          convertedArgs.add(new FunctionCall.Argument(PropertyAccess.of("lhs")));
        } else if (asRL.equals(rhsExpression.getExpression().toSource())) {
          convertedArgs.add(new FunctionCall.Argument(PropertyAccess.of("rhs")));
        } else {
          convertedArgs.add(arg);
        }
      }

      return new FunctionCall(expression.getIdentifier(), convertedArgs);
    }

    @Override
    public Expression visit(BinaryOperation expression, Object data) {
      //The left and right side of the binary expression may not match the left and right side of the join.
      //So we need to work that out
      boolean leftIsLeft = lhsExpression.getExpression().toSource().equals(expression.getLhs().toSource());
      return new BinaryOperation(
          leftIsLeft ? PropertyAccess.of("lhs") : PropertyAccess.of("rhs"),
          expression.getOperator(),
          leftIsLeft ? PropertyAccess.of("rhs") : PropertyAccess.of("lhs"));
    }

  };

  private final Map<Object, List<Tuple>> index;

  /**
   * True if the join condition is an equality match. In this case
   * {@link #createRhsIterator(nz.org.riskscape.engine.Tuple) } can index the map directly for lhs
   * value and short circuit other operations.
   */
  private boolean equalityJoin = false;

  /**
   * {@link Struct} that will be populated with the left and right side values to test in
   * {@link #createRhsIterator(nz.org.riskscape.engine.Tuple) }
   */
  private Struct joinTestType = null;

  /**
   * Expression to test if results should be included in {@link #createRhsIterator(nz.org.riskscape.engine.Tuple) }.
   */
  private RealizedExpression joinTest;

  public HashIndexer(Join join, ExpressionRealizer expressionRealizer, int initialIndexSize) {
    super(join, expressionRealizer);
    index = Maps.newHashMapWithExpectedSize(initialIndexSize);
    init();
  }

  private void init() {
    if (! super.isUsable()) {
      return;
    }
    if (operator != null && operator.getExpression() instanceof BinaryOperation) {
      if (((BinaryOperation) operator.getExpression()).getOperator().type == TokenTypes.EQUALS) {
        equalityJoin = true;
      }
    }
    joinTestType = Struct.of("lhs", lhsExpression.getResultType(), "rhs", rhsExpression.getResultType());

    @SuppressWarnings("unchecked")
    ResultOrProblems<RealizedExpression> convertedAndRealized = expressionRealizer.realize(
        joinTestType,
        (Expression) operator.getExpression().accept(
            converter,
            lhsExpression.getExpression().toSource()));

    if (! convertedAndRealized.hasErrors()) {
      joinTest = convertedAndRealized.get();
    }
  }

  @Override
  public void addToIndex(Tuple toCache) {
    Object rhsValue = rhsExpression.evaluate(toCache);

    synchronized (this) {
      index.computeIfAbsent(rhsValue, (k) -> Lists.newArrayList())
          .add(toCache);
    }
  }

  @Override
  public TupleIterator createRhsIterator(Tuple lhs
  ) {
    Object lhsValue = lhsExpression.evaluate(lhs);

    if (equalityJoin) {
      //The simplest case. Access the index for the lhs value and return just those results.
      List<Tuple> rhsValues = index.getOrDefault(lhsValue, Collections.emptyList());
      return TupleIterator.wrapped(rhsValues.iterator(), Optional.empty());
    }

    //Need to iterate over the index keys to find all that match condition
    List<Tuple> combined = Lists.newArrayList();
    for (Map.Entry<Object, List<Tuple>> entry : index.entrySet()) {
      Tuple toTest = Tuple.ofValues(joinTestType, lhsValue, entry.getKey());

      if (Boolean.TRUE.equals(joinTest.evaluate(toTest))) {
        combined.addAll(entry.getValue());
      }
    }
    return TupleIterator.wrapped(combined.iterator(), Optional.empty());
  }

  @Override
  public boolean isUsable() {
    return super.isUsable() && this.joinTest != null;
  }

  @Override
  public String toString() {
    return getClass().getSimpleName();
  }

}
