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

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.ListIterator;
import java.util.Optional;
import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

import com.google.common.collect.Lists;

import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.output.PipelineOutputContainer;
import nz.org.riskscape.pipeline.PipelineMetadata;
import nz.org.riskscape.rl.ast.PipelineDeclaration;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problems;

/**
 * A RealizedPipeline is the result of realizing (going from untyped to typed) a {@link PipelineDeclaration}
 * against a project and a set of parameters via a {@link PipelineRealizer}.  The structure of the pipeline is
 * represented here as directed graph of {@link RealizedStep}s, where each step represents a typed transformation of
 * data from one step to the next.
 *
 * A {@link RealizedPipeline} can be given to a {@link PipelineExecutor} to execute the pipeline, writing results to a
 * {@link PipelineOutputContainer}.
 *
 * # Closing attached resources
 *
 * Closing a RealizedPipeline will close all of its realized steps.  One slightly awkward side effect of using an
 * immutable-prototype pattern for building a realized pipeline is that it muddies the water for who is responsible for
 * closing it.
 *
 * My feeling is that as long as realized steps lazily allocate resources, that is, they only take them once execution
 * starts, then this won't be a problem in practice: any orphaned steps won't need to be closed; only a pipeline and its
 * steps that get executed will need to be closed.
 *
 * If this starts to become a source of bugs, we could:
 *
 * * Split this role up and have a RealizedPipelineBuilder and totally immutable RealizedPipeline.  At this point only
 * the RealizedPipeline is executable.
 *
 * * Remove the {@link AutoCloseable} from RealizedPipeline and {@link RealizedStep} and consider changing the
 * `getResult()` method to be more like a constructor, e.g. `newRealized(ExecutionContext)` and have that be the
 * closable
 */
@RequiredArgsConstructor
@EqualsAndHashCode
public class RealizedPipeline implements AutoCloseable {

  /**
   * @return a pipeline with no steps.
   */
  public static RealizedPipeline empty(ExecutionContext context, PipelineDeclaration pipeline) {
    return new RealizedPipeline(context, pipeline, List.of());
  }

  @Getter
  private final ExecutionContext context;

  /**
   * The source code that was realized to produce this realized pipeline
   */
  @Getter
  private final PipelineDeclaration ast;

  /**
   * The individually realized steps, in the order they were realized
   */
  @Getter
  private final List<RealizedStep> realizedSteps;

  /**
   * A set of problems associated with realizing the pipeline that weren't specific to a single step
   */
  private final List<Problem> problems;

  private final PipelineMetadata metadata;

  public RealizedPipeline(ExecutionContext context, PipelineDeclaration pipeline, List<RealizedStep> steps) {
    this(context, pipeline, steps,  Collections.emptyList(), pipeline.getMetadata());
  }

  public RealizedPipeline(ExecutionContext context, PipelineDeclaration pipeline, List<RealizedStep> steps,
      List<Problem> problems) {
    this(context, pipeline, steps, problems, pipeline.getMetadata());
  }

  @Override
  public void close() {
    List<Exception> failures = Lists.newArrayListWithExpectedSize(0);
    for (RealizedStep step : realizedSteps) {
      try {
        step.close();
      } catch (Exception e) {
        failures.add(e);
      }
    }

    if (!failures.isEmpty()) {
      throw new AutoClosingException(this, failures);
    }
  }

  /**
   * @return the set of steps that have no dependencies
   */
  public Set<RealizedStep> getStartSteps() {
    return realizedSteps.stream().filter(RealizedStep::hasNoDependencies).collect(Collectors.toSet());
  }

  /**
   * @return the set of {@link RealizedStep}s that no other steps depend on
   */
  public Set<RealizedStep> getEndSteps() {
    Set<RealizedStep> set = new HashSet<>(realizedSteps);

    // If a step is a dependency, then it is not an end step - we clear these out from the set of all sets
    for (RealizedStep realizedStep : realizedSteps) {
      realizedStep.getDependencies().forEach(set::remove);
    }

    return set;
  }

  /**
   * @return a Set of {@link RealizedStep}s that have a dependency on the given step
   */
  public Set<RealizedStep> getDependents(RealizedStep step) {
    Set<RealizedStep> dependents = new HashSet<>();

    for (RealizedStep realizedStep : realizedSteps) {
      if (realizedStep.getDependencies().contains(step)) {
        dependents.add(realizedStep);
      }
    }

    return dependents;
  }

  /**
   * @return a list of Problems that directly caused this pipeline to fail
   */
  public List<Problem> getFailures() {
    List<Problem> failures = Lists.newArrayList();

    this.problems.stream().filter(Problem::isError).forEach(failures::add);

    for (RealizedStep step : realizedSteps) {
      // Only include failures from steps that are directly failed. Error output can become over bearing
      // pretty quickly when you have to filter out of the failed deps errors to.
      if (step.isDirectlyFailed()) {
        step.getFailureProblem().ifPresent(p -> failures.add(p));
      }
    }
    return failures;
  }

  /**
   * @return true if this realized pipeline has failed
   */
  public boolean hasFailures() {
    return !getFailures().isEmpty();
  }

  /**
   * @return a step by its name within the pipeline (c.f. NamedStep), or `empty` if none can be found
   */
  public Optional<RealizedStep> getStep(String stepName) {
    for (RealizedStep realizedStep : realizedSteps) {
      if (realizedStep.getName().equals(stepName)) {
        return Optional.of(realizedStep);
      }
    }
    return Optional.empty();
  }

  /**
   * @return a copy of this realized pipeline, copying references to any immutable members, and cloning anything that's
   * mutable
   */
  @Override
  public RealizedPipeline clone() {
    return clone(stepToCopy -> new RealizedStep(
        stepToCopy.getName(),
        stepToCopy.getImplementation(),
        stepToCopy.getAst(),
        stepToCopy.getResult(),
        // make a copy of the parameter map
        new HashMap<>(stepToCopy.getBoundParameters()),
        // copy the original steps to start with, but we'll need to 'sub them out' with the new cloned steps
        new ArrayList<>(stepToCopy.getDependencies()),
        // types should also immutable
        stepToCopy.getProduces()),
        new ArrayList<>(this.problems)
      );
  }

  private RealizedPipeline clone(Function<RealizedStep, RealizedStep> stepCloner, List<Problem> newProblems) {
    List<RealizedStep> clonedSteps = new ArrayList<>(realizedSteps.size());

    // go through each step in original order (which was possibly traversal order?)
    for (RealizedStep toCopy : realizedSteps) {
      clonedSteps.add(stepCloner.apply(toCopy));
    }

    // now fill in dependencies and dependents on our clones
    for (RealizedStep clonedStep : clonedSteps) {

      // replace the original dependencies with the clones
      ListIterator<RealizedStep> dependencies = clonedStep.getDependencies().listIterator();
      while (dependencies.hasNext()) {
        RealizedStep originalDependency = dependencies.next();
        // because we traverse and build in the same order, the clone is in the same position as the original - no need
        // to search twice
        int index = realizedSteps.indexOf(originalDependency);
        RealizedStep clonedDependency = clonedSteps.get(index);
        dependencies.set(clonedDependency);
      }
    }

    return new RealizedPipeline(
        context,
        ast,
        clonedSteps,
        newProblems,
        metadata
    );
  }

  /**
   * @return the RealizedPipeline with any warnings removed and passed to the problemConsumer
   */
  public RealizedPipeline drainWarnings(Consumer<Problem> problemConsumer) {
    if (hasFailures()) {
      // pipeline has errors, so don't bother removing any warnings
      return this;
    } else {

      this.problems.stream().filter(p -> !p.isError()).forEach(problemConsumer::accept);

      return clone(stepToCopy -> new RealizedStep(
          stepToCopy.getName(),
          stepToCopy.getImplementation(),
          stepToCopy.getAst(),
          stepToCopy.getResult().drainWarnings(problemConsumer,
              // group the problems so it's easy to see what step they relate to
              (severity, ps) -> Problems.foundWith(stepToCopy.getAst(), ps)),
          new HashMap<>(stepToCopy.getBoundParameters()),
          new ArrayList<>(stepToCopy.getDependencies()),
          stepToCopy.getProduces()),
          Collections.emptyList() // NB there shouldn't be any errors, otherwise the method would have short-circuited
      );
    }
  }
  /**
   * @return metadata associated with this {@link RealizedPipeline}
   */
  public PipelineMetadata getMetadata() {
    return this.metadata;
  }

  /**
   * @return a copy of this pipeline, with different metadata to the original.
   */
  public RealizedPipeline withMetadata(final PipelineMetadata newMetadata) {
    return new RealizedPipeline(context, ast, realizedSteps, problems, newMetadata);
  }

  /**
   * @return a new {@link RealizedPipeline} that includes the extra step.
   * @throws IllegalArgumentException if the step has dependencies not in this pipeline, or if it already contains a
   * step with the same name.
   */
  public RealizedPipeline add(RealizedStep newStep) {
    // assert name unique
    for (RealizedStep existing : realizedSteps) {
      if (newStep.getName().equals(existing.getName())) {
        String errorMessage = "Can not add step %s, name '%s' is not unique, already taken by %s"
            .formatted(newStep, newStep.getName(), existing);

        throw new IllegalArgumentException(errorMessage);
      }
    }

    // assert dependencies exist - this is done by instance, not by equality
    check:
    for (RealizedStep depStep : newStep.getDependencies()) {
      for (RealizedStep existingStep : realizedSteps) {
        if (existingStep == depStep) {
          continue check;
        }
      }

      throw new IllegalArgumentException(
          "Step %s refers to a dependency not in this pipeline: %s".formatted(newStep, depStep)
      );
    }

    List<RealizedStep> newSteps = new ArrayList<>(realizedSteps.size() + 1);
    newSteps.addAll(realizedSteps);
    newSteps.add(newStep);

    return new RealizedPipeline(context, ast, newSteps, problems, metadata);
  }

  /**
   * @return a new pipeline that contains these extra problems
   */
  public RealizedPipeline addProblems(Problem... rest) {
    List<Problem> newProblems = new ArrayList<>(this.problems.size() + rest.length);
    newProblems.addAll(this.problems);
    newProblems.addAll(List.of(rest));

    return new RealizedPipeline(context, ast, realizedSteps, newProblems, metadata);
  }
}
