/*
 * @(#)SshTask.java
 *
 * Copyright (c) 2001-2002 JScape
 * 1147 S. 53rd Pl., Mesa, Arizona, 85206, U.S.A.
 * All rights reserved.
 *
 * This software is the confidential and proprietary information of
 * JScape. ("Confidential Information").  You shall not disclose such
 * Confidential Information and shall use it only in accordance with
 * the terms of the license agreement you entered into with JScape.
 */

package com.jscape.inet.ssh;


import com.jscape.util.Clock;

import java.io.*;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Timer;
import java.util.TimerTask;
import java.util.UUID;
import java.util.regex.Pattern;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * The SshTask class represents a command to be executed in batch (non-interactive) mode on SSH server.
 * In order to execute a command in batch mode the SshTask must know the following:
 * <p/>
 * Start Prompt - The start prompt is the data that is sent back by the SSH server indicating that it is ready to
 * execute the command. e.g. $
 * <p/>
 * Message - The command is the command to be executed. e.g. ls -al
 * <p/>
 * End Prompt - The end prompt is the data that is sent back by the SSH server indicating that it has completed execution
 * of the command.  In most cases this is the same as the Start Prompt.  e.g. $
 * <p/>
 * When creating a new SshTask you must at a MINIMUM provide both the Message and End Prompt.  In most cases a Start Prompt should be
 * provided although you may provide a <code>null</code> value in cases where you are not waiting for SSH server data but want the
 * command to be executed immediately.
 */
public class SshTask
        implements Serializable {

    /**
     * SshTask is waiting to be activated
     */
    public static final int PENDING = 0;

    /**
     * SshTask is active and waiting for start prompt to execute command
     */
    public static final int ACTIVE = 1;

    /**
     * SshTask has received start prompt and is currently executing command
     */
    public static final int RUNNING = 2;

    /**
     * SshTask has finished execution of command and has received end prompt
     */
    public static final int COMPLETE = 3;

    private static final String[] STATUS_STRINGS = {"Pending", "Active", "Running", "Complete"};

    private static final long serialVersionUID = 1141550602041941125L;

    private static final String DEFAULT_EOL = "\n";

    private static final long DEFAULT_TIMEOUT_MILLIS = 120000;

    private static final int DEFAULT_REGEX_MATCH_LENGTH = 20;

    private static final PrintStream DEFAULT_DEBUG_STREAM = System.out;

    private static final int TIMER_PERIOD = 500;

    private String name;
    private String promptPattern;
    private boolean promptPatternIsRegex;
    private String command;
    private String responsePattern;
    private boolean responsePatternIsRegex;
    private int regexMatchLength;
    private String endOfLine;
    private String characterSet;
    private long promptTimeout;
    private long responseTimeout;
    private boolean debug;
    private transient PrintStream debugStream;
    private transient Listener listener;
    private transient OutputStream out;
    private StringBuilder promptBuffer;
    private StringBuilder responseBuffer;
    private Pattern promptRegexpPattern;
    private Pattern responseRegexpPattern;
    private volatile int status;
    private volatile Exception error;
    private volatile boolean timedOut;
    private transient long timeoutDate;
    private transient Timer timer;

    public SshTask(String name,
                   String promptPattern,
                   boolean promptPatternIsRegex,
                   String command,
                   String responsePattern,
                   boolean responsePatternIsRegex,
                   int regexMatchLength,
                   String endOfLine,
                   String characterSet,
                   long promptTimeout,
                   long responseTimeout) {
        assertNotNull(name);
        this.name = name;

        this.promptPattern = promptPattern;
        this.promptPatternIsRegex = promptPatternIsRegex;
        this.command = command;

        assertNotNull(responsePattern);
        this.responsePattern = responsePattern;

        this.responsePatternIsRegex = responsePatternIsRegex;
        this.regexMatchLength = regexMatchLength;

        assertNotNull(endOfLine);
        this.endOfLine = endOfLine;

        this.characterSet = characterSet;

        this.promptTimeout = promptTimeout;
        this.responseTimeout = responseTimeout;

        this.debugStream = DEFAULT_DEBUG_STREAM;
        this.status = PENDING;
        this.promptBuffer = new StringBuilder();
        this.responseBuffer = new StringBuilder();
    }

    public SshTask(String name, String promptPattern, boolean promptPatternIsRegex, String command, String responsePattern, boolean responsePatternIsRegex) {
        this(name, promptPattern, promptPatternIsRegex, command, responsePattern, responsePatternIsRegex, DEFAULT_REGEX_MATCH_LENGTH, DEFAULT_EOL, null, DEFAULT_TIMEOUT_MILLIS, DEFAULT_TIMEOUT_MILLIS);
    }

    /**
     * Constructs a new <code>SshTask</code> instance.
     *
     * @param name        a name to uniquely identify this <code>SshTask</code>.
     * @param promptPattern the shell prompt this task should wait for before executing command.
     * @param command     the command to execute
     * @param responsePattern   the end shell prompt this task should wait for to indicate that this task is finished.
     */
    public SshTask(String name, String promptPattern, String command, String responsePattern) {
        this(name, promptPattern, false, command, responsePattern, false);
    }

    /**
     * Constructs a new <code>SshTask</code> instance.
     *
     * @param promptPattern the start data this task should wait for before executing command.
     * @param command     the command to execute
     * @param responsePattern   the end data this task should wait for to indicate that this task is finished.
     */
    public SshTask(String promptPattern, String command, String responsePattern) {
        this(UUID.randomUUID().toString(), promptPattern, command, responsePattern);
    }

    /**
     * Gets name of this <code>SshTask</code> as defined in constructor.
     *
     * @return name of this <code>SshTask</code>, <code>Unknown</code> if not defined.
     */
    public String getName() {
        return this.name;
    }

    /**
     * Gets the shell prompt this <code>SshTask</code> should wait for before executing command.
     *
     * @return the start shell prompt
     */
    public String getStartPrompt() {
        return this.promptPattern;
    }

    /**
     * Specified if start prompt is a regular expression
     *
     * @return true if is a regular expression, false otherwise
     */
    public boolean isStartPromptRegex() {
        return this.promptPatternIsRegex;
    }

    /**
     * Specifies if start prompt is a regular expression
     *
     * @param value true if start prompt is a regular expression, false otherwise
     *
     * @throws SshException if the start prompt is an invalid regular expression
     * @see #setRegexMatchLength
     */
    public void setStartPromptRegex(boolean value) {
        this.promptPatternIsRegex = value;
    }

    /**
     * Gets the command to execute
     *
     * @return the command to execute
     */
    public String getCommand() {
        return this.command;
    }

    /**
     * Gets the end shell prompt this <code>SshTask</code> should wait for to indicate that this <code>SshTask</code> is complete.
     *
     * @return the end shell prompt
     */
    public String getEndPrompt() {
        return this.responsePattern;
    }

    /**
     * Checks if end prompt is a regular expression
     *
     * @return true if is a regular expression, false otherwise
     */
    public boolean isEndPromptRegex() {
        return this.responsePatternIsRegex;
    }

    /**
     * Specifies if end prompt is a regular expression
     *
     * @param value true if end prompt is a regular expression, false otherwise
     *
     * @throws SshException if the end prompt is an invalid regular expression
     * @see #setRegexMatchLength
     */
    public void setEndPromptRegex(boolean value) {
        this.responsePatternIsRegex = value;
    }

    /**
     * Gets the number of characters to read from the end of the response buffer when using regular expressions
     * to detect start or end prompts.  Default value is 20.
     *
     * @return match length
     */
    public int getRegexMatchLength() {
        return this.regexMatchLength;
    }

    /**
     * Sets the number of characters to read from the end of the response buffer when using regular expressions
     * to detect start or end prompts.  Default value is 20.
     *
     * @param value
     */
    public void setRegexMatchLength(int value) {
        this.regexMatchLength = value;
    }

    /**
     * Gets line terminator used when sending commands.
     * Default is LF (Line Feed).
     *
     * @return line terminator
     */
    public String getTerminator() {
        return this.endOfLine;
    }

    /**
     * Sets line terminator used when sending commands.
     * Default is LF (Line Feed).
     *
     * @param value line terminator
     */
    public void setTerminator(String value) {
        assertNotNull(value);
        this.endOfLine = value;
    }

    /**
     * Sets character set used when write data
     *
     * @param value the character set
     */
    public void setCharacterSet(String value) {
        this.characterSet = value;
    }

    /**
     * Gets the character set
     *
     * @return The character set used
     */
    public String getCharacterSet() {
        return this.characterSet;
    }

    /**
     * Gets the maximum timeout between the time the task is in active state and the expected start prompt is received.
     *
     * @return the timeout in milliseconds
     */
    public long getStartPromptTimeout() {
        return promptTimeout;
    }

    /**
     * Sets the maximum timeout between the time the task is in active state and the expected start prompt is received.
     *
     * @param value the timeout in milliseconds
     */
    public void setStartPromptTimeout(long value) {
        this.promptTimeout = value;
    }

    /**
     * Gets the maximum timeout between the time the task is in running state and the expected end prompt is received.
     *
     * @return the timeout in milliseconds
     */
    public long getEndPromptTimeout() {
        return responseTimeout;
    }

    /**
     * Sets the maximum timeout between the time the task is in running state and the expected end prompt is received.
     *
     * @param value the timeout in milliseconds
     */
    public void setEndPromptTimeout(long value) {
        this.responseTimeout = value;
    }

    /**
     * Checks whether debugging is enabled for this task.
     *
     * @return true if enabled, false otherwise
     */
    public boolean getDebug() {
        return this.debug;
    }

    /**
     * Enables/disables additional debugging for this task.
     *
     * @param value true if debugging is enabled, false otherwise
     */
    public void setDebug(boolean value) {
        this.debug = value;
    }

    /**
     * Sets debug stream to be used when printing debug messages.
     * Default debug stream is System.out
     *
     * @param value a PrintStream
     */
    public void setDebugStream(PrintStream value) {
        assertNotNull(value);
        this.debugStream = value;
    }

    /**
     * Gets debug stream used when printing debug messages.
     * Default debug stream is System.out
     *
     * @return a PrintStream
     */
    public PrintStream getDebugStream() {
        return this.debugStream;
    }

    /**
     * Gets completion status of this <code>SshTask</code>.
     *
     * @return the completion status of this <code>SshTask</code>
     *
     * @see #PENDING
     * @see #ACTIVE
     * @see #RUNNING
     * @see #COMPLETE
     */
    public int getStatus() {
        return this.status;
    }

    /**
     * Gets status string based on task status
     *
     * @param status the status code
     *
     * @return the status
     *
     * @see #PENDING
     * @see #ACTIVE
     * @see #RUNNING
     * @see #COMPLETE
     */
    public String getStatusString(int status) {
        return 0 <= status && status < STATUS_STRINGS.length ? STATUS_STRINGS[status] : "Unknown";
    }

    /**
     * Checks if this <code>SshTask</code> is pending activation.
     *
     * @return <code>true</code> if pending activation, <code>false</code> otherwise
     */
    public boolean isPending() {
        return this.status == PENDING;
    }

    /**
     * Checks if this <code>SshTask</code> is active.
     *
     * @return <code>true</code> if active, <code>false</code> otherwise
     */
    public boolean isActive() {
        return this.status == ACTIVE;
    }

    /**
     * Checks if this <code>SshTask</code> is running.
     *
     * @return <code>true</code> if running, <code>false</code> otherwise
     */
    public boolean isRunning() {
        return this.status == RUNNING;
    }

    /**
     * Checks if this <code>SshTask</code> is complete.
     *
     * @return <code>true</code> if complete, <code>false</code> otherwise
     */
    public boolean isComplete() {
        return this.status == COMPLETE;
    }

    /**
     * Gets response to command executed.  Response will be empty if command has not executed.
     *
     * @return the response to command executed
     *
     * @see #isComplete
     */
    public String getResponse() {
        return this.responseBuffer.toString();
    }

    /**
     * Checks if task has exceeded timeout
     *
     * @return true if timeout exceeded, false otherwise
     */
    public boolean isTimedOut() {
        return this.timedOut;
    }

    public Exception getError() {
        return this.error;
    }

    public synchronized void activate(Listener listener, OutputStream out, String previousTaskResponse, String endOfLine) {
        assertNotNull(listener);
        this.listener = listener;

        assertNotNull(out);
        this.out = out;

        if (endOfLine != null) {
            this.endOfLine = endOfLine;
        }

        this.promptBuffer = new StringBuilder(previousTaskResponse);
        this.responseBuffer = new StringBuilder();
        this.promptRegexpPattern = this.promptPatternIsRegex && this.promptPattern != null ? Pattern.compile(this.promptPattern) : null;
        this.responseRegexpPattern = this.responsePatternIsRegex && this.responsePattern != null ? Pattern.compile(this.responsePattern) : null;

        this.timedOut = false;
        this.error = null;

        status(ACTIVE);

        this.timeoutDate = Clock.after(this.promptTimeout);
        startTimeoutTimer();

        if (promptMatched()) {
            debug(this.promptPattern != null ? String.format("Found start prompt matching \"%s\" in response \"%s\".", this.promptPattern, this.promptBuffer) : "");
            sendCommand();
        }
    }

    public synchronized void onData(char data) {
        if (isActive()) {
            onPromptData(data);
        } else if (isRunning()) {
            onResponseData(data);
        }
    }

    public synchronized void onEof() {
        this.error = new EOFException();
        status(COMPLETE);
        stopTimeoutTimer();
        this.listener.onTaskError(this, new Exception("Unexpected EOF reached."));
    }

    /**
     * Cancels this task and subsequent tasks by setting it's status to complete.
     */
    public void cancel() {
        synchronized (this) {
            if (isComplete()) {
                return;
            }

            debug(String.format("Task running command \"%s\" has been cancelled.", this.command != null ? this.command : "<none>"));
            status(COMPLETE);

            stopTimeoutTimer();
        }

        if (this.listener != null) {
            this.listener.onTaskCancelled(this);
        }
    }

    private void onTimer() {
        if (isComplete()) {
            synchronized (this) {
                stopTimeoutTimer();
            }
        } else if (Clock.isAfter(this.timeoutDate)) {
            synchronized (this) {
                this.timedOut = true;
                status(COMPLETE);
                stopTimeoutTimer();
            }
            this.listener.onTaskTimeout(this);
        }
    }

    private void sendCommand() {
        if (this.command == null) {
            taskCompleted();
            return;
        }

        debug(String.format("Executing command \"%s\"...", this.command));
        status(RUNNING);

        String fullCommand = this.command + this.endOfLine;

        try {
            this.out.write(this.characterSet != null ? fullCommand.getBytes(Charset.forName(this.characterSet)) : fullCommand.getBytes(StandardCharsets.UTF_8));
            this.out.flush();
            this.timeoutDate = Clock.after(this.responseTimeout);
        } catch (Exception e) {
            status(COMPLETE);
            stopTimeoutTimer();
            this.listener.onTaskError(this, e);
        }
    }

    private void onPromptData(char data) {
        this.promptBuffer.append(data);

        if (promptMatched()) {
            debug(String.format("Found start prompt matching \"%s\" in response \"%s\".", this.promptPattern, this.promptBuffer));
            sendCommand();
        }
    }

    private void onResponseData(char data) {
        this.responseBuffer.append(data);

        if (responseMatched()) {
            debug(String.format("Found end prompt matching \"%s\" fin response \"%s\".", this.responsePattern, this.responseBuffer));
            taskCompleted();
        }
    }

    private boolean promptMatched() {
        if (this.promptPattern == null) {
            return true;
        } else if (this.promptRegexpPattern != null) {
            String data = this.promptBuffer.substring(this.promptBuffer.length() > this.regexMatchLength ? this.promptBuffer.length() - this.regexMatchLength : 0);
            return data.length() > 0 && this.promptRegexpPattern.matcher(data).find();
        } else {
            String data = this.promptBuffer.substring(this.promptBuffer.length() > this.promptPattern.length() ? this.promptBuffer.length() - this.promptPattern.length() : 0);
            return data.length() > 0 && data.equals(this.promptPattern);
        }
    }

    private boolean responseMatched() {
        if (this.responseRegexpPattern != null) {
            String data = this.responseBuffer.substring(this.responseBuffer.length() > this.regexMatchLength ? this.responseBuffer.length() - this.regexMatchLength : 0);
            return data.length() > 0 && this.responseRegexpPattern.matcher(data).find();
        } else {
            String data = this.responseBuffer.substring(this.responseBuffer.length() > this.responsePattern.length() ? this.responseBuffer.length() - this.responsePattern.length() : 0);
            return data.length() > 0 && data.equals(this.responsePattern);
        }
    }

    private void taskCompleted() {
        status(COMPLETE);
        stopTimeoutTimer();
        this.listener.onTaskCompleted(this);
    }

    private void status(int value) {
        this.status = value;
        debug(String.format("Changed task status running command \"%s\" to \"%s\".", this.command != null ? this.command : "<none>", STATUS_STRINGS[this.status]));
    }

    private void startTimeoutTimer() {
        stopTimeoutTimer();

        this.timer = new Timer(String.format("SSH Task '%s' Timer", this.name), true);
        this.timer.schedule(new TimerTask() {
            @Override
            public void run() {
                onTimer();
            }
        }, TIMER_PERIOD, TIMER_PERIOD);
    }

    private void stopTimeoutTimer() {
        if (this.timer != null) {
            this.timer.cancel();
            this.timer = null;
        }
    }

    private void debug(String message) {
        if (this.debug && !message.isEmpty()) {
            this.debugStream.println(String.format("Task: %s. %s", this.name, message));
        }
    }

    private void readObject(ObjectInputStream stream)
            throws IOException,
                   ClassNotFoundException {
        stream.defaultReadObject();
        this.debugStream = DEFAULT_DEBUG_STREAM;
    }

    public interface Listener {

        void onTaskCompleted(SshTask task);

        void onTaskCancelled(SshTask task);

        void onTaskError(SshTask task, Exception error);

        void onTaskTimeout(SshTask task);
    }

}
