package nxm.webvnc.lib;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.net.MalformedURLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.logging.Level;
import java.util.stream.Collectors;

import nxm.sys.lib.LogUtilities;
import nxm.sys.lib.Table;
import nxm.sys.lib.TextFile;

/**
 * Handles Opening and closing of VNC sessions, along with the necessary websocket proxy
 *
 * @author adw
 * @since NeXtMidas 4.1.1
 *
 */
public class VNCSessionManager {
  private List<Integer> activeDisplays = new ArrayList<>();
  private List<Integer> externalDisplays = new ArrayList<>();

  private int maxSessions = Integer.MAX_VALUE;

  private boolean ssl = false;
  private String certPath = null;
  private String keyPath = null;
  private String geometry = null;
  private String xstartup = null;
  private String securityTypes = null;

  private boolean manual = false;
  private boolean debug  = false;
  private boolean ansiColors = false;

  /** Instantiates a VNCSessionManager and makes note internally of all currently active
   *  VNC sessions on the host to avoid collisions */
  public VNCSessionManager () {
    // Find any already running vnc displays and track them so we don't collide
    externalDisplays.addAll(getActiveVNCServers());
  }

  /** Instantiates a VNCSessionManager and makes note internally of all currently active
   *  VNC sessions on the host to avoid collisions
   *  @param ansiColors true to use Midas logging with ANSI Colors */
  public VNCSessionManager (boolean ansiColors) {
    this();
    this.ansiColors = ansiColors;
  }

  /**  Set if DEBUG is active
   *   @param debug true to turn on DEBUG level messages
   */
  public void setDebug ( boolean debug) {
    this.debug = debug;
  }

  /** Set the maximum number of allowed concurrent VNC sessions managed by this instance
   *
   *  @param maxSessions the maximum number of sessions to start
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager setMaxSessions (int maxSessions) {
    this.maxSessions = maxSessions;
    return this;
  }

  /** Enable SSL for the VNC Webproxy
   *
   *  @param certPath full path to a certificate file in PEM format
   *  @param keyPath full path to a key file in PEM format
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager enableSSL (String certPath, String keyPath) {
    this.ssl = true;
    this.certPath = certPath;
    this.keyPath = keyPath;
    return this;
  }

  /** Prevents automatic creation of vnc sessions and websocket proxies.
   *  This is useful for debugging, as output from the vncserver and websocket proxy is output to the terminal
   *
   *  @param manual true to enable manual mode, false to disable
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager setManual (boolean manual) {
    this.manual = manual;
    return this;
  }

  /** Set the -geometry argument to the vncserver command. This sets the screen resolution of the desktop presented to the user
   *
   *  @param geometry  [width_pixels]x[height_pixels] i.e. 1280x800
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager setGeometry (String geometry) {
    this.geometry = geometry;
    return this;
  }

  /** Set a custom xstartup file for the vncserver
   *
   *  @param xstartup  the path to the custom xstartup file
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager setXStartup (String xstartup) {
    this.xstartup = xstartup;
    return this;
  }

  /** Set the securitytypes parameter for vncserver. See the xvnc man page for more info
   *
   *  @param securityTypes  the value to set
   *  @return the instance of VNCSessionManager
   */
  public VNCSessionManager setSecurityTypes (String securityTypes) {
    this.securityTypes = securityTypes;
    return this;
  }

  /** Start a VNC session and an associated websocket proxy
   *
   *  @param userDN the DN of the user who requested the session.
   *  @param requestHeader a {@link Table} containing the HTTP request header information.
   *  @return an int id number for the session, or -1 if something went wrong
   */
  public int getVNCSession (String userDN, Table requestHeader) {
    if (debug) LogUtilities.midasLog(null, Level.INFO, LogUtilities.buildMessage("[DEBUG] Get Session"), null);
    updateExternalDisplays();
    if (activeDisplays.size() >= maxSessions) return -1;
    for (int i = 1; i <= maxSessions + externalDisplays.size(); i++) {
      if (!activeDisplays.contains(i) && !externalDisplays.contains(i)) {
        if (startvnc(i, userDN, requestHeader)) {
          if (startwss(i)) {
            activeDisplays.add(i);
            if (debug) LogUtilities.midasLog(null, Level.INFO, LogUtilities.buildMessage("[DEBUG] Session id = {}", i), null);
            return i;
          } else {
            if (debug) LogUtilities.midasLog(null, Level.INFO, LogUtilities.buildMessage("[DEBUG] Failed to start websocket"), null);
            stopvnc(i);
          }
        } else {
          if (debug) LogUtilities.midasLog(null, Level.INFO, LogUtilities.buildMessage("[DEBUG] Failed to start vncserver"), null);
        }
      }
    }
    return -1;
  }

  /** Close the VNC session with the given id and the associated websocket proxy
   *
   *  @param i the id of the session to close
   */
  public void closeVNCSession (int i) {
    if (i == 0) return;
    String msg = LogUtilities.buildMessage("Requested session end for VNC display :{}", i);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    if (stopwss(i) && stopvnc(i)) activeDisplays.remove(Integer.valueOf(i));
  }

  /** Close all active VNC sessions started by this instance and associated websocket proxies
   */
  public void closeAllVNCSessions () {
    List<Integer> toClose = new ArrayList<>(activeDisplays);
    toClose.forEach(i -> closeVNCSession(i));
  }

  /** Is ssl enabled
   *  @return true if enabled */
  boolean isSSL () {
    return ssl;
  }

  private void updateExternalDisplays() {
    List<Integer> active = getActiveVNCServers();
    active.forEach(display -> {
      if (!externalDisplays.contains(display) && !activeDisplays.contains(display)) {
        if (debug) LogUtilities.midasLog(null, Level.INFO, LogUtilities.buildMessage("[DEBUG] Found New External Display {}", display), null);
        externalDisplays.add(display);
      }
    });
  }

  private List<Integer> getActiveVNCServers () {
    String[] command = {"/bin/sh", "-c", "ps -ef | grep -Eo 'vnc :[0-9]+' | grep -Eo '[0-9]+' | awk -v ORS=',' '{print $1}'"};
    SystemCommandOutput out = executeSystemCommand(command);
    List<Integer> activeDisplays = new ArrayList<>();
    Arrays.asList(out.stdout.split(",")).forEach(sint -> {
      if (!sint.isEmpty()) activeDisplays.add(Integer.valueOf(sint));
    });
    if (!out.isSuccess()) {
      String msg = LogUtilities.buildMessage("Error finding active vnc sessions{}", LogUtilities.indentBlock(out.stderr));
      if (ansiColors) LogUtilities.midasLog(null, Level.WARNING, msg, null);
      else LogUtilities.log(Level.WARNING, msg);
      return activeDisplays;
    }
    String msg = LogUtilities.buildMessage("Found {} active VNC sessions: {}", activeDisplays.size(), activeDisplays);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    return activeDisplays;
  }

  private boolean startvnc (int i, String userDN, Table requestHeader) {
    if (manual) return true;
    String msg = LogUtilities.buildMessage("Starting VNC session on display :{}", i);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    TextFile tbl = new TextFile(null, getTableFileName(i)); // can not be all UC so Midas not matter
    addClientInfoToTable(requestHeader, userDN);
    if (debug) {
      LogUtilities.midasLog(null, Level.INFO,
        LogUtilities.buildMessage("[DEBUG] Session Info Table: {} {}", getTableFilePath(getTableFileName(i)),
            LogUtilities.indentBlock(requestHeader.toString())), null);
    }
    requestHeader.toTextFile(tbl);
    String[] command = {"/bin/sh", "-c", String.format("vncserver%s :%s", getVNCOptions(i), i)};
    return executeSystemCommand(command).isSuccess();
  }

  private boolean startwss (int i) {
    if (manual) return true;
    int vncport = 5900+i;
    int wssport = 6080+i;
    String msg = LogUtilities.buildMessage("Starting websocket proxy to {} on {}", vncport, wssport);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    String[] command = {"/bin/sh", "-c", String.format("websockify %s %s localhost:%s", getWSSOptions(), wssport, vncport)};
    return executeSystemCommand(command).isSuccess();
  }

  private boolean stopvnc (int i) {
    if (manual) return true;
    String msg = LogUtilities.buildMessage("Ending VNC session on display :{}", i);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    TextFile.delete(getTableFilePath(getTableFileName(i)));
    String[] command = {"/bin/sh", "-c", String.format("vncserver -kill :%s", i)};
    return executeSystemCommand(command).isSuccess();
  }

  private boolean stopwss (int i) {
    if (manual) return true;
    int wssport = 6080+i;
    String msg = LogUtilities.buildMessage("Ending websocket proxy on {}", wssport);
    if (ansiColors) LogUtilities.midasLog(null, Level.INFO, msg, null);
    else LogUtilities.log(Level.INFO, msg);
    String[] command = {"/bin/sh", "-c", String.format("pkill -f 'websockify %s %s'", getWSSOptions(), wssport)};
    return executeSystemCommand(command).isSuccess();
  }

  private String getVNCOptions (int i) {
    return ((geometry == null) ? "" : " -geometry " + geometry)
        + " -xstartup '" + xstartup + " " + getTableFilePath(getTableFileName(i)) + "'"
        + ((securityTypes == null) ? "" : " -securitytypes " + securityTypes);
  }

  private String getWSSOptions () {
    String options = "-D";
    if (ssl) options += " --ssl-only --cert=" + certPath + " --key=" + keyPath;
    return options;
  }

  private void addClientInfoToTable(Table table, String userDN) {
    if (userDN == null || userDN.isEmpty()) return;
    table.addIfNotPresent("clientDN", userDN);
    String userCN = Arrays.asList(userDN.split(",")).stream()
        .map(s -> s.split("="))
        .filter(s -> s[0].equals("CN"))
        .map(s -> s[1])
        .collect(Collectors.toList()).get(0);
    table.addIfNotPresent("clientCN", userCN);
  }

  private String getTableFileName(int i) {
    return "webvnc_" + i + ".tbl";
  }

  private String getTableFilePath(String filename) {
    try {
      return TextFile.getFileNameFor(null, filename).toLocalFileName();
    } catch (MalformedURLException e) {
      // TODO Auto-generated catch block
      e.printStackTrace();
    }
    return filename;
  }

  /**
   * Executes a bash command
   *
   * @param command The bash command
   * @return The OS response of the bash command
   * @throws AirtightException Any Exception or error response from OS
   */
  private SystemCommandOutput executeSystemCommand(String[] command) {
    if (debug) LogUtilities.midasLog(null, Level.INFO,
        LogUtilities.buildMessage("[DEBUG] Executing shell command: {}", command[2]), null);

    int out = 1;
    StringBuilder stdout = new StringBuilder();
    StringBuilder stderr = new StringBuilder();

    try {
      String line = null;
      Process proc = Runtime.getRuntime().exec(command);

      BufferedReader br = new BufferedReader(new InputStreamReader(proc.getInputStream()));
      while ((line = br.readLine()) != null) {
        if (!stdout.toString().isEmpty()) stdout.append(System.lineSeparator());
        stdout.append(line);
      }
      br.close();

      br = new BufferedReader(new InputStreamReader(proc.getErrorStream()));
      while ((line = br.readLine()) != null) {
        if (!stderr.toString().isEmpty()) stdout.append(System.lineSeparator());
        stderr.append(line);
      }
      out = proc.waitFor();
      proc.destroy();
    } catch (Exception e) {
      String msg = LogUtilities.buildMessage("Error executing system command:\n\t{}\n\t{}", command[2]);
      if (ansiColors) LogUtilities.midasLog(null, Level.SEVERE, msg + "\n " + e, null);
      else LogUtilities.log(Level.SEVERE, msg, e);
      LogUtilities.log(Level.SEVERE, LogUtilities.buildMessage("Error executing system command:\n\t{}\n\t{}", command[2]), e);
    }
    if (debug) {
      if (!stdout.toString().isEmpty()) {
        LogUtilities.midasLog(null, Level.INFO,
            LogUtilities.buildMessage("[DEBUG] STDOUT{}",LogUtilities.indentBlock(stdout.toString())), null);
      }
      if (!stderr.toString().isEmpty()) {
      }
    }
    return new SystemCommandOutput(out, stdout.toString(), stderr.toString());
  }

  private class SystemCommandOutput {
    public int exitCode;
    public String stdout;
    public String stderr;

    public SystemCommandOutput(int exitCode, String stdout, String stderr) {
      this.exitCode = exitCode;
      this.stdout = stdout;
      this.stderr = stderr;
    }

    public boolean isSuccess() {
      return exitCode == 0;
    }
  }

}
