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

import java.io.EOFException;
import java.io.IOException;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.apache.commons.pool2.impl.GenericObjectPool;

import com.codahale.metrics.Counter;
import com.codahale.metrics.ExponentiallyDecayingReservoir;
import com.codahale.metrics.Histogram;
import com.codahale.metrics.Meter;
import com.google.common.collect.Lists;

import lombok.Getter;
import nz.org.riskscape.cpython.CPythonSpawner.CPythonProcess;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.function.FunctionCallException;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.types.Anything;
import nz.org.riskscape.engine.types.Bool;
import nz.org.riskscape.engine.types.Enumeration;
import nz.org.riskscape.engine.types.Floating;
import nz.org.riskscape.engine.types.Geom;
import nz.org.riskscape.engine.types.Integer;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Smallfloat;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Text;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.StandardCodes;

/**
 * Executes a user's CPython code from RiskScape. We can't just do this in the
 * same process space as Java, like we do for Jython functions. We need to spawn
 * a separate CPython process, serialize the function args we pass it, and then
 * deserialize the result returned.
 *
 * Package scoped - it's not meant to be used directly, but as the result of realization.  I'd have made this an inner
 * class, but it existed first.
 */
class CPythonFunction implements RiskscapeFunction {

  /**
   * Types that are supported for CPython function return values.
   */
  // this should line up with the supported_types in the Python serializer.py code
  public static final List<Class<? extends Type>> SUPPORTED_RETURN_TYPES = Lists.newArrayList(
      Bool.class,
      Integer.class,
      Enumeration.class,
      Floating.class,
      RSList.class,
      Nullable.class,
      Smallfloat.class,
      Struct.class,
      Text.class,
      Geom.class
  );

  /**
   * Types that are supported for CPython function arguments.
   */
  public static final List<Class<? extends Type>> SUPPORTED_ARGUMENT_TYPES = Stream.concat(
      SUPPORTED_RETURN_TYPES.stream(),
      // also allow 'Anything' as an argument type, as we will just end up serializing
      // whatever type the function was realized with
      Stream.of(Anything.class)
  ).collect(Collectors.toList());

  @Getter
  private final List<Type> argumentTypes;

  private final CPythonRealizableFunction realizable;

  private final GenericObjectPool<CPythonProcess> pool;

  private final CPythonSerializer serializer = new CPythonSerializer();

  /**
   * Counter of function call runtime in milliseconds.
   */
  private final Counter runtime;

  protected final Histogram runtimeAverage;

  protected final Meter in;


  /**
   * Similar to {@link #runtime} but excludes the time taken to serialize/de-serialize function arguments and result.
   * If this counter is a lot smaller than runtime then the user most likely is sending too much data to the
   * function (maybe large geometries).
   */
  private final Counter runtimeExcludingSerialization;

  CPythonFunction(List<Type> argumentTypes, CPythonRealizableFunction realizable,
      GenericObjectPool<CPythonProcess> pool, RealizationContext context) {
    this.argumentTypes = argumentTypes;
    this.realizable = realizable;
    this.pool = pool;

    // make some metrics to appear in model-run-stats.csv
    // We use the metric names/types that are expected by StatsWriter.java so these stats will be
    // written out without any drama.
    String metricName = "Cpython_Function_" + realizable.identified().getId();
    this.runtime = context.getMetricRegistry().counter(metricName + ".runtime");
    this.runtimeAverage = context.getMetricRegistry().histogram(metricName + ".runtime-average",
        () -> new Histogram(new ExponentiallyDecayingReservoir()));
    // excluding serialization is not output in model-run-stats.
    // but it's left in as it may be useful for devs debugging slow functions (maybe)
    this.runtimeExcludingSerialization = context.getMetricRegistry().counter(
        metricName + ".excluding-serialization.runtime-ms");

    this.in = context.getMetricRegistry().meter(metricName + ".tuples-in");

    // register `in` as `out` (they are the same) for consistency in the model-run-stats output (might not be worth it?)
    context.getMetricRegistry().meter(metricName + ".tuples-out", () -> in);
  }

  @Override
  public Object call(List<Object> args) {
    // CPython needs to run in a separate process that we IPC to
    in.mark();
    CPythonProcess process;
    try {
      process = pool.borrowObject();
    } catch (RiskscapeException e) {
      throw e;
    } catch (Exception e) {
      throw new RiskscapeException(Problems.caught(e));
    }

    long startTime = System.nanoTime();
    try {
      // serialize the args we want to pass to the function onto the child process's stdin
      for (int i = 0; i < argumentTypes.size(); i++) {
        Type argType = argumentTypes.get(i);
        Object value = args.get(i);
        serializer.serialize(process.getChildStdIn(), argType, value);
      }
      process.getChildStdIn().flush();

      long startTimeFunctionCall = System.nanoTime();

      // first byte is a signal to say whether the function call succeeded or not
      Long status = (Long) Types.INTEGER.fromBytes(process.getChildStdOut());

      // NB do a cheap check for any output on stderr, if it looks like some is there then wake up the spawner's
      // reaping thread so that each function call is likely to have its own output kept near the actual invocation of
      // the function
      if (process.hasErrorText()) {
        realizable.runner.checkForOutput();
      }

      long stopTimeFunctionCall = System.nanoTime();
      runtimeExcludingSerialization
          .inc(TimeUnit.MILLISECONDS.convert(stopTimeFunctionCall - startTimeFunctionCall, TimeUnit.NANOSECONDS));

      if (status == 1) {
        // if it did, we deserialize the expect return value from the child process's stdout
        return getReturnType().fromBytes(process.getChildStdOut());
      } else {
        // if not, we read a string containing the error message and throw an exception containing it - this should
        // help the user fix/debug their function
        String errorMesssage = (String) Types.TEXT.fromBytes(process.getChildStdOut());
        throw new FunctionCallException(
          Problems.foundWith(realizable.identified(), CPythonProblems.get().functionCallException(errorMesssage)),
          null
        );
      }
    } catch (EOFException ex) {
      // give stderr a chance to be read - this is a more patient version of checkForOutput - we really, really want to
      // be sure that stderr is read as this might contain the reason why we got an EOF
      realizable.runner.checkForOutputAndWait();
      // if we got EOF reading the return-type, then it probably means that the child
      // process exited unexpectedly. Include the last bit of stderr from the child process in the error in case it
      // helps?
      throw new FunctionCallException(
          Problems.foundWith(realizable.identified(),
              Problem.error(StandardCodes.CAUGHT_EXCEPTION, process.lastErrorText)), ex);
    } catch (IOException ex) {
      throw new FunctionCallException(Problems.foundWith(realizable.identified(), Problems.caught(ex)), ex);
    } finally {
      long stopTime = System.nanoTime();
      long runtimeMillis = TimeUnit.MILLISECONDS.convert(stopTime - startTime, TimeUnit.NANOSECONDS);
      runtime.inc(runtimeMillis);
      runtimeAverage.update(runtimeMillis);
      try {
        pool.returnObject(process);
      } catch (RiskscapeException e) {
        throw e;
      } catch (Exception e) {
        throw new RiskscapeException(Problems.caught(e));
      }
    }
  }


  @Override
  public Type getReturnType() {
    return realizable.getMetadata().getReturnType();
  }

  @Override
  public void close() {
    realizable.runner.checkForOutputAndWait();
    pool.close();
  }

  @Override
  public String toString() {
    return String.format("CPythonFunction[%s]", realizable.getMetadata());
  }

}
