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

import java.io.IOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.jar.Manifest;

import com.google.common.collect.Lists;

import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.config.BootstrapIniSettings;
import nz.org.riskscape.config.ConfigUtils;
import nz.org.riskscape.defaults.auth.KoordinatesKeySecret;
import nz.org.riskscape.defaults.auth.PlatformSecret;
import nz.org.riskscape.defaults.curves.ContinuousLinearFitter;
import nz.org.riskscape.defaults.curves.CurveFitter;
import nz.org.riskscape.defaults.curves.FitCurveFunction;
import nz.org.riskscape.defaults.curves.LinearFitter;
import nz.org.riskscape.defaults.curves.PowerLawFitter;
import nz.org.riskscape.defaults.curves.TrapezoidIntegrationFunction;
import nz.org.riskscape.defaults.function.AALHazardBasedFunction;
import nz.org.riskscape.defaults.function.BucketFunction;
import nz.org.riskscape.defaults.function.BucketRange;
import nz.org.riskscape.defaults.function.CombineCoverages;
import nz.org.riskscape.defaults.function.ListToColumns;
import nz.org.riskscape.defaults.function.LookupFunction;
import nz.org.riskscape.defaults.function.MapCoverage;
import nz.org.riskscape.defaults.function.Segment;
import nz.org.riskscape.defaults.function.SegmentByGrid;
import nz.org.riskscape.defaults.function.ToListFunction;
import nz.org.riskscape.defaults.function.ToLookupTable;
import nz.org.riskscape.defaults.function.expression.ExpressionFunctionFramework;
import nz.org.riskscape.defaults.interp.ApplyContinuousFunction;
import nz.org.riskscape.defaults.interp.CreateContinuousFunction;
import nz.org.riskscape.defaults.interp.StackContinuousFunction;
import nz.org.riskscape.engine.Engine;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.auth.SecretBuilders;
import nz.org.riskscape.engine.bind.DefaultBindingContext;
import nz.org.riskscape.engine.data.BookmarkResolvers;
import nz.org.riskscape.engine.data.coverage.CoverageFileBookmarkResolver;
import nz.org.riskscape.engine.defaults.data.CsvResolver;
import nz.org.riskscape.engine.defaults.data.DataURILoader;
import nz.org.riskscape.engine.defaults.data.GeoJSONResolver;
import nz.org.riskscape.engine.defaults.data.GeoPackageRelationResolver;
import nz.org.riskscape.engine.defaults.data.KmlResolver;
import nz.org.riskscape.engine.defaults.data.ShapefileBookmarkResolver;
import nz.org.riskscape.engine.defaults.data.WfsBookmarkResolver;
import nz.org.riskscape.engine.defaults.data.jdbc.GeoPackageFormat;
import nz.org.riskscape.engine.defaults.function.IsExposed;
import nz.org.riskscape.engine.defaults.function.LookupBookmark;
import nz.org.riskscape.engine.defaults.resource.HttpResourceLoader;
import nz.org.riskscape.engine.defaults.resource.SwiftObjectStorageResourceLoader;
import nz.org.riskscape.engine.function.ApplyFunction;
import nz.org.riskscape.engine.function.FunctionClassLoader;
import nz.org.riskscape.engine.function.FunctionContext;
import nz.org.riskscape.engine.function.FunctionProvider;
import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.function.IdentifiedFunction.Category;
import nz.org.riskscape.engine.function.JavaFunction;
import nz.org.riskscape.engine.function.StringFunctions;
import nz.org.riskscape.engine.function.geometry.LayerIntersections;
import nz.org.riskscape.engine.function.geometry.SampleClosest;
import nz.org.riskscape.engine.function.geometry.SampleCoverageAtCentroid;
import nz.org.riskscape.engine.function.geometry.SampleIntersections;
import nz.org.riskscape.engine.function.geometry.SampleValues;
import nz.org.riskscape.engine.function.geometry.ToTypedCoverage;
import nz.org.riskscape.engine.function.lang.Call;
import nz.org.riskscape.engine.function.lang.MergeStruct;
import nz.org.riskscape.engine.output.CsvFormat;
import nz.org.riskscape.engine.output.FileSystemPipelineOutputStore;
import nz.org.riskscape.engine.output.GeoJSONFormat;
import nz.org.riskscape.engine.output.GeoPackagePipelineOutputStore;
import nz.org.riskscape.engine.output.KmlFormat;
import nz.org.riskscape.engine.output.QGSWriter;
import nz.org.riskscape.engine.output.ShapefileFormat;
import nz.org.riskscape.engine.output.ShapefileGeotoolsFormat;
import nz.org.riskscape.engine.pipeline.ParameterizedPipelineModelFramework;
import nz.org.riskscape.engine.plugin.PluginDescriptor;
import nz.org.riskscape.engine.plugin.PluginFeature;
import nz.org.riskscape.engine.resource.Resource;
import nz.org.riskscape.engine.rl.DefaultOperators;
import nz.org.riskscape.engine.rl.GeometryFunctions;
import nz.org.riskscape.engine.rl.LanguageFunctions;
import nz.org.riskscape.engine.rl.LogicFunctions;
import nz.org.riskscape.engine.rl.MathsFunctions;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.rl.agg.AggregationFunction;
import nz.org.riskscape.engine.spi.EngineBootstrapper;
import nz.org.riskscape.engine.steps.EnlargeStep;
import nz.org.riskscape.engine.steps.FilterStep;
import nz.org.riskscape.engine.steps.GroupByStep;
import nz.org.riskscape.engine.steps.JoinStep;
import nz.org.riskscape.engine.steps.RelationInputStep;
import nz.org.riskscape.engine.steps.SaveStep;
import nz.org.riskscape.engine.steps.SelectStep;
import nz.org.riskscape.engine.steps.SortStep;
import nz.org.riskscape.engine.steps.UnionStep;
import nz.org.riskscape.engine.steps.UnnestStep;
import nz.org.riskscape.engine.types.TypeRegistry;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.ResultOrProblems;

@Slf4j
public class Plugin extends nz.org.riskscape.engine.plugin.Plugin implements FunctionProvider {

  private boolean useV2ClassifierFormat = true;

  private boolean outputQgisProject = true;

  /**
   * Manifest file attribute name expected to contain a space separated list of class names to load
   * as {@link IdentifiedFunction}s.
   */
  public static final String FUNCTION_CLASSES_ATTRIBUTE_NAME = "Riskscape-Function-Classes";
  private static final String CLASSIFIER_FUNCTIONS_USE_NEW_FORMAT = "classifier-functions.v2-enabled";

  private static final String OUTPUT_QGIS_PROJECT = "output.qgis-project";

  public static final URI DEFAULTS_SOURCE_URI = URI.create("riskscape:default-plugin");

  /**
   * Dummy {@link Resource} to identify things as coming from the defaults plugin.
   */
  public static final Resource DEFAULTS_SOURCE = new Resource.PseudoResource(DEFAULTS_SOURCE_URI);

  public static void addDefaultTypeInformation(Engine engine) {
    TypeRegistry.addDefaults(engine.getTypeRegistry());
  }

  public Plugin(PluginDescriptor pDescriptor) {
    super(pDescriptor);
  }

  @Override
  public List<Problem> initializeEngine(Engine engine) {
    BookmarkResolvers resolvers = engine.getBookmarkResolvers();
    resolvers.add(new ShapefileBookmarkResolver(engine));
    resolvers.add(new CoverageFileBookmarkResolver(engine));
    resolvers.add(new WfsBookmarkResolver(engine));
    resolvers.add(new CsvResolver(engine));
    resolvers.add(new GeoJSONResolver(engine));
    resolvers.add(new KmlResolver(engine));
    resolvers.add(new GeoPackageRelationResolver(engine));

    engine.getModelFrameworks().add(new ParameterizedPipelineModelFramework(engine));

    engine.getBinders().addAll(DefaultBindingContext.DEFAULT_BINDERS);

    engine.getFunctionFrameworks().add(new ExpressionFunctionFramework());

    addOutputFormats(engine);
    addDefaultTypeInformation(engine);

    engine.getResourceFactory().add(new DataURILoader(engine));
    engine.getPipelineOutputStores().add(new FileSystemPipelineOutputStore());
    engine.getPipelineOutputStores().add(new GeoPackagePipelineOutputStore());

    if (outputQgisProject) {
      // update the FileSystemPipelineOutputStore to save the project.qgs file as well
      engine.getPipelineOutputStores().getAll().stream()
        .filter(store -> store instanceof FileSystemPipelineOutputStore)
        .map(store -> (FileSystemPipelineOutputStore) store)
        .findFirst()
        .ifPresent(store -> store.onCompletion(QGSWriter.WRITE_PROJECT_FILE));
    }

    engine.getCollection(SecretBuilders.class).add(PlatformSecret.BUILDER);
    engine.getCollection(SecretBuilders.class).add(KoordinatesKeySecret.BUILDER);

    return Collections.emptyList();
  }

  @Override
  public void startUp(Map<String, List<String>> settingsConfig, ProblemSink sink) {
    super.startUp(settingsConfig);
    BootstrapIniSettings settings = new BootstrapIniSettings(settingsConfig, sink);
    this.useV2ClassifierFormat = settings.getOrDefault(CLASSIFIER_FUNCTIONS_USE_NEW_FORMAT, true);
    this.outputQgisProject = settings.getOrDefault(OUTPUT_QGIS_PROJECT, true);

    if (!useV2ClassifierFormat) {
      log.warn("Support for v1 classifier functions has been removed, ignoring setting {}",
          CLASSIFIER_FUNCTIONS_USE_NEW_FORMAT);
    }

  }

  @Override
  public List<Problem> initializeProject(Project project, Engine engine) {
    project.add(new IsExposed().asFunction().builtin("is_exposed", Category.RISK_MODELLING));

    return Collections.emptyList();
  }

  @Override
  public void addFunctions(Project project, Consumer<Problem> problemSink) {

    addFunctionsFromConfig(project).forEach(problemSink);

    project.getFunctionSet().insertFirst(DefaultOperators.INSTANCE);
    GeometryFunctions geometryFunctions = new GeometryFunctions(project.getEngine());
    project.getFunctionSet().addAll(geometryFunctions.getPredicates(), DEFAULTS_SOURCE);
    project.getFunctionSet().addAll(geometryFunctions.getFunctions(), DEFAULTS_SOURCE);
    project.getFunctionSet().addAll(LogicFunctions.LOGIC_FUNCTIONS, DEFAULTS_SOURCE);
    project.getFunctionSet().addAll(MathsFunctions.FUNCTIONS, DEFAULTS_SOURCE);
    project.getFunctionSet().addAll(StringFunctions.FUNCTIONS, DEFAULTS_SOURCE);
    project.getFunctionSet().addAll(LanguageFunctions.FUNCTIONS, DEFAULTS_SOURCE);


    // TODO add this to LookupBookmark to DRY up with tests
    LookupBookmark lookupBookmark = new LookupBookmark();
    project.getFunctionSet().add(
        RealizableFunction.asFunction(lookupBookmark, lookupBookmark.getArguments(), Types.ANYTHING)
            .builtin("bookmark", IdentifiedFunction.Category.MISC),
        DEFAULTS_SOURCE);

    project.getFunctionSet().add(new ToTypedCoverage()
        .builtin("relation_to_coverage", Category.MISC), DEFAULTS_SOURCE);

    project.getFunctionSet().add(new MergeStruct().builtin("merge", Category.LANGUAGE), DEFAULTS_SOURCE);
    project.getFunctionSet().add(new Call().builtin("call", Category.LANGUAGE), DEFAULTS_SOURCE);

    SampleIntersections sampleIntersections = new SampleIntersections();
    project.getFunctionSet().add(
        sampleIntersections.builtin("sample_intersections",Category.GEOMETRY_PROCESSING),
        DEFAULTS_SOURCE
    );
    // sample intersections is also added under the older name sample
    project.getFunctionSet().add(
        sampleIntersections.builtin("sample", Category.GEOMETRY_PROCESSING),
        DEFAULTS_SOURCE
    );

    project.getFunctionSet().add(new SampleCoverageAtCentroid()
        .builtin("sample_centroid", Category.GEOMETRY_PROCESSING), DEFAULTS_SOURCE);

    SampleClosest sampleClosest = new SampleClosest();
    project.getFunctionSet().add(
        sampleClosest.builtin("sample_closest", Category.GEOMETRY_PROCESSING),
        DEFAULTS_SOURCE
    );
    // sample closest is also added under the older name sample_one
    project.getFunctionSet().add(
        sampleClosest.builtin("sample_one", Category.GEOMETRY_PROCESSING),
        DEFAULTS_SOURCE
    );
    project.getFunctionSet().add(
            new SampleValues().builtin("sample_values", Category.GEOMETRY_PROCESSING),
            DEFAULTS_SOURCE
    );

//    project.getFunctionSet().add(new ToRelation()
//        .builtin("to_relation", Category.GEOMETRY_PROCESSING), DEFAULTS_SOURCE);

    project.getFunctionSet().add(new Segment()
        .builtin("segment", IdentifiedFunction.Category.GEOMETRY_PROCESSING), DEFAULTS_SOURCE
    );
    project.getFunctionSet().add(new SegmentByGrid()
        .builtin("segment_by_grid", IdentifiedFunction.Category.GEOMETRY_PROCESSING), DEFAULTS_SOURCE
    );

    project.getFunctionSet().add(new LayerIntersections()
        .builtin("layer_intersection", IdentifiedFunction.Category.GEOMETRY_PROCESSING), DEFAULTS_SOURCE
    );

    project.getFunctionSet().add(
      new CombineCoverages().builtin("combine_coverages", IdentifiedFunction.Category.GEOMETRY_PROCESSING)
    );
    project.getFunctionSet().add(new MapCoverage().builtin("map_coverage", Category.GEOMETRY_PROCESSING));

    project.getFunctionSet().add(new ToListFunction().builtin("to_list", Category.LANGUAGE));

    project.getFunctionSet().add(new LookupFunction().builtin("lookup", Category.MISC));
    project.getFunctionSet().add(new ListToColumns().builtin("list_to_columns", Category.MISC));

    project.getFunctionSet().add(new ToLookupTable().asRiskscapeFunction().builtin("to_lookup_table", Category.MISC));

    // bucketing functions
    project.getFunctionSet().add(AggregationFunction.asFunction(new BucketFunction()).builtin("bucket", Category.MISC));
    project.getFunctionSet().add(
        AggregationFunction.asFunction(new BucketRange(new BucketFunction())).builtin("bucket_range", Category.MISC)
    );

    // continuous functions
    project.getFunctionSet().addAll(Arrays.asList(
      new CreateContinuousFunction().asFunction().builtin("create_continuous", Category.MATHS),
      new ApplyContinuousFunction().asFunction().builtin("apply_continuous", Category.MATHS),
      AggregationFunction.asFunction(new StackContinuousFunction()).builtin("stack_continuous", Category.MATHS)
    ));


    project.getFunctionSet().add(new ApplyFunction().asFunction().builtin("apply", Category.MATHS));

    // TODO it would be nice to be able to have plugins be able to define their own identified collections, and allow
    // them to be looked up by class.  This plugin could just pass the entire identified collection to the
    // function, rather than having a statically defined list.  Other plugins could add their own fitters independently
    // of the construction of this function.
    List<CurveFitter<?>> fitters = Arrays.asList(
      new LinearFitter(),
      new PowerLawFitter(),
      new ContinuousLinearFitter()
    );
    project.getFunctionSet().add(
        AggregationFunction.asFunction(new FitCurveFunction(fitters)).builtin("fit_curve", Category.MATHS));

    project.getFunctionSet().add(new TrapezoidIntegrationFunction().builtin("trapz", Category.MATHS));
    project.getFunctionSet().add(new AALHazardBasedFunction().builtin("aal_trapz", Category.MATHS));
  }

  private Collection<? extends Problem> addFunctionsFromConfig(Project project) {
    List<Problem> problems = new ArrayList<>();

    ConfigUtils.getResources(project, p -> problems.add(p), cs -> cs.getName().startsWith("project"), "functions")
      .forEach(resource -> {
        if (resource.getLocation().getPath().endsWith(".jar")) {
          loadFunctionsFromJar(resource, project, problems);
        }
      }
    );

    return problems;
  }

  public void addOutputFormats(Engine engine) {
    engine.getFormats().add(new CsvFormat());
    engine.getFormats().add(new ShapefileFormat());
    engine.getFormats().add(new ShapefileGeotoolsFormat());
    engine.getFormats().add(new GeoJSONFormat());
    engine.getFormats().add(new KmlFormat());
    engine.getFormats().add(new GeoPackageFormat());
  }

  /**
   * Loads functions contained within a jar file and adds them to the project.
   *
   * @param jarFunctionResource resource expected to be a jar containing functions
   * @param project to add loaded functions to
   */
  void loadFunctionsFromJar(Resource jarFunctionResource, Project project, List<Problem> problems) {
    FunctionContext context = new FunctionContext(project.getTypeSet());
    try {
      // jar files need to be local if possible
      ResultOrProblems<Path> localPath = jarFunctionResource.ensureLocal(Resource.SECURE_OPTIONS);
      problems.addAll(localPath.getProblems());
      if (!localPath.isPresent()) {
        return;
      }
      URL jarURL = localPath.get().toUri().toURL();
      FunctionClassLoader functionLoader = new FunctionClassLoader(new URL[]{jarURL}, Engine.class.getClassLoader());

      URL manifestUrl = functionLoader.findResource("META-INF/MANIFEST.MF");
      if (manifestUrl == null) {
        functionLoader.close();
        problems.add(Problem.error("Could not find required META-INF/MANIFEST.MF in %s",
            jarFunctionResource.getLocation()));

      } else {
        Manifest manifest = new Manifest(manifestUrl.openStream());
        String functionClasses = manifest.getMainAttributes().getValue(FUNCTION_CLASSES_ATTRIBUTE_NAME);
        Class<?> loadedClass = null;
        List<Problem> jarProblems = Lists.newArrayList();
        for (String functionClassName : functionClasses.split("\\s")) {
          try {
            loadedClass = functionLoader.loadClass(functionClassName);
            if (!IdentifiedFunction.class.isAssignableFrom(loadedClass)) {
              jarProblems.add(Problem.error("Function class %s must extend %s",
                  functionClassName, IdentifiedFunction.class));
              continue;
            }

            @SuppressWarnings("unchecked")
            Class<IdentifiedFunction> functionClass = (Class<IdentifiedFunction>) loadedClass;
            Constructor<IdentifiedFunction> constructor = functionClass.getConstructor(FunctionContext.class);
            IdentifiedFunction function = constructor.newInstance(context);
            if (function instanceof JavaFunction) {
              ((JavaFunction)function).setSourceURI(jarFunctionResource.getLocation());
            }
            project.getFunctionSet().add(function, jarFunctionResource);

          } catch (NoSuchMethodException e) {
            jarProblems.add(Problem.error(e, "Function class %s does not define %s(%s)",
                functionClassName, loadedClass.getSimpleName(), FunctionContext.class));

          } catch (ClassNotFoundException e) {
            jarProblems.add(Problem.error(e, "Function class %s does not exist", functionClassName));

          } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) {
            jarProblems.add(Problem.error(e, "Function class %s could not be instantiated. Cause: %s",
                functionClassName, e.getMessage()));

          }
        }
        if (! jarProblems.isEmpty()) {
          problems.add(Problem.composite(jarProblems, "Problems found in function jar %s",
              jarFunctionResource.getLocation()));
        }
      }
    } catch (MalformedURLException e) {
      problems.add(Problem.error(e, "Could not find a jar file at %s", jarFunctionResource.getLocation()));
    } catch (IOException e) {
      problems.add(Problem.error(e, "Could not read manifest file from %s", jarFunctionResource.getLocation()));
    }
  }


  public void addSteps(Engine engine) {
    engine.getPipelineSteps().addAll(Arrays.asList(
      new RelationInputStep(engine),
      new FilterStep(engine),
      new JoinStep(engine),
      new UnnestStep(engine),
      new SelectStep(engine),
      new SortStep(engine),
      new GroupByStep(engine),
      new SaveStep(engine),
      new EnlargeStep(engine),
      new UnionStep(engine)
    ));
  }

  public void addResourceLoaders(Engine engine) {
    engine.getResourceFactory().add(new HttpResourceLoader(engine));
    engine.getResourceFactory().add(new SwiftObjectStorageResourceLoader());
  }

  @Override
  public List<PluginFeature> getFeatures() {
    return Arrays.asList(
      new EngineBootstrapper("default-pipeline-steps", this::addSteps),
      new EngineBootstrapper("default-resource-loaders", this::addResourceLoaders)
    );
  }
}
