/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.elasticsearch.gradle.vagrant;

import org.gradle.api.DefaultTask;
import org.gradle.api.tasks.Input;
import org.gradle.api.tasks.TaskAction;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;

import static org.elasticsearch.gradle.vagrant.VagrantMachine.convertLinuxPath;
import static org.elasticsearch.gradle.vagrant.VagrantMachine.convertWindowsPath;

/**
 * A shell script to run within a vagrant VM.
 *
 * The script is run as root within the VM.
 */
public abstract class VagrantShellTask extends DefaultTask {

    private final VagrantExtension extension;
    private final VagrantMachine service;
    private UnaryOperator<String> progressHandler = UnaryOperator.identity();

    public VagrantShellTask() {
        extension = getProject().getExtensions().findByType(VagrantExtension.class);
        if (extension == null) {
            throw new IllegalStateException("elasticsearch.vagrant-base must be applied to create " + getClass().getName());
        }
        service = getProject().getExtensions().getByType(VagrantMachine.class);
    }

    @Input
    protected abstract List<String> getWindowsScript();

    @Input
    protected abstract List<String> getLinuxScript();

    @Input
    public UnaryOperator<String> getProgressHandler() {
        return progressHandler;
    }

    public void setProgressHandler(UnaryOperator<String> progressHandler) {
        this.progressHandler = progressHandler;
    }

    @TaskAction
    public void runScript() {
        String rootDir = getProject().getRootDir().toString();
        if (extension.isWindowsVM()) {
            service.execute(spec -> {
                spec.setCommand("winrm");

                List<String> script = new ArrayList<>();
                script.add("try {");
                script.add("cd " + convertWindowsPath(getProject(), rootDir));
                extension.getVmEnv().forEach((k, v) -> script.add("$Env:" + k + " = \"" + v + "\""));
                script.addAll(getWindowsScript().stream().map(s -> "    " + s).collect(Collectors.toList()));
                script.addAll(Arrays.asList(
                    "    exit $LASTEXITCODE",
                    "} catch {",
                    // catch if we have a failure to even run the script at all above, equivalent to set -e, sort of
                    "    echo $_.Exception.Message",
                    "    exit 1",
                    "}"));
                spec.setArgs("--elevated", "--command", String.join("\n", script));
                spec.setProgressHandler(progressHandler);
            });
        } else {
            service.execute(spec -> {
                spec.setCommand("ssh");

                List<String> script = new ArrayList<>();
                script.add("sudo bash -c '"); // start inline bash script
                script.add("pwd");
                script.add("cd " + convertLinuxPath(getProject(), rootDir));
                extension.getVmEnv().forEach((k, v) -> script.add("export " + k + "=" + v));
                script.addAll(getLinuxScript());
                script.add("'"); // end inline bash script
                spec.setArgs("--command", String.join("\n", script));
                spec.setProgressHandler(progressHandler);
            });
        }
    }


}

@@ -3471,16 +3471,38 @@
  else %7B%0A
+            try %7B%0A    
         
@@ -3523,32 +3523,36 @@
 ecute(spec -%3E %7B%0A
+    
                 
@@ -3584,32 +3584,36 @@
                 
+    
 List%3CString%3E scr
@@ -3629,32 +3629,36 @@
  ArrayList%3C%3E();%0A
+    
                 
@@ -3723,32 +3723,36 @@
                 
+    
 script.add(%22pwd%22
@@ -3750,24 +3750,28 @@
 add(%22pwd%22);%0A
+    
             
@@ -3843,32 +3843,36 @@
                 
+ 
+   
 extension.getVmE
@@ -3940,32 +3940,36 @@
                 
+    
 script.addAll(ge
@@ -3993,32 +3993,36 @@
                 
+ 
+   
 script.add(%22'%22);
@@ -4040,32 +4040,36 @@
 ine bash script%0A
+    
                 
@@ -4130,32 +4130,36 @@
                 
+  
+  
 spec.setProgress
@@ -4188,35 +4188,511 @@
 r);%0A            
-%7D);
+    %7D);%0A            %7D catch (Exception e) %7B%0A                getLogger().error(%22Failed command, dumping dmesg%22, e);%0A                service.execute(spec -%3E %7B%0A                    spec.setCommand(%22ssh%22);%0A                    spec.setArgs(%22--command%22, %22dmesg%22);%0A                    spec.setProgressHandler(line -%3E %7B%0A                        getLogger().error(line);%0A                        return null;%0A                    %7D);%0A                %7D);%0A                throw e;%0A            %7D
 %0A        %7D%0A    %7D

Add dmesg dump on failure of vagrant ssh command (#45349)