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

import java.io.BufferedReader;
import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.Reader;
import java.net.URI;
import java.util.AbstractMap;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.function.Predicate;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.python.core.Py;
import org.python.core.PyCode;
import org.python.core.PyException;
import org.python.core.PyFunction;
import org.python.core.PyJavaType;
import org.python.core.PyList;
import org.python.core.PyObject;
import org.python.core.PyString;
import org.python.core.PyStringMap;
import org.python.util.PythonInterpreter;

import com.google.common.collect.Lists;

import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeIOException;
import nz.org.riskscape.engine.function.FunctionMetadata;
import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.function.IdentifiedFunction.Category;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.resource.Resource;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.typeset.IdentifiedType;
import nz.org.riskscape.engine.typeset.MissingTypeException;
import nz.org.riskscape.engine.typeset.TypeSet;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

@Slf4j
/**
 * Creates Riskscape objects from python scripts
 */
public class JythonFactory implements Closeable, AutoCloseable, JavaJythonHelpers {

  private static final Object FUNCTION_OBJECT = JythonRiskscapeFunctionDelegate.class;

  /**
   * {@link Category} that will apply to jython functions that have not specifically set a category.
   */
  public static final Category DEFAULT_CATEGORY = Category.UNASSIGNED;

  private static final Predicate<String> ID_PATTERN_PRESENCE = Pattern.compile("^\\h*ID\\h*=\\h*").asPredicate();
  private static final Pattern ID_PATTERN_SINGLE_QUOTES = Pattern.compile("^\\h*ID\\h*=\\h*\\'(?<ID>.+)'\\h*$");
  private static final Pattern ID_PATTERN_DOUBLE_QUOTES = Pattern.compile("^\\h*ID\\h*=\\h*\\\"(?<ID>.+)\"\\h*$");

  private PythonInterpreter interpreter;
  private TypeSet typeSet;

  /**
   *  Fake map that only responds to get, calling the functionset to find a {@link IdentifiedFunction} before
   *  wrapping it with a {@link NestedFunction} adapter to allow seamless calling from jython.
   */
  private Map<String, NestedFunction> pythonFunctionMap;

  private Project project;

  public JythonFactory(Project project) {
    this.project = project;
    initialize();
    this.pythonFunctionMap = new AbstractMap<String, NestedFunction>() {

      @Override
      public Set<Entry<String, NestedFunction>> entrySet() {
        return Collections.emptySet();
      }

      @Override
      public NestedFunction get(Object key) {
        IdentifiedFunction function = project.getFunctionSet().get(key.toString(),
            project.getEngine().getProblemSink());

        return new NestedFunction(function);
      }

    };
    this.interpreter = new PythonInterpreter();
    this.typeSet = project.getTypeSet();
  }

  private void initialize() {
    Properties props = new Properties();
    props.put("python.import.site", "false");
    File jythonHome = new File(System.getProperty("user.home"), ".riskscape-jython");
    if (! jythonHome.exists()) {
      jythonHome.mkdir();
    }
    props.setProperty("python.home", jythonHome.toString());
    Properties preprops = System.getProperties();
    PythonInterpreter.initialize(preprops, props, new String[0]);
  }

  private PyObject loadGlobals(Resource resource) {
    Reader reader = resource.getContentReader();
    log.info("Loading function from: {}", resource.getLocation());
    // We can keep this in memory, but we may never use it - if it turns out that these are really
    // big (in memory terms) we might want to null the pointer and re parse inside the supplier.
    PyObject globals = compileToGlobals(reader, resource.getLocation(), FUNCTION_OBJECT);

    try {
      reader.close();
    } catch (IOException e) {
      throw new JythonScriptException(FUNCTION_OBJECT, resource.getLocation(), e);
    }
    return globals;
  }

  /**
   * Builds a Python function from the given Resource. The function's metadata (argument-types, etc)
   * are read from the Python code itself by looking for specially named global variables
   */
  public ResultOrProblems<IdentifiedFunction> buildFunction(String functionId, Resource resource) {
    try {
      PyObject globals = loadGlobals(resource);

      // the main function should be called "function"
      Optional<PyFunction> functionDefinition = findItem(globals, "function", PyFunction.class, FUNCTION_OBJECT,
          resource.getLocation());

      if (functionDefinition.isPresent()) {
        // snoop the function's return-type/etc from the Python global variables ('RETURN-TYPE', etc)
        FunctionMetadata metadata = snoopFunctionMetadata(functionId, globals, resource.getLocation());

        return ResultOrProblems.of(buildFunction(functionDefinition.get(), metadata, resource));
      } else {
        // no 'function' definition - this is only allowed if the user is wrapping a RiskScape function
        return ResultOrProblems.of(wrapRiskscapeFunction(functionId, globals, resource.getLocation()));
      }
    } catch (JythonScriptException ex) {
      return ResultOrProblems.failed(GeneralProblems.get().badResource(IdentifiedFunction.class, resource.getLocation())
          .withChildren(ex.getProblem()));
    }
  }

  /**
   * Builds a Python function from the given Resource with pre-defined function metadata.
   */
  public ResultOrProblems<IdentifiedFunction> buildFunction(Resource resource, FunctionMetadata metadata) {
    Problem badFunction = GeneralProblems.get().badResource(IdentifiedFunction.class, resource.getLocation());
    try {
      PyObject globals;
      try {
        globals = loadGlobals(resource);
      } catch (RiskscapeIOException e) {
        // the resource may throw a RiskscapeIOException. this will definitely happen if the resource
        // is pointing to a https resource that 404s.
        return ResultOrProblems.failed(Problems.caught(e));
      }

      // the main function should be called "function" and *must* be present
      Optional<PyFunction> functionDefinition = findItem(globals, "function", PyFunction.class, FUNCTION_OBJECT,
          resource.getLocation());
      if (!functionDefinition.isPresent()) {
        return ResultOrProblems.failed(
            badFunction.withChildren(JythonProblems.get().mustHaveFunction()));
      }
      return ResultOrProblems.of(buildFunction(functionDefinition.get(), metadata, resource));

    } catch (JythonScriptException ex) {
      return ResultOrProblems.failed(badFunction.withChildren(ex.getProblem()));
    }
  }

  /**
   * Searches through the given resource for an ID declaration - allows an ID to be extracted without having to parse
   * and (and potentially execute parts of) the script.  Falls back to parsing the script if the ID can't be found
   */
  protected String getFunctionID(Resource resource) {
    // First try to get the function ID by pattern matching. This will only work if the function
    // ID is defined on a single line. Which should be in 99.9% of the time
    try (BufferedReader reader = new BufferedReader(resource.getContentReader())) {
      String line;
      while ((line = reader.readLine()) != null) {

        if (ID_PATTERN_PRESENCE.test(line)) {
          String id = extractID(line, ID_PATTERN_SINGLE_QUOTES);
          if (id == null) {
            id = extractID(line, ID_PATTERN_DOUBLE_QUOTES);
          }

          if (id != null) {
            return id;
          }
        }
      }
    } catch (IOException e) {
      throw new JythonScriptException(FUNCTION_OBJECT, resource.getLocation(), e);
    }

    throw new JythonScriptException(FUNCTION_OBJECT, resource.getLocation(),
        JythonProblems.get().missingConstant("function", "ID"));
  }

  /**
   * @param lineFromScript to search
   * @param pattern containing an ID capture group to return if matched
   * @return ID if found in the input, null otherwise
   */
  private String extractID(String lineFromScript, Pattern pattern) {
    Matcher matcher = pattern.matcher(lineFromScript);
    if (matcher.matches()) {
      return matcher.group("ID");
    }
    return null;
  }

  private FunctionMetadata snoopFunctionMetadata(String functionId, PyObject globals, URI pathToScript) {
    String description = findItem(globals, "DESCRIPTION", PyString.class, FUNCTION_OBJECT, pathToScript)
        .map(d -> d.toString()).orElse(""); //DESCRIPTION is optional

    PyList arguments = findRequiredItem(globals, "ARGUMENT_TYPES", PyList.class, FUNCTION_OBJECT, pathToScript);
    List<Type> argumentTypes = Lists.newArrayList();
    for (Object rawArgument : arguments) {
      argumentTypes.add(toType(rawArgument, "function", pathToScript, "ARGUMENT_TYPES members"));
    }

    Type returnType = toType(
        findRequiredItem(globals, "RETURN_TYPE", PyObject.class, "function", pathToScript).__tojava__(Object.class),
        "function",
        pathToScript,
        "RETURN_TYPE"
    );

    Optional<PyString> categoryString = findItem(globals, "CATEGORY", PyString.class, FUNCTION_OBJECT, pathToScript);
    Category category = categoryString
        .map(cat -> {
          String catName = cat.toString().toUpperCase();
          try {
            return Category.valueOf(catName);
          } catch (IllegalArgumentException e) {
            throw new JythonScriptException("CATEGORY", pathToScript,
                GeneralProblems.get().
                notAnOption(catName, Category.class, Arrays.asList(Category.values())));
          }
        })
        .orElse(DEFAULT_CATEGORY);

    return new FunctionMetadata(functionId, argumentTypes, returnType, description, category, pathToScript);
  }

  private IdentifiedFunction buildFunction(PyFunction pythonFunction, FunctionMetadata metadata, Resource source) {
    return new JythonRealizableFunction(metadata, source, pythonFunction).identified();
  }

  /**
   * This handles the uncommon case where a user wants to create their own 'wrapper function' around
   * an existing RiskScape function (in particular {@link MathsFunction}). For example:
   * ```
   * from nz.org.riskscape.engine.function import Maths
   * ID = 'foo'
   * FUNCTION = Maths.newPolynomial([2, 1])
   * ```
   */
  private IdentifiedFunction wrapRiskscapeFunction(String functionId, PyObject globals, URI pathToScript) {
    String description = findItem(globals, "DESCRIPTION", PyString.class, FUNCTION_OBJECT, pathToScript)
        .map(d -> d.toString()).orElse(""); // DESCRIPTION is optional

    RiskscapeFunction exportedFunction = (RiskscapeFunction) findItem(globals,
        "FUNCTION", PyObject.class, FUNCTION_OBJECT, pathToScript)
        .orElseThrow(() -> {
          return new JythonScriptException(
              FUNCTION_OBJECT,
              pathToScript,
              JythonProblems.get().missingFunctionObject()
          );
        })
        .__tojava__(RiskscapeFunction.class);

    return exportedFunction.identified(functionId, description, pathToScript, DEFAULT_CATEGORY);
  }

  private Type toType(Object rawArgument, String typeOfThing, URI source, String explanation) {
    if (rawArgument instanceof String) {
      try {
        return typeSet.getRequired(rawArgument.toString());
      } catch (MissingTypeException e) {
        throw new JythonScriptException(explanation, source, e.getProblem());
      }
    } else if (rawArgument instanceof Type) {
      return (Type) rawArgument;
    } else {
      throw new JythonScriptException(typeOfThing, source,
          JythonProblems.get().thingWasWrongType(explanation, rawArgument.getClass())
      );
    }
  }

  /**
   * Create an asset type from jython.
   * @param resource containing the jython script
   * @return asset type
   */
  public Type createType(Resource resource) {
    log.info("Loading type from: {}", resource.getLocation());
    PyObject globals = compileToGlobals(resource.getContentReader(), resource.getLocation(), "type");

    // extract bits we are interested in
    PyString id = findRequiredItem(globals, "ID", PyString.class, "type", resource.getLocation());
    PyObject type = findRequiredItem(globals, "TYPE", PyObject.class, "type", resource.getLocation());

    try {
      Type toAdd = (Type)type.__tojava__(Type.class);
      typeSet.addType(id.getString(), resource, () -> ResultOrProblems.of(toAdd));
      return toAdd;
    } catch (ClassCastException e) {
      throw new JythonScriptException(
          "type",
          resource.getLocation(),
          JythonProblems.get().couldNotConvertToType().withChildren(Problems.caught(e)));
    }
  }

  /**
   * Looks for either prefix_TYPE or prefix_TYPE_ID from globals and attempts to
   * convert either in to an IdentifiedType.  Ids return an UnresolvedType
   * whereas a TYPE returns a ResolvedType with a bogus id based on the
   * functions id
   * TODO consider making that id actually resolvable
   * If either both or none are defined, an exception is raised
   */
  protected IdentifiedType extractType(String functionId, PyObject globals, String prefix, URI scriptSource) {
    Optional<PyString> typeId = findItem(globals, prefix + "_TYPE_ID", PyString.class, "type", scriptSource);
    Optional<PyObject> typeObject = findItem(globals, prefix + "_TYPE", PyObject.class, "type", scriptSource);

    if (typeId.isPresent() && typeObject.isPresent()) {
      throw new RuntimeException("Function has defined an id and a type for for " + prefix + ". Only one is allowed");
    }
    if (!typeId.isPresent() && !typeObject.isPresent()) {
      throw new RuntimeException("Function has defined neither id or type for for " + prefix
          + ". One of these must be defined");
    }

    if (typeId.isPresent()) {
      return typeSet.getLinkedType(typeId.get().toString());
    } else {
      try {
        Type type = (Type) typeObject.get().__tojava__(Type.class);

          String fakeId = functionId + "_" + prefix + "_ANON";
          return typeSet.add(fakeId,  type);
      } catch (ClassCastException ex) {
        throw new RuntimeException(prefix + "_TYPE is not a Type object", ex);
      }
    }
  }

  private PyObject compileToGlobals(Reader script, URI source, Object typeOfThing) {
    // compile code

    PyCode compiled;
    try {
      compiled = interpreter.compile(script, source.getPath());
    } catch (PyException ex) {
      throw new JythonScriptException(typeOfThing, source, ex);
    }

    // evaluate code against a new namespace
    PyStringMap globals = Py.newStringMap();
    globals.__setitem__("typeset", PyJavaType.wrapJavaObject(typeSet));
    globals.__setitem__("functions", PyJavaType.wrapJavaObject(pythonFunctionMap));

    // wraps a riskscape function in a callable
    globals.__setitem__("nested_function", toPyCallable(RiskscapeFunction.class,
        NestedFunction.class,
        (RiskscapeFunction rsFunction) -> new NestedFunction(rsFunction)));

    try {
      Py.exec(compiled, globals, globals);
    } catch (PyException ex) {
      throw new JythonScriptException(typeOfThing, source, ex);
    }

    return globals;
  }

  // TODO maybe do __tojava__ conversion here?
  protected <T extends PyObject> Optional<T> findItem(PyObject source, String key, Class<T> clazz, Object typeOfThing,
      URI scriptSource) {

    PyObject item = source.__finditem__(key.intern());

    if (item == null) {
      return Optional.empty();
    } else {
      if (!clazz.isInstance(item)) {
        throw new JythonScriptException(typeOfThing, scriptSource, JythonProblems.get().
            unexpectedType(key, clazz, item.getClass()));
      }

      @SuppressWarnings("unchecked")
      T cast = (T) item;
      return Optional.of(cast);
    }
  }

  protected <T extends PyObject> T findRequiredItem(PyObject source, String key, Class<T> clazz,
      Object typeOfThing, URI scriptSource) {

    return findItem(source, key, clazz, typeOfThing, scriptSource).orElseThrow(() -> {

      return new JythonScriptException(typeOfThing, scriptSource,
          JythonProblems.get().missingConstant(typeOfThing, key));
    });
  }

  @Override
  public void close() {
    this.interpreter.close();
  }
}
