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

import java.util.List;
import java.util.Optional;
import java.util.Set;

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

import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.bind.ParameterField;
import nz.org.riskscape.engine.join.DefaultJoiner;
import nz.org.riskscape.engine.join.ExpressionJoinCondition;
import nz.org.riskscape.engine.join.HashIndexer;
import nz.org.riskscape.engine.join.JoinCondition;
import nz.org.riskscape.engine.join.JoinIndexer.Constructor;
import nz.org.riskscape.engine.join.JoinType;
import nz.org.riskscape.engine.join.Joiner;
import nz.org.riskscape.engine.join.NoIndexIndexer;
import nz.org.riskscape.engine.join.SpatialIndexer;
import nz.org.riskscape.engine.pipeline.RealizationInput;
import nz.org.riskscape.engine.pipeline.Realized;
import nz.org.riskscape.engine.pipeline.RealizedPipeline;
import nz.org.riskscape.engine.pipeline.RealizedStep;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.rl.ExpressionRealizer;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ExpressionParser;
import nz.org.riskscape.rl.ast.Expression;
import nz.org.riskscape.rl.ast.InputChain;
import nz.org.riskscape.rl.ast.PipelineDeclaration;
import nz.org.riskscape.rl.ast.StepDeclaration;
import nz.org.riskscape.rl.ast.StepReference;

public class JoinStep extends BaseStep<JoinStep.Parameters> {

  public interface LocalProblems extends ProblemFactory {

    Problem noLeftHandSide();

    Problem noRightHandSide();

    /**
     * When the RHS has been provide two sets of input, from named input and the rhs arg.
     */
    Problem twoRightHandSides();

    /**
     * Validation error when the input chain contains step references - these are a likely source of errors and
     * complexity that we can't currently see a use case for, so for now we're going to disallow them instead (unless)
     * they are the entire pipeline.
     *
     * FYI StepDeclaration will always be a StepReference, was just easier to keep the parameter generic
     */
    Problem stepReferencesNotAllowed(StepDeclaration step);
  }

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

  /**
   * {@link Constructor}s in priority order.
   */
  public static final List<Constructor> INDEX_CONSTRUCTORS = Lists.newArrayList(
      SpatialIndexer.CONSTRUCTOR,
      HashIndexer.CONSTRUCTOR,
      NoIndexIndexer.CONSTRUCTOR
  );

  public static class Parameters {

    @ParameterField
    public Expression on = ExpressionParser.parseString("true");

    @ParameterField
    public JoinType joinType = JoinType.INNER;

    @ParameterField
    public Optional<Integer> initialIndexSize;

    @ParameterField
    public Optional<InputChain> rhs;

    public RealizedStep lhsStep;

    public RealizedStep rhsStep;

    //Injected field
    public RealizationInput rInput;

  }

  public JoinStep(Engine engine) {
    super(engine);
  }

  @Override
  protected ResultOrProblems<Parameters> bindInputs(RealizationInput input, Parameters parameters) {
    List<RealizedStep> inputs = input.getDependencies();

    if (!getInputArity().contains(inputs.size())) {
      return ResultOrProblems.failed(
          GeneralProblems.get().badArity(Input.class, this, getInputArity(), inputs.size()));
    }

    // first go by name
    parameters.lhsStep = input.getNamedDependencies().get("lhs");
    parameters.rhsStep = input.getNamedDependencies().get("rhs");

    // then fall back to order
    if (parameters.lhsStep == null && !inputs.get(0).equals(parameters.rhsStep)) {
      // we don't assign the lhs step from the first input in the case that it is already in use
      // as the rhsStep from a named input. because if only the rhs named input has been used it will
      // also be the first (and only) input
      parameters.lhsStep = inputs.get(0);
    }

    if (parameters.rhsStep == null && inputs.size() == 2) {
      parameters.rhsStep = inputs.get(1);
    }

    if (parameters.lhsStep == null) {
      return ResultOrProblems.failed(PROBLEMS.noLeftHandSide());
    }

    return ResultOrProblems.of(parameters);
  }

  @Override
  protected ResultOrProblems<Parameters> bindParameters(RealizationInput input) {
    return super.bindParameters(input).flatMap(parameters -> {

      // validate that there are no step references is in our input chain - see Problems.stepReferencesNotAllowed
      // for more info
      if (parameters.rhs.isPresent()) {
        InputChain inputChain = parameters.rhs.get();

        if (inputChain.getAsStepReference().isEmpty()) {
          for (StepDeclaration decl : inputChain.getChain().getSteps()) {
            if (decl.isA(StepReference.class).isPresent()) {
              // found one, failed
              return ResultOrProblems.failed(PROBLEMS.stepReferencesNotAllowed(decl));
            }
          }
        }
      }

      return ResultOrProblems.of(parameters);
    });
  }

  @Override
  public List<String> getInputNames() {
    return List.of("lhs", "rhs");
  }

  @Override
  public Range<Integer> getInputArity() {
    return Range.closed(1, 2);
  }

  @Override
  public RealizedPipeline realize(RealizationInput input) {
    RealizedPipeline startFrom = input.getRealizedPipeline();

    ResultOrProblems<Parameters> built = buildParametersObject(input);
    if (built.hasProblems()) {
      RealizedStep newStep = input.newPrototypeStep().realizedBy(this).withResult(
          ResultOrProblems.failed(built.getProblems())
      );
      return startFrom.add(newStep);
    }
    Parameters parameters = built.get();

    InputChain rhs = parameters.rhs.orElse(null);

    if (rhs != null) {
      // we have two ways to provide the RHS of the join.
      // but only one of them should be used at any given time
      if (parameters.rhsStep != null) {
        RealizedStep newStep = input.newPrototypeStep().realizedBy(this).withResult(
            ResultOrProblems.failed(PROBLEMS.twoRightHandSides())
        );
        return startFrom.add(newStep);
      }

      StepReference rhsReference = rhs.getAsStepReference().orElse(null);
      if (rhsReference != null) {
        // no need to build a child pipeline, just set rhsStep to it
        parameters.rhsStep = startFrom.getStep(rhsReference.getIdent()).get();
      } else {
        PipelineDeclaration rhsAsPipeline = new PipelineDeclaration(List.of(rhs.getChain()));

        ChildPipelineBuilder builder = new ChildPipelineBuilder(input, rhsAsPipeline);
        Set<RealizedStep> childEndSteps = builder.realizeChild().getEndSteps();

        startFrom = builder.combine();

        if (startFrom.hasFailures()) {
          return startFrom;
        }

        assert childEndSteps.size() == 1; // has to be this - we built it from a single chain
        parameters.rhsStep = builder.getRebuiltStep(childEndSteps.iterator().next());
      }
    }

    // we need to ensure that we have a RHS
    if (parameters.rhsStep == null) {
      RealizedStep newStep = input.newPrototypeStep().realizedBy(this).withResult(
          ResultOrProblems.failed(PROBLEMS.noRightHandSide())
      );

      return startFrom.add(newStep);
    }

    ExpressionRealizer expressionRealizer = parameters.rInput.getExecutionContext().getExpressionRealizer();

    //Seed the join with a cartesian join condition. This is to get the shape of the joined data
    ResultOrProblems<DefaultJoiner> joinerPr = Joiner.build(
        parameters.lhsStep.getProduces(),
        parameters.rhsStep.getProduces(),
        JoinCondition.CARTESIAN,
        parameters.joinType).flatMap(join -> {

      //now realize the expression join condition for the joined data
      return ExpressionJoinCondition.realize(parameters.on, join.getJoinedType(), expressionRealizer)
          .flatMap(on -> {
            return ResultOrProblems.of(
                //and update join with the realized expression join condition
                new DefaultJoiner(join.withJoinConditon(on), INDEX_CONSTRUCTORS, expressionRealizer,
                    parameters.initialIndexSize));
          });
      });

    return startFrom.add(
        input
          .newPrototypeStep()
          .withResult(joinerPr)
          .withDependencies(parameters.lhsStep, parameters.rhsStep)
      );

  }

  @Override
  public ResultOrProblems<? extends Realized> realize(Parameters parameters) {
    throw new UnsupportedOperationException("join step does all the realizing itself");
  }

}
