/*
 * 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.util.List;
import java.util.stream.Collectors;

import com.codahale.metrics.Meter;
import com.google.common.collect.Lists;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.pipeline.Realized;
import nz.org.riskscape.engine.pipeline.RealizedStep;
import nz.org.riskscape.engine.pipeline.Sink;
import nz.org.riskscape.engine.pipeline.SinkConstructor;
import nz.org.riskscape.engine.projection.FlatProjector;
import nz.org.riskscape.engine.projection.Projector;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.restriction.Restrictor;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ResultComputationException;


public class ChainTask extends WorkerTask {

  @RequiredArgsConstructor
  public static class NamedOp {
    public final String name;
    public final Object op;

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

  @RequiredArgsConstructor
  public static class StackElement {
    public final TupleIterator source;
    public final int stepIndex;
  }


  @Getter
  private final List<NamedOp> operations;

  private final List<StackElement> visitStack = Lists.newArrayList();
  private StackElement currentElement;

  private Meter tuplesSunk;

  public ChainTask(TaskSpec spec) throws ProblemException {
    super(spec);

    try {
    this.operations = spec.getForSteps().stream()
        .map(rs -> new NamedOp(rs.getStepName(), getOp(rs)))
        .collect(Collectors.toList());
    } catch (ResultComputationException ex) {
      throw new ProblemException(ex.getProblem());
    }

    if (pageWriter == null) {
      tuplesSunk = spec.newMetric("tuples-sunk", () -> new Meter());
    }
  }


  private Object getOp(RealizedStep rs) {
    Realized r = rs.getResult().get();

    if (r instanceof SinkConstructor) {
      Sink sink = ((SinkConstructor) r).newInstance(getSpec().getJobContext()).get();

      // sanity check / assertion -
      // only unbounded sinks make sense being attached to a chain task - if they are bounded they should be fed via a
      // buffer
      if (!sink.isUnbounded()) {
        throw new IllegalArgumentException("SinkConstructor returned a bounded sink - not allowed");
      }

      return sink;
    } else {
      return r;
    }

  }


  private boolean finishedLastPage() {
    // here we check currentElement, in case we stopped in the middle
    // of processing the last tuple in the page (in which case the
    // PageReader would think we've finished that page)
    return currentElement == null;
  }

  @Override
  public boolean hasPageInProgress() {
    return !finishedLastPage() || super.hasPageInProgress();
  }

  @Override
  public boolean isReadyToRun() {
    boolean isReady;
    if (pageWriter == null) {
      isReady = isInputReady();
    } else {
      boolean hasStackElement = !finishedLastPage() && pageWriter.hasSpace();
      isReady = hasStackElement || super.isReadyToRun();
    }

    return isReady;
  }

  @Override
  public ReturnState run() {

    while (true) {
      if (isFull()) {
        return ReturnState.OUTPUT_FULL;
      }

      // check if we need to read a new page
      if (finishedLastPage()) {
        if (pageReader.isComplete()) {
          // if the last operation was a sink, it's possible that other tasks downstream will depend on it
          if (getLastOperation().op instanceof Sink) {
            processingResult = getLastOperation().op;

            // mark it as complete
            ((Sink)processingResult).finish();
          }

          return taskComplete();
        }

        if (!pageReader.hasNext()) {
          return ReturnState.INPUT_EMPTY;
        }
      }
      processPage();
    }
  }

  private boolean isFull() {
    return pageWriter != null && pageWriter.isFull();
  }

  private void processPage() {
    if (currentElement == null) {
      // set up "virtual stack" state if a previous call to processPage exhausted the chain, otherwise we
      // will throw away tuples
      currentElement = new StackElement(pageReader, 0);
    }

    whileloop:
    while (!isFull()) {

      Tuple transforming;

      if (!currentElement.source.hasNext()) {
        currentElement.source.close();
        if (visitStack.isEmpty()) {
          // we've finished the current page
          currentElement = null;
          break whileloop;
        } else {
          currentElement = visitStack.remove(visitStack.size() - 1);
          continue whileloop;
        }
      }
      transforming = currentElement.source.next();

      // if the stack is empty, it means we just pulled a tuple from the input buffer
      if (visitStack.isEmpty()) {
        in.mark();
      }

      for (int i = currentElement.stepIndex; i < operations.size(); i++) {

        Object op = operations.get(i).op;
        if (op instanceof Projector) {
          transforming = ((Projector) op).apply(transforming);
        } else if (op instanceof Restrictor) {
          boolean skip = !((Restrictor) op).test(transforming);
          if (skip) {
            transforming = null;
            continue whileloop;
          }
        } else if (op instanceof FlatProjector) {
          TupleIterator flatProjecting = ((FlatProjector) op).apply(transforming);

          if (flatProjecting.hasNext()) {
            // we have a flat projector with results, get the first
            transforming = flatProjecting.next();

            // if we have more to consume from this, pop the current iterator on the stack and next iteration of the
            // loop starts from here - this is a little bit like loop unrolling - it avoids a needless pop on to the
            // stack
            if (flatProjecting.hasNext()) {
              visitStack.add(currentElement);
              currentElement = new StackElement(flatProjecting, i+1);
            } else {
              flatProjecting.close();
            }
          } else {
            flatProjecting.close();
            // empty flat projected iterator - back to start
            continue whileloop;
          }
        } else if (op instanceof Sink) {
          boolean accepted = ((Sink)op).accept(transforming);
          tuplesSunk.mark();
          if (!accepted) {
            // we could return output full, but at the moment, we're assuming we're only being used with unbounded
            // sinks, so let's fail instead
            throw new IllegalStateException("sink is full - not supposed to happen");
          }
          transforming = null;
          continue whileloop;
        } else {
          throw new RiskscapeException("Bad step type " + op);
        }
      }

      if (transforming != null) {
        out.mark();
        pageWriter.add(transforming);
      }
    }
  }

  @Override
  public void close() {
    if (currentElement != null) {
      currentElement.source.close();
    }
    for (StackElement stackElement : visitStack) {
      stackElement.source.close();
    }
  }

  @Override
  public boolean producesResult() {
    return getLastOperation().op instanceof Sink;
  }

  public NamedOp getLastOperation() {
    return operations.get(operations.size() - 1);
  }

  @Override
  public String getSpecNameBrief() {
    if (spec.getForSteps().size() > 1) {
      // hopefully just the first and last step should make the pipeline branch clear
      return "steps-" + spec.getFirstStep().getStepName() + "~>" + spec.getLastStep().getStepName();
    } else {
      return super.getSpecNameBrief();
    }
  }
}
