/*
 * 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 static nz.org.riskscape.engine.Assert.*;
import static nz.org.riskscape.engine.Matchers.*;
import static org.hamcrest.Matchers.*;
import static org.junit.Assert.*;

import java.util.List;

import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import com.google.common.collect.Lists;

import nz.org.riskscape.engine.OutputProblems;
import nz.org.riskscape.engine.cli.ExitException;
import nz.org.riskscape.engine.cli.tests.BaseModelRunCommandTest;
import nz.org.riskscape.engine.resource.ResourceProblems;
import nz.org.riskscape.engine.test.EngineTestPlugins;
import nz.org.riskscape.engine.test.EngineTestSettings;
import nz.org.riskscape.problem.Problem;

@EngineTestPlugins({"defaults", "cpython"})
@EngineTestSettings({"cpython.python3-bin=python3", "cpython.lib-dir=lib"})
public class CPythonStepIntegrationTest extends BaseModelRunCommandTest {

  CPythonExecStep execStep;

  @Before
  public void addExecStep() {
    // while the python step is in beta, we need to do add the step to the engine 'manually'
    CPythonPlugin plugin =
        (CPythonPlugin) engine.getPlugins().stream().filter(p -> p.getId().equals("cpython")).findAny().get();
    execStep = new CPythonExecStep(engine, plugin.getSpawner());
    engine.getPipelineSteps().add(execStep);
  }

  @After
  public void removeExecStep() {
    engine.getPipelineSteps().remove(execStep.getId());
  }

  @Test
  public void canRunAModelThatUsesThePythonStepToProcessRows() throws Exception {
    runCommand.modelId = "dataframe-aal";
    runCommand.parameters = List.of();
    runCommand.run();

    List<List<String>> rows = openCsv("python.csv",
        "aal"
    );

    assertThat(rows, contains(contains("3410.0")));
  }

  @Test
  public void canAcceptParametersInPythonStep() throws Exception {
    // this model is basically the same as dataframe-aal (above test) except it passes
    // parameters through to the python script
    runCommand.modelId = "dataframe-aal-params";
    runCommand.parameters = List.of();
    runCommand.run();

    List<List<String>> rows = openCsv("python.csv",
        "aal"
    );

    assertThat(rows, contains(contains("3410.0")));
  }

  @Test
  public void canPlumbModelParametersThroughToPythonScript() throws Exception {
    // same as previous test except we check we can override the python parameters passed though with model $parameters
    runCommand.modelId = "dataframe-aal-params";
    // change the input data to have a different 'probability' attribute name ('prob'). We need to pass that
    // attribute name through to the CPython code in order for the AAL calc to still work
    runCommand.parameters = List.of("input=event-loss2", "probability=prob");
    runCommand.run();

    List<List<String>> rows = openCsv("python.csv",
        "aal"
    );

    assertThat(rows, contains(contains("3500.0")));
  }

  @Test
  public void pythonErrorIfScriptParametersAreMandatory() throws Exception {
    // similar again, except this time the Python script *must* be passed parameters, but the pipeline doesn't
    runCommand.modelId = "dataframe-aal";
    runCommand.parameters = List.of("script='mandatory_params.py'");
    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());
    String errors = render(ex.getProblem());
    assertThat(errors, allOf(
        containsString("Failed to load your CPython script for `python` step on line 1."),
        containsString("Your python() pipeline step is passing a different number of arguments "
                + "to what your Python script expects."),
        containsString("Your Python function takes 2 positional arguments, but the python() "
                + "pipeline step is currently only passing 1 argument "
                + "(which is the rows of pipeline data, via a generator function)."),
        containsString("Your Python function may be expecting parameters that are missing "
                + "from the python() pipeline step")
    ));
  }

  @Test
  public void canPassMandatoryParametersThroughToPythonScript() throws Exception {
    // sanity-check the python script in the previous test works in a model where we *do* supply parameters
    runCommand.modelId = "dataframe-aal-params";
    runCommand.parameters = List.of("script='mandatory_params.py'");
    runCommand.run();

    List<List<String>> rows = openCsv("python.csv",
            "aal"
    );

    assertThat(rows, contains(contains("3410.0")));
  }

  @Test
  public void pythonErrorIfUserSuppliedParametersDoNotMatchUp() throws Exception {
    runCommand.modelId = "raw-params";
    // parameters is missing the probability attribute needed by the function.
    // So the pipeline step is syntactically valid, but it will cause a Python programming error
    runCommand.parameters = List.of("input=event-loss2", "parameters={loss: 'loss'}");

    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());
    String errors = render(ex.getProblem());
    assertThat(errors, allOf(
        containsString("Problems found with `python` step on line 2"),
        containsString("Your CPython function raised an error: KeyError: 'probability'"),
        // this is where the error is in python parameters['probability']
        containsString("df['expected_loss'] = df[parameters['probability']] * df[parameters['loss']]")
    ));
  }

  @Test
  public void cannotPassParametersToFunctionThatDoesNotAcceptThem() {
    runCommand.modelId = "raw-params";
    // count.py does not accept parameters
    runCommand.parameters = List.of("script=count.py", "parameters={loss: 'loss'}");

    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());
    String errors = render(ex.getProblem());
    assertThat(errors, allOf(
        containsString("Failed to load your CPython script for `python` step on line 2"),
        containsString("Your python() pipeline step is passing a different number of arguments "
             + "to what your Python script expects."),
        containsString("The python() pipeline step uses parameters, but the Python function is "
             + "missing a second 'parameters' argument")
    ));
  }

  @Test
  public void pythonErrorWithBadNumberOfArgsInScript() {
    runCommand.modelId = "raw-params";
    // confused.py just has the plain wrong number of args
    // we give the user a more targeted message here - it's not just a case of forgetting the parameters
    runCommand.parameters = List.of("script=confused.py", "parameters={foo: 'bar'}");

    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());
    String errors = render(ex.getProblem());
    assertThat(errors, allOf(
        containsString("Failed to load your CPython script for `python` step on line 2"),
        containsString("Your python() pipeline step is passing a different number of arguments "
                + "to what your Python script expects."),
        containsString("Your Python function takes 3 positional arguments, but the python() "
                + "pipeline step can only pass 1-2 arguments to your Python script"),
        containsString("The first argument is the rows of pipeline data (via a generator function)."),
        containsString("A second optional argument can pass through parameters from your pipeline python() step")
    ));
  }

  @Test
  public void canRunAModelThatUsesThePythonStepToYieldAllRows() throws Exception {
    doDataframeEcho("event-loss");
  }

  @Test
  public void canRunAModelThatUsesThePythonStepToYieldAllRowsWithLookupResult() throws Exception {
    doDataframeEcho("lookup(\'event-loss\')");
  }

  private void doDataframeEcho(String resultType) throws Exception {
    runCommand.modelId = "dataframe-echo";
    runCommand.parameters = List.of(String.format("results_type='%s'", resultType));
    runCommand.run();

    List<List<String>> rows = openCsv("python.csv",
        "loss", "event", "probability"
    );

    assertThat(rows, contains(
        contains("1000", "1", "0.01"),
        contains("2000", "2", "0.1"),
        contains("4000", "1", "0.2"),
        contains("8000", "2", "0.3")
    ));
  }

  @Test
  public void canRunAModelThatDoesNotProcessAnyRows() throws Exception {
    // A bit of a weird case, but I could see it being useful for debugging or cleanup tasks?
    runCommand.modelId = "no-rows";
    runCommand.parameters = List.of();

    runCommand.run();

    assertThat(openFile("foo.txt"), equalTo("foo"));
  }

  @Test
  public void failsGracefullyIfFileNotFound() throws Exception {
    runCommand.modelId = "not-found";
    runCommand.parameters = List.of();

    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());

    assertThat(ex.getProblem(),
        hasAncestorProblem(equalTo(
            ResourceProblems.get().notFound(
                stdhome().resolve("does-not-exist.py").toUri()
            )
        ))
    );
  }

  @Test
  public void scriptsCanAddArbitraryFilesToTheModelsOutputs() throws Exception {
    runCommand.modelId = "outputs";
    runCommand.parameters = List.of();
    runCommand.run();

    // Check that the reported file location is in the output directory - not where it was originally saved.
    assertThat(outBytes.toString(), containsString(runCommand.runnerOptions.output));

    assertThat(openFile("foo.csv"), equalTo("4"));
  }

  @Test
  public void warningDisplayedIfNoResultTypeOrModelOutputCall() {
    runCommand.modelId = "useless";
    runCommand.parameters = List.of();
    runCommand.run();

    assertThat(collectedSinkProblems,
        contains(CPythonProblems.get().noOutput().withSeverity(Problem.Severity.WARNING)));
  }

  @Test
  public void niceErrorIfYieldWithNoResultType() {
    runCommand.modelId = "no-result-type";
    runCommand.parameters = List.of();
    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());

    assertThat(
        render(ex.getProblem()),
        containsString("No result-type was specified, but your function yielded values")
    );
  }

  @Test
  public void warningIfSameOutputFileRegisteredTwice() {
    runCommand.modelId = "name-reuse";
    runCommand.run();

    assertThat(collectedSinkProblems,
        contains(OutputProblems.get().outputAlreadyExists("foo.txt")));
  }

  @Test
  public void warningIfSameNameInDifferentSubdirectories() {
    runCommand.modelId = "subfolders";
    runCommand.run();

    // See comment in BasePipelineOutputContainer.registerLocalFile()

    assertThat(collectedSinkProblems,
        contains(OutputProblems.get().outputAlreadyExists("foo.txt")));
  }

  @Test
  public void handleCollisionBetweenNormalAndPythonOutput() throws Exception {
    runCommand.modelId = "reuse-with-normal-output";
    runCommand.run();

    // No warning here, there's no risk of the user overwriting a normal output
    // with the python output (or the other way around). We just add a suffix
    // if needed.

    assertThat(outBytes.toString(), allOf(
        containsString("foo.csv"),
        containsString("foo-1.csv")
    ));

    assertThat(openFile("foo.csv"), containsString("loss,event,probability"));
    assertThat(openFile("foo-1.csv"), containsString("4"));
  }

  @Test
  public void modelRunRepeatedWithSetOutputFolder() throws Exception {
    runCommand.modelId = "outputs";
    runCommand.run();

    assertThat(outBytes.toString(), containsString("foo.csv"));

    runCommand.run();
    assertThat(outBytes.toString(), containsString("foo-1.csv"));

    // Check that both outputs are there at the end
    assertThat(openFile("foo.csv"), containsString("4"));
    assertThat(openFile("foo-1.csv"), containsString("4"));
  }

  @Test(timeout = 1000 * 30)
  public void canRunWithLimitedThreads() {

    // We had an issue where the python step would block, causing the engine to hang
    // if we didn't have more threads than concurrent python steps

    runCommand.getEngine().getPipelineExecutor().setNumThreads(2);
    runCommand.modelId = "two-python-steps";

    runCommand.run();
  }

  @Test
  public void doesNotLoseAnyRows() throws Exception {

    // Make sure that we don't somehow skip an output from python when we're reading them back in

    runCommand.modelId = "many-rows";
    runCommand.run();

    assertThat(openFile("python.csv").split("\n"), arrayWithSize(1001));

  }

  @Test
  public void unsupportedTypesInInputIsProblem() throws Exception {
    runCommand.modelId = "unsupported-input";
    ExitException ex = assertThrows(ExitException.class, () -> runCommand.run());
    List<String> errors = Lists.newArrayList(render(ex.getProblem()).split("\n"));
    assertThat(errors, contains(
        is("Failed to validate model for execution"),
        containsString("Failed to validate `python` step on line 2 for execution"),
        containsString("Input from step 'input' includes type(s) that can not be sent to "
            + "CPython: [linearcontinuousfunctiontype]. Try removing"),
        containsString("Step 'input' produces {curve=>ContinuousCurve(xvalues=[1.0, 2.0, 3.0], returnType=Floating)}")
    ));
  }
}
