import { EventEmitter } from 'events';
import {
  ChildProcess,
  spawn,
  SpawnOptions,
  execFile,
  execFileSync,
} from 'child_process';
import { EOL as newline, tmpdir } from 'os';
import { join, sep } from 'path';
import { Readable, Transform, TransformCallback, Writable } from 'stream';
import { writeFile, writeFileSync } from 'fs';
import { promisify } from 'util';

function toArray<T>(source?: T | T[]): T[] {
  if (typeof source === 'undefined' || source === null) {
    return [];
  } else if (!Array.isArray(source)) {
    return [source];
  }
  return source;
}

/**
 * adds arguments as properties to obj
 */
function extend(obj: {}, ...args) {
  Array.prototype.slice.call(arguments, 1).forEach(function (source) {
    if (source) {
      for (let key in source) {
        obj[key] = source[key];
      }
    }
  });
  return obj;
}

/**
 * gets a random int from 0-10000000000
 */
function getRandomInt() {
  return Math.floor(Math.random() * 10000000000);
}

const execFilePromise = promisify(execFile);

export interface Options extends SpawnOptions {
  /**
   * if binary is enabled message and stderr events will not be emitted
   */
  mode?: 'text' | 'json' | 'binary';
  formatter?: string | ((param: string) => any);
  parser?: string | ((param: string) => any);
  stderrParser?: string | ((param: string) => any);
  encoding?: BufferEncoding;
  pythonPath?: string;
  /**
   * see https://docs.python.org/3.7/using/cmdline.html
   */
  pythonOptions?: string[];
  /**
   * overrides scriptPath passed into PythonShell constructor
   */
  scriptPath?: string;
  /**
   * arguments to your program
   */
  args?: string[];
}

export class PythonShellError extends Error {
  traceback: string | Buffer;
  exitCode?: number;
}

export class PythonShellErrorWithLogs extends PythonShellError {
  logs: any[];
}

/**
 * Takes in a string stream and emits batches seperated by newlines
 */
export class NewlineTransformer extends Transform {
  // NewlineTransformer: Megatron's little known once-removed cousin
  private _lastLineData: string;
  _transform(chunk: any, encoding: string, callback: TransformCallback) {
    let data: string = chunk.toString();
    if (this._lastLineData) data = this._lastLineData + data;
    const lines = data.split(newline);
    this._lastLineData = lines.pop();
    lines.forEach(this.push.bind(this));
    callback();
  }
  _flush(done: TransformCallback) {
    if (this._lastLineData) this.push(this._lastLineData);
    this._lastLineData = null;
    done();
  }
}

/**
 * An interactive Python shell exchanging data through stdio
 * @param {string} script    The python script to execute
 * @param {object} [options] The launch options (also passed to child_process.spawn)
 * @param [stdoutSplitter] Optional. Splits stdout into chunks, defaulting to splitting into newline-seperated lines
 * @param [stderrSplitter] Optional. splits stderr into chunks, defaulting to splitting into newline-seperated lines
 * @constructor
 */
export class PythonShell extends EventEmitter {
  scriptPath: string;
  command: string[];
  mode: string;
  formatter: (param: string | Object) => any;
  parser: (param: string) => any;
  stderrParser: (param: string) => any;
  terminated: boolean;
  childProcess: ChildProcess;
  stdin: Writable;
  stdout: Readable;
  stderr: Readable;
  exitSignal: string;
  exitCode: number;
  private stderrHasEnded: boolean;
  private stdoutHasEnded: boolean;
  private _remaining: string;
  private _endCallback: (
    err: PythonShellError,
    exitCode: number,
    exitSignal: string,
  ) => any;

  // starting 2020 python2 is deprecated so we choose 3 as default
  static defaultPythonPath = process.platform != 'win32' ? 'python3' : 'python';

  static defaultOptions: Options = {}; //allow global overrides for options

  /**
   * spawns a python process
   * @param scriptPath path to script. Relative to current directory or options.scriptFolder if specified
   * @param options
   * @param stdoutSplitter Optional. Splits stdout into chunks, defaulting to splitting into newline-seperated lines
   * @param stderrSplitter Optional. splits stderr into chunks, defaulting to splitting into newline-seperated lines
   */
  constructor(
    scriptPath: string,
    options?: Options,
    stdoutSplitter: Transform = null,
    stderrSplitter: Transform = null,
  ) {
    super();

    /**
     * returns either pythonshell func (if val string) or custom func (if val Function)
     */
    function resolve(type, val: string | Function) {
      if (typeof val === 'string') {
        // use a built-in function using its name
        return PythonShell[type][val];
      } else if (typeof val === 'function') {
        // use a custom function
        return val;
      }
    }

    if (scriptPath.trim().length == 0)
      throw Error(
        'scriptPath cannot be empty! You must give a script for python to run',
      );

    let self = this;
    let errorData = '';
    EventEmitter.call(this);

    options = <Options>extend({}, PythonShell.defaultOptions, options);
    let pythonPath: string;
    if (!options.pythonPath) {
      pythonPath = PythonShell.defaultPythonPath;
    } else pythonPath = options.pythonPath;
    let pythonOptions = toArray(options.pythonOptions);
    let scriptArgs = toArray(options.args);

    this.scriptPath = join(options.scriptPath || '', scriptPath);
    this.command = pythonOptions.concat(this.scriptPath, scriptArgs);
    this.mode = options.mode || 'text';
    this.formatter = resolve('format', options.formatter || this.mode);
    this.parser = resolve('parse', options.parser || this.mode);
    // We don't expect users to ever format stderr as JSON so we default to text mode
    this.stderrParser = resolve('parse', options.stderrParser || 'text');
    this.terminated = false;
    this.childProcess = spawn(pythonPath, this.command, options);

    ['stdout', 'stdin', 'stderr'].forEach(function (name) {
      self[name] = self.childProcess[name];
      self.parser &&
        self[name] &&
        self[name].setEncoding(options.encoding || 'utf8');
    });

    // Node buffers stdout&stderr in batches regardless of newline placement
    // This is troublesome if you want to recieve distinct individual messages
    // for example JSON parsing breaks if it recieves partial JSON
    // so we use newlineTransformer to emit each batch seperated by newline
    if (this.parser && this.stdout) {
      if (!stdoutSplitter) stdoutSplitter = new NewlineTransformer();
      // note that setting the encoding turns the chunk into a string
      stdoutSplitter.setEncoding(options.encoding || 'utf8');
      this.stdout.pipe(stdoutSplitter).on('data', (chunk: string) => {
        this.emit('message', self.parser(chunk));
      });
    }

    // listen to stderr and emit errors for incoming data
    if (this.stderrParser && this.stderr) {
      if (!stderrSplitter) stderrSplitter = new NewlineTransformer();
      // note that setting the encoding turns the chunk into a string
      stderrSplitter.setEncoding(options.encoding || 'utf8');
      this.stderr.pipe(stderrSplitter).on('data', (chunk: string) => {
        this.emit('stderr', self.stderrParser(chunk));
      });
    }

    if (this.stderr) {
      this.stderr.on('data', function (data) {
        errorData += '' + data;
      });
      this.stderr.on('end', function () {
        self.stderrHasEnded = true;
        terminateIfNeeded();
      });
    } else {
      self.stderrHasEnded = true;
    }

    if (this.stdout) {
      this.stdout.on('end', function () {
        self.stdoutHasEnded = true;
        terminateIfNeeded();
      });
    } else {
      self.stdoutHasEnded = true;
    }

    this.childProcess.on('error', function (err: NodeJS.ErrnoException) {
      self.emit('error', err);
    });
    this.childProcess.on('exit', function (code, signal) {
      self.exitCode = code;
      self.exitSignal = signal;
      terminateIfNeeded();
    });

    function terminateIfNeeded() {
      if (
        !self.stderrHasEnded ||
        !self.stdoutHasEnded ||
        (self.exitCode == null && self.exitSignal == null)
      )
        return;

      let err: PythonShellError;
      if (self.exitCode && self.exitCode !== 0) {
        if (errorData) {
          err = self.parseError(errorData);
        } else {
          err = new PythonShellError(
            'process exited with code ' + self.exitCode,
          );
        }
        err = <PythonShellError>extend(err, {
          executable: pythonPath,
          options: pythonOptions.length ? pythonOptions : null,
          script: self.scriptPath,
          args: scriptArgs.length ? scriptArgs : null,
          exitCode: self.exitCode,
        });
        // do not emit error if only a callback is used
        if (self.listeners('pythonError').length || !self._endCallback) {
          self.emit('pythonError', err);
        }
      }

      self.terminated = true;
      self.emit('close');
      self._endCallback &&
        self._endCallback(err, self.exitCode, self.exitSignal);
    }
  }

  // built-in formatters
  static format = {
    text: function toText(data): string {
      if (!data) return '';
      else if (typeof data !== 'string') return data.toString();
      return data;
    },
    json: function toJson(data) {
      return JSON.stringify(data);
    },
  };

  //built-in parsers
  static parse = {
    text: function asText(data): string {
      return data;
    },
    json: function asJson(data: string) {
      return JSON.parse(data);
    },
  };

  /**
   * checks syntax without executing code
   * @returns rejects promise w/ string error output if syntax failure
   */
  static async checkSyntax(code: string) {
    const randomInt = getRandomInt();
    const filePath = tmpdir() + sep + `pythonShellSyntaxCheck${randomInt}.py`;

    const writeFilePromise = promisify(writeFile);
    return writeFilePromise(filePath, code).then(() => {
      return this.checkSyntaxFile(filePath);
    });
  }

  static getPythonPath() {
    return this.defaultOptions.pythonPath
      ? this.defaultOptions.pythonPath
      : this.defaultPythonPath;
  }

  /**
   * checks syntax without executing code
   * @returns {Promise} rejects w/ stderr if syntax failure
   */
  static async checkSyntaxFile(filePath: string) {
    const pythonPath = this.getPythonPath();
    return execFilePromise(pythonPath, ['-m', 'py_compile', filePath]);
  }

  /**
   * Runs a Python script and returns collected messages as a promise.
   * If the promise is rejected, the err will probably be of type PythonShellErrorWithLogs
   * @param scriptPath   The path to the script to execute
   * @param options  The execution options
   */
  static run(scriptPath: string, options?: Options): Promise<any[]> {
    return new Promise((resolve, reject) => {
      let pyshell = new PythonShell(scriptPath, options);
      let output = [];

      pyshell
        .on('message', function (message) {
          output.push(message);
        })
        .end(function (err) {
          if (err) {
            (err as PythonShellErrorWithLogs).logs = output;
            reject(err);
          } else resolve(output);
        });
    });
  }

  /**
   * Runs the inputted string of python code and returns collected messages as a promise. DO NOT ALLOW UNTRUSTED USER INPUT HERE!
   * @param code   The python code to execute
   * @param options  The execution options
   * @return a promise with the output from the python script
   */
  static runString(code: string, options?: Options) {
    // put code in temp file
    const randomInt = getRandomInt();
    const filePath = tmpdir + sep + `pythonShellFile${randomInt}.py`;
    writeFileSync(filePath, code);

    return PythonShell.run(filePath, options);
  }

  static getVersion(pythonPath?: string) {
    if (!pythonPath) pythonPath = this.getPythonPath();
    return execFilePromise(pythonPath, ['--version']);
  }

  static getVersionSync(pythonPath?: string) {
    if (!pythonPath) pythonPath = this.getPythonPath();
    return execFileSync(pythonPath, ['--version']).toString();
  }

  /**
   * Parses an error thrown from the Python process through stderr
   * @param  {string|Buffer} data The stderr contents to parse
   * @return {Error} The parsed error with extended stack trace when traceback is available
   */
  private parseError(data: string | Buffer) {
    let text = '' + data;
    let error: PythonShellError;

    if (/^Traceback/.test(text)) {
      // traceback data is available
      let lines = text.trim().split(newline);
      let exception = lines.pop();
      error = new PythonShellError(exception);
      error.traceback = data;
      // extend stack trace
      error.stack +=
        newline + '    ----- Python Traceback -----' + newline + '  ';
      error.stack += lines.slice(1).join(newline + '  ');
    } else {
      // otherwise, create a simpler error with stderr contents
      error = new PythonShellError(text);
    }

    return error;
  }

  /**
   * Sends a message to the Python shell through stdin
   * Override this method to format data to be sent to the Python process
   * @returns {PythonShell} The same instance for chaining calls
   */
  send(message: string | Object) {
    if (!this.stdin) throw new Error('stdin not open for writing');
    let data = this.formatter ? this.formatter(message) : message;
    if (this.mode !== 'binary') data += newline;
    this.stdin.write(data);
    return this;
  }

  /**
   * Closes the stdin stream. Unless python is listening for stdin in a loop
   * this should cause the process to finish its work and close.
   * @returns {PythonShell} The same instance for chaining calls
   */
  end(
    callback: (
      err: PythonShellError,
      exitCode: number,
      exitSignal: string,
    ) => any,
  ) {
    if (this.childProcess.stdin) {
      this.childProcess.stdin.end();
    }
    this._endCallback = callback;
    return this;
  }

  /**
   * Sends a kill signal to the process
   * @returns {PythonShell} The same instance for chaining calls
   */
  kill(signal?: NodeJS.Signals) {
    this.terminated = this.childProcess.kill(signal);
    return this;
  }

  /**
   * Alias for kill.
   * @deprecated
   */
  terminate(signal?: NodeJS.Signals) {
    // todo: remove this next breaking release
    return this.kill(signal);
  }
}

// This interface is merged in with the above class definition
export interface PythonShell {
  addListener(event: string, listener: (...args: any[]) => void): this;
  emit(event: string | symbol, ...args: any[]): boolean;
  on(event: string, listener: (...args: any[]) => void): this;
  once(event: string, listener: (...args: any[]) => void): this;
  prependListener(event: string, listener: (...args: any[]) => void): this;
  prependOnceListener(event: string, listener: (...args: any[]) => void): this;

  addListener(event: 'message', listener: (parsedChunk: any) => void): this;
  emit(event: 'message', parsedChunk: any): boolean;
  on(event: 'message', listener: (parsedChunk: any) => void): this;
  once(event: 'message', listener: (parsedChunk: any) => void): this;
  prependListener(event: 'message', listener: (parsedChunk: any) => void): this;
  prependOnceListener(
    event: 'message',
    listener: (parsedChunk: any) => void,
  ): this;

  addListener(event: 'stderr', listener: (parsedChunk: any) => void): this;
  emit(event: 'stderr', parsedChunk: any): boolean;
  on(event: 'stderr', listener: (parsedChunk: any) => void): this;
  once(event: 'stderr', listener: (parsedChunk: any) => void): this;
  prependListener(event: 'stderr', listener: (parsedChunk: any) => void): this;
  prependOnceListener(
    event: 'stderr',
    listener: (parsedChunk: any) => void,
  ): this;

  addListener(event: 'close', listener: () => void): this;
  emit(event: 'close'): boolean;
  on(event: 'close', listener: () => void): this;
  once(event: 'close', listener: () => void): this;
  prependListener(event: 'close', listener: () => void): this;
  prependOnceListener(event: 'close', listener: () => void): this;

  addListener(
    event: 'error',
    listener: (error: NodeJS.ErrnoException) => void,
  ): this;
  emit(event: 'error', error: NodeJS.ErrnoException): boolean;
  on(event: 'error', listener: (error: NodeJS.ErrnoException) => void): this;
  once(event: 'error', listener: (error: NodeJS.ErrnoException) => void): this;
  prependListener(
    event: 'error',
    listener: (error: NodeJS.ErrnoException) => void,
  ): this;
  prependOnceListener(
    event: 'error',
    listener: (error: NodeJS.ErrnoException) => void,
  ): this;

  addListener(
    event: 'pythonError',
    listener: (error: PythonShellError) => void,
  ): this;
  emit(event: 'pythonError', error: PythonShellError): boolean;
  on(event: 'pythonError', listener: (error: PythonShellError) => void): this;
  once(event: 'pythonError', listener: (error: PythonShellError) => void): this;
  prependListener(
    event: 'pythonError',
    listener: (error: PythonShellError) => void,
  ): this;
  prependOnceListener(
    event: 'pythonError',
    listener: (error: PythonShellError) => void,
  ): this;
}