/*
 * 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.io.Closeable;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

import com.codahale.metrics.Counter;
import com.codahale.metrics.ExponentiallyDecayingReservoir;
import com.codahale.metrics.Histogram;
import com.codahale.metrics.Meter;

import lombok.Getter;
import lombok.NoArgsConstructor;
import nz.org.riskscape.engine.HasMeter;
import nz.org.riskscape.engine.pipeline.ExecutionContext;
import nz.org.riskscape.engine.pipeline.Realized;
import nz.org.riskscape.engine.pipeline.RealizedStep;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.sched.Scheduler;

/**
 * A WorkerTask is based off a TaskSpec and is what actually carries out the
 * work for the pipeline step(s). When a TaskSpec supports parallel processing,
 * it gets split into multiple WorkerTasks that all can run at the same time.
 * The WorkerTasks get managed by the Scheduler and get assigned to a Worker
 * thread to run.
 */
@NoArgsConstructor(force = true)
public abstract class WorkerTask implements AutoCloseable {

  @Getter
  protected final TaskSpec spec;
  protected final int id;

  protected final PageWriter pageWriter;
  protected final PageReader pageReader;

  @Getter
  private final ExecutionContext context;

  protected final Meter in;
  protected final Meter out;

  protected final Counter runtime;
  protected final Histogram runtimeAverage;
  protected final Counter contextSwitches;

  /**
   * Set this during `run` to pass the processing result to any dependents.  It's good hygiene to set this only once
   * processing is complete, but the {@link Scheduler} won't consume it until the task signals it is complete.
   */
  protected Object processingResult;

  private TaskState state = TaskState.CREATED;

  public WorkerTask(TaskSpec spec) {
    this.spec = spec;
    this.context = spec.getExecutionContext();
    // give each WorkerTask a unique ID (for debugging)
    this.id = spec.getWorkerTasks().size() + 1;

    // setup page read/writer helpers (if the task has input and/or output buffers)
    this.pageReader = spec.getInput().map(b -> new PageReader(b)).orElse(null);
    this.pageWriter = spec.getOutput().map(b -> new PageWriter(b)).orElse(null);


    in = pageReader != null ? spec.newMetric("tuples-in", "in", () -> new Meter()) : null;
    out = pageWriter != null ? spec.newMetric("tuples-out", "out", () -> new Meter()) : null;

    runtime = spec.newMetric("runtime", () -> new Counter());
    runtimeAverage = spec.newMetric("runtime-average", () -> new Histogram(new ExponentiallyDecayingReservoir()));
    contextSwitches = spec.newMetric("context-switches", () -> new Counter());

    for (RealizedStep step: spec.getForSteps()) {
      Realized realized = step.getResult().get();
      if (realized instanceof HasMeter) {
        HasMeter hasMeter = (HasMeter)realized;
        for (String keyMetric: hasMeter.getProgressMetricNames()) {
          spec.getProgressMetrics().put(
              step.getStepName() + "-" + keyMetric,
              hasMeter.getRegistry().getMetrics().get(keyMetric));
        }
      }
    }
  }

  /**
   * Processes the work that the task has to do. The task doesn't necessarily run
   * to completion in one go - it's likely that the task will run out of input or
   * output first, so it'll need to keep coming back and chipping away at the work.
   */
  public abstract ReturnState run();

  public ReturnState runPublic() {
    long startTime = System.nanoTime();

    try {
      return run();
    } finally {
      long stopTime = System.nanoTime();
      long runtimeMillis = TimeUnit.MILLISECONDS.convert(stopTime - startTime, TimeUnit.NANOSECONDS);

      runtime.inc(runtimeMillis);
      runtimeAverage.update(runtimeMillis);
    }
  }

  /**
   * Returns an optional helper for reading pages of input Tuples from the
   * ReadPageBuffer (note that not all tasks have an input buffer though).
   */
  public Optional<PageReader> getPageReader() {
    return Optional.ofNullable(pageReader);
  }

  /**
   * Returns an optional helper for writing pages of output Tuples to the
   * WritePageBuffer (note that not all tasks have an output buffer though).
   */
  public Optional<PageWriter> getPageWriter() {
    return Optional.ofNullable(pageWriter);
  }

  protected boolean hasInputPage() {
    return getPageReader().map(PageReader::hasPageInProgress).orElse(false);
  }

  protected boolean hasOutputPage() {
    return getPageWriter().map(PageWriter::hasPageInProgress).orElse(false);
  }

  /**
   * Returns true if the task currently has a page of either input or output that
   * it hasn't finished with yet.
   */
  public boolean hasPageInProgress() {
    // check whether we're halfway through either reading or writing a page
    return hasInputPage() || hasOutputPage();
  }

  private boolean isInputComplete() {
    return getPageReader().map(PageReader::isComplete).orElse(false);
  }

  /**
   * Test for whether this task should run based on the state of the input buffer.  Note that because of an awkward
   * case in the scheduling logic, this method will return true if there are no tuples *but* the input buffer has been
   * marked as complete.  This allows the task to run once more and be run to completion.  This is something we should
   * probably be handling in the scheduler logic instead of this test.
   *
   * @return true if there are tuples in the input buffer to be read, or an incomplete page, or the input buffer has
   * been marked as complete.
   */
  public final boolean isInputReady() {
    // note that a complete input buffer is technically empty, but the
    // task still needs to wake up to action that it's complete
    boolean hasInput = getPageReader().map(PageReader::hasInput).orElse(true);
    return hasInput || isInputComplete();
  }

  /**
   * Test for whether this task should run based on the state of the output buffer.
   *
   * It will ignore the state of the output buffer if there's no more input - but this check may disappear once we have
   * better test coverage and can confirm it's not relied upon.  If it is, it should be in a task specific check.
   *
   * @return true if there is capacity in the output or the input is complete
   */
  public final boolean isOutputReady() {
    // if the input buffer is complete then we know there's no more output,
    // so it doesn't matter if the output buffer is full or not - TODO dig further in to whether this makes sense or
    // not - it feels like that's quite a task-specific check that should probably be added to isReadyToRun rather than
    // this final method
    boolean noMoreOutput = isInputComplete();
    boolean hasSpace = getPageWriter().map(PageWriter::hasSpace).orElse(true);
    return hasSpace || noMoreOutput;
  }

  /**
   * Returns true if the task has work it can do. Returns false if the task is
   * blocked waiting on either more input, more output buffers to free up, or it's
   * dependent on other tasks that haven't completed yet.
   */
  public boolean isReadyToRun() {
    return !spec.hasOutstandingDependencies() && isInputReady() && isOutputReady();
  }

  protected ReturnState taskComplete() {
    // we must always make sure any partially written output gets flushed
    // when the task finishes - otherwise we could drop Tuples or deadlock
    getPageWriter().ifPresent(PageWriter::flushPage);
    return ReturnState.COMPLETE;
  }

  public RealizedStep getLastStep() {
    List<RealizedStep> steps = spec.getForSteps();
    return steps.get(steps.size() - 1);
  }

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

  protected Realized getFirstStepRealizedResult() {
    return getFirstStep().getResult().get();
  }

  @Override
  public String toString() {
    return String.format("%s-%d(steps=%s, input=%s, output=%s)", getClass().getSimpleName(), id,
        spec.getStepsSummary(), getPageReader().orElse(null), getPageWriter().orElse(null));
  }

  /**
   * Override this method to clean up any resources that were allocated/created by this task for use during execution.
   * Will be called from the scheduler once the task has signaled it is complete, but it may also get called if
   * a job that this task was part of has failed.
   *
   * Thread safety should be ensured by requiring the scheduler to only call close on a task that is not currently being
   * run.
   *
   * Implementations shouldn't need to do any buffer management in this method, it's meant for closing things like
   * {@link TupleIterator}s or other sorts of resources that follow the {@link Closeable} pattern.
   */
  @Override
  public void close() {
  }


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

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

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


  /**
   * @return a unique name for this worker task, based on the step(s) it covers, and the type of
   *         worker it is.
   */
  public final String getName() {
    // some steps (i.e. join) can have multiple workers (index, emit, etc) so we need the
    // class name too in order to make it unique
    return spec.getCombinedStepsName() + "." + getClass().getSimpleName().toLowerCase() + "-" + id;
  }

  /**
   * @return a unique name for the worker's TaskSpec that is brief and user-friendly. This can be
   *         used as a simple way to represent the work that this task is doing to the user.
   */
  public String getSpecNameBrief() {
    return spec.getFirstStep().getStepName();
  }

  public abstract boolean producesResult();

  /**
   * @return a processingResult that has been set, clearing its value in the process (to allow the result to be garbage
   * collected in the fullness of time)
   */
  public Object consumeProcessingResult() {
    Object result = processingResult;
    processingResult = null;
    return result;
  }

  public void markStarted() {
    state = TaskState.STARTED;
    // every time the task is started counts as a context switch
    contextSwitches.inc();
  }
  public void markComplete() {
    state = TaskState.COMPLETE;
  }
}
