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

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import com.google.common.collect.Lists;

import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.bind.Parameter;
import nz.org.riskscape.engine.cli.ApplicationCommand;
import nz.org.riskscape.engine.cli.ExitException;
import nz.org.riskscape.engine.cli.pipeline.CliPipelineRunner;
import nz.org.riskscape.engine.model.Model;
import nz.org.riskscape.engine.pipeline.ExecutionContext;
import nz.org.riskscape.engine.pipeline.RealizedPipeline;
import nz.org.riskscape.picocli.CommandLine.Command;
import nz.org.riskscape.picocli.CommandLine.Option;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problems;

@Command(name = "batch")
public class BatchCommand extends BaseModelRunCommand {

  public enum BatchFailureMode {
    WARN, //Warns if a batch error occurs, but continue the batch run
    ABORT, //Aborts that batch run if a batch error occurs
    EXIT_ERROR; //Exits with error if a batch error has occured, but all batches are attempted
  }

  @RequiredArgsConstructor
  public static class Vary {

    final Parameter modelParameter;
    final String replacementPattern;
  }

  public static final String VARY_TOKEN = "{}";
  public static final String VARY_TOKEN_PATTERN = Pattern.quote(VARY_TOKEN);

  @Option(names = {"--vary-parameter"}, required = false)
  public String varyParameter;

  @Option(names = {"--vary-input"}, required = true)
  public String varyInputString;

  @Option(names = {"--failure-mode"}, required = false)
  public BatchFailureMode failureMode = BatchFailureMode.WARN;

  @Override
  public Object doCommand(Project useProject) {
    if (runnerOptions.output != null) {
      // we check that the output will be usable as a URI. Of course we need to replace any '{}' values
      // first because they aren't allowed in a URI.
      String mockOutputPath = runnerOptions.output.replaceAll(VARY_TOKEN_PATTERN, "VARY-VALUE");
      bindOutput(mockOutputPath);
    }

    int batchesToRun = 0, batchesAttempted = 0, batchesSuccessful = 0, batchesFailed = 0;

    // we start by getting the model and populating it with the standard parameter values from
    // --parameter and --parameters arguments
    Model model = ApplicationCommand.getObject(useProject.getIdentifiedModels(), modelId, getTerminal()).getModel();

    Map<String, List<String>> cliParameterMap = getCliParameterMap();
    Vary varying = findVary(model, cliParameterMap);
    if (!(parameters.isEmpty() && parametersFile == null)) {
      model = updateParameters(model, cliParameterMap);
    }

    List<String> varyInputs = processVaryInput();
    batchesToRun = varyInputs.size();

    // now we make a specific model for each batch.
    List<Problem> parameterErrors = new ArrayList<>();
    List<Model> batches = Lists.newArrayListWithCapacity(varyInputs.size());
    for (String varyInput : varyInputs) {
      Map<String, List<String>> varyParameterMap = new HashMap<>();
      varyParameterMap.put(varying.modelParameter.getName(), Arrays.asList(
          varying.replacementPattern.replaceAll(VARY_TOKEN_PATTERN, Matcher.quoteReplacement(varyInput)))
      );
      try {
        batches.add(updateParameters(model, varyParameterMap));
      } catch (ExitException e) {
        parameterErrors.add(Problems.foundWith(varyInput, e.getProblem()));
      }
    }
    if (!parameterErrors.isEmpty()) {
      throw new ExitException(Problems.foundWith("--vary-parameter", parameterErrors));
    }

    // Stash the time that the batch job was started. We might need this to build a sensible output location later.
    String batchStartTimestamp = CliPipelineRunner.createDirFriendlyTimestamp(getCurrentTime());

    // Start setting up for the model runs

    // we are going to mutate this in the loop below, but need to remember what the original state was
    final boolean outputSwitchGiven = runnerOptions.output != null;

    // the out pattern is set when the --output switch contains a '{}' - this tells us that we need to change the output
    // per run based on the vary input
    final String outPattern;
    if (outputSwitchGiven && runnerOptions.output.contains(VARY_TOKEN)) {
      outPattern = runnerOptions.output;
    } else {
      outPattern = null;
    }

    for (int i = 0; i < batches.size(); i++) {
      ++batchesAttempted;
      Model toRun = batches.get(i);

      // NB this command runs many models - each one needs a new executionContext (so don't use command provided one)
      try (ExecutionContext executionContext = useProject.newExecutionContext()) {
        CliPipelineRunner runner = new CliPipelineRunner(executionContext, this);
        RealizedPipeline realizedPipeline = realize(toRun, executionContext);

        if (i == 0 && runnerOptions.format != null) {
          // we only want the format override warning to come out for the first batch. It's just noise
          // after that.
          warnIfCannotOverrideFormat(realizedPipeline, useProject.getEngine().getProblemSink());
        }

        if (outputSwitchGiven) {
          // if output has been specified as a switch, we need to mung the output using the current vary input
          if (outPattern != null) {
            String pvAsBasename = com.google.common.io.Files.getNameWithoutExtension(varyInputs.get(i));
            runnerOptions.output = outPattern.replaceAll(VARY_TOKEN_PATTERN, pvAsBasename);
          }
          // the implicit else case here is that everything is going to pile in to the same directory.  If the user has
          // given --replace, it's going trample over itself, but otherwise it'll create enumerated files in the same
          // dir, e.g. event-table-1.csv etc
        } else {
          // if the output dir was not been specified we want a better default than the one the normal model run
          // commands give - We make a directory (similar to default cli pipeline runner) but add `batchx` to the end.
          // note that we add one to the index because human's tend to be one-based
          runnerOptions.output = useProject.getOutputBaseLocation()
              .resolve(String.format("%s/%s/batch%d/", modelId, batchStartTimestamp, i + 1)).toString();
        }

        runner.run(realizedPipeline, executionContext.getPipelineExecutor(), useProject, runnerOptions);
        ++batchesSuccessful;
      } catch (RiskscapeException ex) {
        ++batchesFailed;
        stderr().format("Could not execute batch %d of %d%n", batchesAttempted, batchesToRun);
        stderr().format("  Vary parameter value: %s%n", varyInputs.get(i));
        stderr().format("  Reason: %s%n%n", getMessages().renderProblem(ex));
        if (failureMode == BatchFailureMode.ABORT) {
          throw new ExitException(1, "Batch failed with --vary-parameter value: %s", varyInputs.get(i));
        }
      }
    }

    stdout().format("Batch complete. Success %d of %d%n", batchesSuccessful, batchesToRun);
    if (batchesFailed > 0) {
      String message = String.format("There were %d batch(s) that did not complete successfully", batchesFailed);
      if (failureMode == BatchFailureMode.EXIT_ERROR) {
        throw new ExitException(1, message);
      }
      stderr().println(message);
    }

    return null;
  }

  /**
   * Looks through given parameters to decide what to vary between runs
   */
  private Vary findVary(Model model, Map<String, List<String>> cliParameterMap) {
    // still none? give up
    String paramName = varyParameter;
    String varyPattern;

    if (paramName == null) {
      List<Entry<String, List<String>>> withToken = cliParameterMap.entrySet().stream()
          .filter(entry -> {
        return entry.getValue().stream().filter(v -> v.contains(VARY_TOKEN)).findFirst().isPresent();
      }).collect(Collectors.toList());

      if (withToken.size() == 1) {
        Entry<String, List<String>> entry = withToken.get(0);
        if (entry.getValue().size() != 1) {
          throw new ExitException(1, "The batch command does not support being run against parameters with more "
              + "than one argument.  Parameter %s has been supplied %d time(s)", entry.getKey(), entry.getValue());
        }

        paramName = entry.getKey();
        varyPattern = entry.getValue().get(0);
        cliParameterMap.remove(entry.getKey());
      } else if (withToken.isEmpty()) {
        throw new ExitException(1, "--vary-parameter not set, and no %s pattern found in any parameters", VARY_TOKEN);
      } else {
        throw new ExitException(1, "Only one parameter may contain the vary-input pattern `%s`", VARY_TOKEN);
      }
    } else {
      varyPattern = VARY_TOKEN;
    }

    final String toFind = paramName;
    Parameter param = model.getBoundParameters().getBoundTo().getDeclared()
        .stream()
        .filter(p -> p.getName().equals(toFind))
        .findFirst()
        .orElseThrow(()
            -> new ExitException(
            1,
            "--vary-parameter '%s' does not exist on model %s - did you mean one of %s?",
            varyParameter,
            modelId,
            model.getBoundParameters().getBoundTo().getDeclared()
                .stream()
                .map(Parameter::getName)
                .collect(Collectors.toList())));

    return new Vary(param, varyPattern);
  }

  private List<String> processVaryInput() {
    List<String> beingBuilt;
    try {
      File varyInput = new File(varyInputString);
      if (varyInput.isDirectory()) {
        //Harvest files from it.
        beingBuilt = Files.list(varyInput.toPath())
            .map(Path::toUri)
            .map(p -> p.toString())
            .collect(Collectors.toList());
        beingBuilt.sort(String::compareTo);
      } else if (!containsGlobs(varyInputString)) {
        //If it is not a directory and does not contain globs we expect it to be a file that exists
        if (!varyInput.exists()) {
          throw new ExitException(1, String.format("vary-input file: %s does not exist", varyInput));
        }
        //Vary input is a file. We harvest from file content
        if (!(varyInput.canRead() && varyInput.isFile())) {
          throw new ExitException(1,
              String.format("vary-input file: %s cannot be read is or not a file", varyInputString));
        }
        //Read the vary file. Empty lines are removed.
        beingBuilt = Files.readAllLines(varyInput.toPath())
            .stream().filter(l -> !"".equals(l.trim())).collect(Collectors.toList());
      } else {
        //Not a directory or file.
        String namePart = varyInput.getName();
        //Convert namePart to a regex. First escape any dots, then convert our * glob
        String namePattern = namePart.replace(".", "\\.").replace("*", ".*").replace(VARY_TOKEN, ".*");
        File parent = varyInput.getParentFile();
        if (parent == null) {
          //parent will be null if there was no path component in the input "./" and the file doesn't exist
          //which is expected if it is a glob pattern
          parent = new File("");
        }
        beingBuilt = Files.list(parent.toPath())
            .filter(p -> p.toFile().isFile())
            .filter(p -> p.toFile().getName().matches(namePattern))
            .sorted()
            .map(Path::toUri)
            .map(p -> p.toString()).collect(Collectors.toList());
      }

    } catch (IOException e) {
      throw new ExitException(1, String.format("Could not read vary-input file: %s", varyInputString), e);
    }
    if (beingBuilt.isEmpty()) {
      throw new ExitException(1, String.format("vary-input: %s is empty", varyInputString));
    }
    return beingBuilt;
  }

  private boolean containsGlobs(String toCheck) {
    return toCheck.contains("*") || toCheck.contains(VARY_TOKEN);
  }

}
