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

import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricSet;
import com.google.common.collect.Maps;

import lombok.EqualsAndHashCode;
import lombok.Getter;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.output.PipelineJobContext;
import nz.org.riskscape.engine.output.PipelineOutputContainer;
import nz.org.riskscape.engine.pipeline.ExecutionContext;
import nz.org.riskscape.engine.pipeline.ExecutionOptions;
import nz.org.riskscape.engine.pipeline.RealizedStep;
import nz.org.riskscape.engine.sched.Scheduler;
import nz.org.riskscape.problem.ProblemException;

/**
 * The TaskSpec defines the work for a step or series of steps in a pipeline.
 * This is a just blueprint for the work that needs to be done. One or more
 * WorkerTasks are then created for the TaskSpec, which will actually carry out
 * the work.
 */
@EqualsAndHashCode(of = {"impl", "name"})
public final class TaskSpec implements MetricSet {

  /**
   * Constant for {@link #getProcessingResultsFrom(Class, Class, int)}
   */
  public static final int NO_MINIMUM = 0;

  private final Class<? extends WorkerTask> impl;

  /**
   * The underlying pipeline step(s) that this task is executing
   */
  @Getter
  private final List<RealizedStep> forSteps;

  @Getter
  private final List<TaskSpec> dependsOn = new ArrayList<>();
  private final List<TaskSpec> dependsOnSatisified = new ArrayList<>();
  private final Map<TaskSpec, List<Object>> dependencyResults = Maps.newHashMap();


  private final ReadPageBuffer input;
  private final WritePageBuffer output;

  @Getter
  private final boolean parallelizable;

  @Getter
  private final PipelineJobContext jobContext;

  private final String name;

  @Getter
  private final Map<String, Metric> metrics = new HashMap<>();

  @Getter
  private final List<WorkerTask> workerTasks = new ArrayList<>();

  /**
   * A map of {@link Metric}s that give an indication of progress through this task. {@link Metric}s from this map
   * will be added to {@link ExecutionOptions#getProgressMetrics()} whilst the task is being processed.
   */
  @Getter
  protected final Map<String, Metric> progressMetrics = new HashMap<>();

  @Getter
  private TaskState state = TaskState.CREATED;

  public TaskSpec(Class<? extends WorkerTask> impl, List<RealizedStep> forSteps, ReadPageBuffer input,
      WritePageBuffer output, boolean parallelizable, PipelineJobContext context) {
    this.impl = impl;
    this.forSteps = forSteps;
    this.input = input;
    this.output = output;
    this.parallelizable = parallelizable;
    this.jobContext = context;

    // a unique name for the spec (used for equality checks)
    this.name = impl.getSimpleName() + "+" + getCombinedStepsName();
  }


  public Optional<ReadPageBuffer> getInput() {
    return Optional.ofNullable(input);
  }

  public Optional<WritePageBuffer> getOutput() {
    return Optional.ofNullable(output);
  }

  /**
   * @return the subclass of WorkerTask that is used to fulfill this taskspec's processing.  One or many of these
   * classes are instantiated, depending on the type of task, to consume/produce/both tuples as part of execution.
   */
  public Class<? extends WorkerTask> getWorkerTaskClass() {
    return impl;
  }

  /**
   * Adds a hard dependency that requires another task to have run to completion, while allowing arbitrary objects
   * (called processing results) to be passed to dependents before they themselves execute.
   *
   * {@link WorkerTask}s can return 'processing results' when they are complete by setting
   * {@link WorkerTask#processingResult}.  The {@link Scheduler} picks these up once a WorkerTask is complete by calling
   * {@link WorkerTask#consumeProcessingResult()}.  This result is then passed to the dependencies by calling
   * {@link #addProcessingResultFromDependency(TaskSpec, Object)}.  This is all done from the scheduler thread on tasks
   * that are either waiting to run or have finished running, so no locking is required.
   *
   * When a {@link WorkerTask} that has dependencies runs, it can access the results of their dependencies by calling
   * {@link #getProcessingResultsFrom(Class, Class)} or {@link #getProcessingResultFrom(Class, Class)}..
   */
  public void addDependency(TaskSpec task) {
    dependsOn.add(task);
  }
  /**
   * @return true if this task depends on the given task spec being complete
   */
  public boolean hasDependency(TaskSpec taskSpec) {
    return dependsOn.contains(taskSpec);
  }

  /**
   * Mark a given dependent task as completed.  This should be called once any and all {@link WorkerTask}s from that
   * dependency have run to completion and their processingResults have been distributed appropriately, as once all
   * dependencies have been satisfied a task may start running.  If it doesn't have all the results it needs, the
   * results of execution are likely to be incomplete.
   */
  public void satisfyDependency(TaskSpec dependency) {
    if (dependsOn.remove(dependency)) {
      dependsOnSatisified.add(dependency);
    }
  }

  /**
   * @return true if this tasks has dependencies that are preventing it from running.
   */
  public boolean hasOutstandingDependencies() {
    return !dependsOn.isEmpty();
  }

  /**
   * Pass the result of execution from a dependency to this task, so that a (or many) future {@link WorkerTask}s can
   * use the results as part of their execution.  {@link WorkerTask}s should use
   * {@link #getProcessingResultsFrom(Class, Class)} or {@link #getProcessingResultFrom(Class, Class)} to access
   * these results within the call to {@link WorkerTask#run()}.
   */
  public void addProcessingResultFromDependency(TaskSpec fromSpec, Object processingResult) {
    // note that there's no checking that a dependency exists - we're assuming the scheduler will behave, which is fair
    // as it's fairly tightly bound with this implementation
    dependencyResults.computeIfAbsent(fromSpec, s -> new ArrayList<>()).add(processingResult);
  }

  /**
   * Shortcut for `getProcessingResultsFrom(taskType, expectedType, NO_MINIMUM)`
   */
  public <T> List<T> getProcessingResultsFrom(Class<? extends WorkerTask> taskType, Class<T> expectedType) {
    return getProcessingResultsFrom(taskType, expectedType, NO_MINIMUM);
  }

  /**
   * Returns all of the processing results from dependencies of the given task type, or an empty list if there is no
   * dependency of that type, or no results of that type.
   * @param taskType the type of the dependency we're expecting there to be.  No checking is done that there ever was a
   * dependency of this type at this point, but the idea is that future use cases might depend on more than one task
   * type and so this is here to distinguish the dependencies.
   * @param expectedType The expected type of processing results from the given task type.
   * @param min the minimum number of results expected
   * @return A list of the processing results, or an empty list of no results were found
   * @throws ClassCastException if there are results but they are not of `expectedType`
   * @throws IllegalStateException if there are less results than the specified minimum
   */
  public <T> List<T> getProcessingResultsFrom(Class<? extends WorkerTask> taskType, Class<T> expectedType, int min) {
    List<T> results = new ArrayList<>();
    dependencyResults.entrySet().forEach(entry -> {
      if (taskType.isAssignableFrom(entry.getKey().impl)) {
        List<?> values = entry.getValue();
        if (values.size() > 0) {
          for (Object value : values) {
            results.add(expectedType.cast(value));
          }
        }
      }
    });

    if (results.size() < min) {
      throw new IllegalStateException(
          "Expected to have at least " + min + " dependency(ies) from task " + taskType + " but had " + results.size()
        );
    }

    return results;
  }

  /**
   * A version of {@link #getProcessingResultsFrom(Class, Class)} that ensures that one and only one processing result
   * exists.
   * @throws IllegalStateException if there isn't exactly one result.  This will halt execution and is considered a
   * programming error, not a user error
   */
  public <T> T getProcessingResultFrom(Class<? extends WorkerTask> taskType, Class<T> expectedType) {
    List<T> results = getProcessingResultsFrom(taskType, expectedType, -1);

    if (results.size() == 1) {
      return results.get(0);
    } else {
      throw new IllegalStateException(
        "Expected to have a single dependency from task " + taskType + " but had " + results.size()
      );
    }
  }

  /**
   * Creates a new WorkerTask for the given Scheduler. The Scheduler needs to
   * translate a TaskSpec into one or more WorkerTasks that will actually carry
   * out the work.
   */
  public WorkerTask newWorkerTask(Scheduler scheduler) throws ProblemException {
    // while the TaskSpec is final, the corresponding WorkerTask will always
    // be a specialized sub-class. So we need to use reflection here to get
    // the appropriate constructor for the sub-class. (We construct the WorkerTasks
    // this way so that the Scheduler can potentially scale the number of worker
    // tasks up or down, if needed)
    Constructor<? extends WorkerTask> taskClass;
    try {
      taskClass = impl.getConstructor(TaskSpec.class);
    } catch (NoSuchMethodException | SecurityException e) {
      throw new RiskscapeException(
          String.format("WorkerTask %s is missing an appropriate constructor", impl.getSimpleName()), e);
    }

    WorkerTask newTask;
    try {
      newTask = taskClass.newInstance(this);
    } catch (InstantiationException | IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
      if (e.getCause() instanceof ProblemException) {
        throw (ProblemException) e.getCause();
      } else if (e.getCause() instanceof RiskscapeException) {
        // RiskscapeException may contain a problem which would be more meaningful to users. So we look
        // for a problem, and if found throw it in a ProblemException
        RiskscapeException rsException = (RiskscapeException)e.getCause();
        if (rsException.hasProblem()) {
          throw new ProblemException(rsException.getProblem());
        }
      }
      // If no problems are found we fall back to throwing a RiskscapeException
      throw new RiskscapeException(String.format("Failed to create WorkerTask for TaskSpec %s", toString()), e);
    }

    workerTasks.add(newTask);
    return newTask;
  }

  protected <T extends Metric> T newMetric(String metricName, Supplier<T> constructor) {
    return newMetric(metricName, null, constructor);
  }
  /**
   * Create or return an existing Metric with the given task-local name.
   *
   * Creating a metric through this method will register the metric against the execution contexts metrics with a name
   * that identifies it uniquely against the
   * executed pipeline (using the task spec's steps to generate an id).  For parallelizable tasks, you may be returned
   * a pre-existing metric - this allows the same metric to record activity across a set of tasks related to the same
   * {@link TaskSpec}
   *
   * All calls to newMetric should be done from the scheduler thread, not from the worker thread

   * @param progressName a possibly alternative name to use to register this metrics as a progress metric.  See
   * {@link #getProgressMetrics()} for more info.  Set to null if this isn't an important progress indicator metric.
   *
   */
  @SuppressWarnings("unchecked")
  protected <T extends Metric> T newMetric(String metricName, String progressName, Supplier<T> constructor) {
//    String fullMetricName = getCombinedStepsName() + "." + metricName;

    if (metrics.containsKey(metricName)) {
      // this is going to throw ClassCastException if used improperly, don't think it's worth a specific
      // check/exception
      return (T) metrics.get(metricName);
    } else {
      T constructed = constructor.get();
      metrics.put(metricName, constructed);

      if (progressName != null) {
        progressMetrics.put(progressName, constructed);
      }
      return constructed;
    }
  }

  /**
   * @return a combined string containing of all the step names that this task handles
   */
  public String getCombinedStepsName() {
    return forSteps.stream()
        .map(RealizedStep::getStepName)
        .collect(Collectors.joining("+"));
  }

  public String getStepsSummary() {
    return forSteps.stream().map(rs -> rs.getStepName()).collect(Collectors.joining(", "));
  }

  @Override
  public String toString() {
    return String.format("%s(%s)", impl.getSimpleName(), getStepsSummary());
  }

  public RealizedStep getLastStep() {
    return forSteps.get(forSteps.size() - 1);
  }

  public RealizedStep getFirstStep() {
    return forSteps.get(0);
  }

  public void changeState(TaskState newState) {
    if (newState.ordinal() > state.ordinal()) {
      this.state = newState;
    } else {
      throw new IllegalStateException("can not change from " + state + " to " + newState);
    }
  }

  public boolean isCreated() {
    return state == TaskState.CREATED;
  }

  public boolean isStarted() {
    return state == TaskState.STARTED;
  }

  public boolean isComplete() {
    return state == TaskState.COMPLETE;
  }

  public void close() {
    forSteps.forEach(RealizedStep::close);
    // it's safe to clear the input from the previous task now, so that it can be GC'd
    dependencyResults.clear();
  }

  public boolean hadDependency(TaskSpec taskSpec) {
    return dependsOnSatisified.contains(taskSpec);
  }

  public boolean allWorkersMatch(Predicate<WorkerTask> predicate) {
    return workerTasks.stream().allMatch(predicate);
  }

  public ExecutionContext getExecutionContext() {
    return jobContext.getExecutionContext();
  }

  public PipelineOutputContainer getContainer() {
    return jobContext.getOutputContainer();
  }
}
