import { ChildProcess, spawn } from "child_process";
import * as path from "path";
import * as WebSocket from "ws";

import * as pty from "node-pty";
import { IPty } from "node-pty";
import PQueue from "p-queue";
import * as rpc from "vscode-jsonrpc";
import { v4 as getUUID } from "uuid";

import { LangConfig, langs } from "./langs";
import { borrowUser } from "./users";
import * as util from "./util";
import { Context, Options, bash } from "./util";

const PACKAGE_MAX_SEARCH_RESULTS = 100;
const PACKAGE_NAME_REGEX = /[-_a-zA-Z0-9.+:]/;

const allSessions: Set<Session> = new Set();

export class Session {
  ws: WebSocket;
  uuid: string;
  lang: string;

  tearingDown: boolean = false;

  // Initialized by setup()
  uidInfo: {
    uid: number;
    returnUID: () => Promise<void>;
  } | null = null;

  // Initialized later or never
  term: { pty: IPty; live: boolean } | null = null;
  lsp: {
    proc: ChildProcess;
    reader: rpc.StreamMessageReader;
    writer: rpc.StreamMessageWriter;
  } | null = null;
  daemon: { proc: ChildProcess } | null = null;
  formatter: {
    proc: ChildProcess;
    live: boolean;
    input: string;
    output: string;
  } | null = null;
  packageSearcher: {
    proc: ChildProcess;
    live: boolean;
    input: string;
    output: string;
  } | null = null;

  logPrimitive: (msg: string) => void;

  msgQueue: PQueue = new PQueue({ concurrency: 1 });

  get homedir() {
    return `/tmp/riju/${this.uuid}`;
  }

  get config() {
    return langs[this.lang];
  }

  get uid() {
    return this.uidInfo!.uid;
  }

  returnUID = async () => {
    this.uidInfo && (await this.uidInfo.returnUID());
  };

  get context() {
    return { uid: this.uid, uuid: this.uuid };
  }

  log = (msg: string) => this.logPrimitive(`[${this.uuid}] ${msg}`);

  constructor(ws: WebSocket, lang: string, log: (msg: string) => void) {
    this.ws = ws;
    this.uuid = getUUID();
    this.lang = lang;
    this.logPrimitive = log;
    this.log(`Creating session, language ${this.lang}`);
  }

  run = async (args: string[], options?: Options) => {
    return await util.run(args, this.log, options);
  };

  privilegedSetup = () => util.privilegedSetup(this.context);
  privilegedSpawn = (args: string[]) =>
    util.privilegedSpawn(this.context, args);
  privilegedUseradd = () => util.privilegedUseradd(this.uid);
  privilegedTeardown = () => util.privilegedTeardown(this.context);

  setup = async () => {
    try {
      allSessions.add(this);
      const { uid, returnUID } = await borrowUser(this.log);
      this.uidInfo = { uid, returnUID };
      this.log(`Borrowed uid ${this.uid}`);
      await this.run(this.privilegedSetup());
      if (this.config.setup) {
        await this.run(this.privilegedSpawn(bash(this.config.setup)));
      }
      await this.runCode();
      if (this.config.daemon) {
        const daemonArgs = this.privilegedSpawn(bash(this.config.daemon));
        const daemonProc = spawn(daemonArgs[0], daemonArgs.slice(1));
        this.daemon = {
          proc: daemonProc,
        };
        for (const stream of [daemonProc.stdout, daemonProc.stderr]) {
          stream.on("data", (data) =>
            this.send({
              event: "serviceLog",
              service: "daemon",
              output: data.toString("utf8"),
            })
          );
          daemonProc.on("close", (code, signal) =>
            this.send({
              event: "serviceFailed",
              service: "daemon",
              error: `Exited with status ${signal || code}`,
            })
          );
          daemonProc.on("error", (err) =>
            this.send({
              event: "serviceFailed",
              service: "daemon",
              error: `${err}`,
            })
          );
        }
      }
      if (this.config.lsp) {
        if (this.config.lsp.setup) {
          await this.run(this.privilegedSpawn(bash(this.config.lsp.setup)));
        }
        const lspArgs = this.privilegedSpawn(bash(this.config.lsp.start));
        const lspProc = spawn(lspArgs[0], lspArgs.slice(1));
        this.lsp = {
          proc: lspProc,
          reader: new rpc.StreamMessageReader(lspProc.stdout),
          writer: new rpc.StreamMessageWriter(lspProc.stdin),
        };
        this.lsp.reader.listen((data: any) => {
          this.send({ event: "lspOutput", output: data });
        });
        lspProc.stderr.on("data", (data) =>
          this.send({
            event: "serviceLog",
            service: "lsp",
            output: data.toString("utf8"),
          })
        );
        lspProc.on("close", (code, signal) =>
          this.send({
            event: "serviceFailed",
            service: "lsp",
            error: `Exited with status ${signal || code}`,
          })
        );
        lspProc.on("error", (err) =>
          this.send({ event: "serviceFailed", service: "lsp", error: `${err}` })
        );
        this.send({ event: "lspStarted", root: this.homedir });
      }
      this.ws.on("message", (msg: string) =>
        this.msgQueue.add(() => this.receive(msg))
      );
      this.ws.on("close", async () => {
        await this.teardown();
      });
      this.ws.on("error", async (err) => {
        this.log(`Websocket error: ${err}`);
        await this.teardown();
      });
    } catch (err) {
      this.log(`Error while setting up environment`);
      console.log(err);
      this.sendError(err);
      await this.teardown();
    }
  };

  send = async (msg: any) => {
    try {
      if (this.tearingDown) {
        return;
      }
      this.ws.send(JSON.stringify(msg));
    } catch (err) {
      this.log(`Failed to send websocket message: ${err}`);
      console.log(err);
      await this.teardown();
    }
  };

  sendError = async (err: any) => {
    await this.send({ event: "terminalClear" });
    await this.send({
      event: "terminalOutput",
      output: `Riju encountered an unexpected error: ${err}
\r
\rYou may want to save your code and refresh the page.
`,
    });
  };

  logBadMessage = (msg: any) => {
    this.log(`Got malformed message from client: ${JSON.stringify(msg)}`);
  };

  receive = async (event: string) => {
    try {
      if (this.tearingDown) {
        return;
      }
      let msg: any;
      try {
        msg = JSON.parse(event);
      } catch (err) {
        this.log(`Failed to parse message from client: ${event}`);
        return;
      }
      switch (msg && msg.event) {
        case "terminalInput":
          if (typeof msg.input !== "string") {
            this.logBadMessage(msg);
            break;
          }
          if (!this.term) {
            this.log("terminalInput ignored because term is null");
            break;
          }
          this.term!.pty.write(msg.input);
          break;
        case "runCode":
          if (typeof msg.code !== "string") {
            this.logBadMessage(msg);
            break;
          }
          await this.runCode(msg.code);
          break;
        case "formatCode":
          if (typeof msg.code !== "string") {
            this.logBadMessage(msg);
            break;
          }
          if (!this.config.format) {
            this.log("formatCode ignored because format is null");
          }
          await this.formatCode(msg.code);
          break;
        case "lspInput":
          if (typeof msg.input !== "object" || !msg) {
            this.logBadMessage(msg);
            break;
          }
          if (!this.lsp) {
            this.log(`lspInput ignored because lsp is null`);
            break;
          }
          this.lsp.writer.write(msg.input);
          break;
        case "ensure":
          if (!this.config.ensure) {
            this.log(`ensure ignored because ensure is null`);
            break;
          }
          await this.ensure(this.config.ensure);
          break;
        case "packageSearch":
          if (!this.config.pkg || !this.config.pkg.search) {
            this.log(`packageSearch ignored because pkg.search is null`);
            break;
          }
          if (typeof msg.search !== "string") {
            this.logBadMessage(msg);
            break;
          }
          await this.packageSearch(msg.search);
          break;
        default:
          this.logBadMessage(msg);
          break;
      }
    } catch (err) {
      this.log(`Error while handling message from client`);
      console.log(err);
      this.sendError(err);
    }
  };

  writeCode = async (code: string) => {
    if (this.config.main.includes("/")) {
      await this.run(
        this.privilegedSpawn([
          "mkdir",
          "-p",
          path.dirname(`${this.homedir}/${this.config.main}`),
        ])
      );
    }
    await this.run(
      this.privilegedSpawn([
        "sh",
        "-c",
        `cat > ${path.resolve(this.homedir, this.config.main)}`,
      ]),
      { input: code }
    );
  };

  runCode = async (code?: string) => {
    try {
      const {
        name,
        repl,
        main,
        suffix,
        createEmpty,
        compile,
        run,
        template,
      } = this.config;
      if (this.term) {
        const pid = this.term.pty.pid;
        const args = this.privilegedSpawn(
          bash(`kill -SIGTERM ${pid}; sleep 1; kill -SIGKILL ${pid}`)
        );
        spawn(args[0], args.slice(1));
        // Signal to terminalOutput message generator using closure.
        this.term.live = false;
        this.term = null;
      }
      this.send({ event: "terminalClear" });
      let cmdline: string;
      if (code) {
        cmdline = run;
        if (compile) {
          cmdline = `( ${compile} ) && ( ${run} )`;
        }
      } else if (repl) {
        cmdline = repl;
      } else {
        cmdline = `echo '${name} has no REPL, press Run to see it in action'`;
      }
      if (code === undefined) {
        code = createEmpty !== undefined ? createEmpty : template;
      }
      if (code && suffix) {
        code += suffix;
      }
      await this.writeCode(code);
      const termArgs = this.privilegedSpawn(bash(cmdline));
      const term = {
        pty: pty.spawn(termArgs[0], termArgs.slice(1), {
          name: "xterm-color",
        }),
        live: true,
      };
      this.term = term;
      this.term.pty.on("data", (data) => {
        // Capture term in closure so that we don't keep sending output
        // from the old pty even after it's been killed (see ghci).
        if (term.live) {
          this.send({ event: "terminalOutput", output: data });
        }
      });
      this.term.pty.on("exit", (code, signal) => {
        if (term.live) {
          this.send({
            event: "serviceFailed",
            service: "terminal",
            error: `Exited with status ${signal || code}`,
          });
        }
      });
    } catch (err) {
      this.log(`Error while running user code`);
      console.log(err);
      this.sendError(err);
    }
  };

  formatCode = async (code: string) => {
    try {
      if (this.formatter) {
        const pid = this.formatter.proc.pid;
        const args = this.privilegedSpawn(
          bash(`kill -SIGTERM ${pid}; sleep 1; kill -SIGKILL ${pid}`)
        );
        spawn(args[0], args.slice(1));
        this.formatter.live = false;
        this.formatter = null;
      }
      const args = this.privilegedSpawn(bash(this.config.format!.run));
      const formatter = {
        proc: spawn(args[0], args.slice(1)),
        live: true,
        input: code,
        output: "",
      };
      formatter.proc.stdin!.end(code);
      formatter.proc.stdout!.on("data", (data) => {
        if (!formatter.live) return;
        formatter.output += data.toString("utf8");
      });
      formatter.proc.stderr!.on("data", (data) => {
        if (!formatter.live) return;
        this.send({
          event: "serviceLog",
          service: "formatter",
          output: data.toString("utf8"),
        });
      });
      formatter.proc.on("close", (code, signal) => {
        if (!formatter.live) return;
        if (code === 0) {
          this.send({
            event: "formattedCode",
            code: formatter.output,
            originalCode: formatter.input,
          });
        } else {
          this.send({
            event: "serviceFailed",
            service: "formatter",
            error: `Exited with status ${signal || code}`,
          });
        }
      });
      formatter.proc.on("error", (err) => {
        if (!formatter.live) return;
        this.send({
          event: "serviceFailed",
          service: "formatter",
          error: `${err}`,
        });
      });
      this.formatter = formatter;
    } catch (err) {
      this.log(`Error while running code formatter`);
      console.log(err);
      this.sendError(err);
    }
  };

  ensure = async (cmd: string) => {
    const code = await this.run(this.privilegedSpawn(bash(cmd)), {
      check: false,
    });
    this.send({ event: "ensured", code });
  };

  packageSearch = async (search: string) => {
    try {
      if (this.packageSearcher) {
        const pid = this.packageSearcher.proc.pid;
        const args = this.privilegedSpawn(
          bash(`kill -SIGTERM ${pid}; sleep 1; kill -SIGKILL ${pid}`)
        );
        spawn(args[0], args.slice(1));
        this.packageSearcher.live = false;
        this.packageSearcher = null;
      }
      if (!search) {
        this.send({
          event: "packageSearched",
          results: [],
          search: "",
        });
        return;
      }
      const args = this.privilegedSpawn(
        bash(this.config.pkg!.search!.replace(/NAME/g, search))
      );
      const packageSearcher = {
        proc: spawn(args[0], args.slice(1)),
        live: true,
        input: search,
        output: "",
      };
      packageSearcher.proc.stdout!.on("data", (data) => {
        if (!packageSearcher.live) return;
        packageSearcher.output += data.toString("utf8");
      });
      packageSearcher.proc.stderr!.on("data", (data) => {
        if (!packageSearcher.live) return;
        this.send({
          event: "serviceLog",
          service: "packageSearch",
          output: data.toString("utf8"),
        });
      });
      packageSearcher.proc.on("close", (code, signal) => {
        if (!packageSearcher.live) return;
        if (code === 0) {
          this.send({
            event: "packageSearched",
            results: packageSearcher.output.split("\n").filter((x) => x),
            search: packageSearcher.input,
          });
        } else {
          this.send({
            event: "serviceFailed",
            service: "packageSearch",
            error: `Exited with status ${signal || code}`,
          });
        }
      });
      packageSearcher.proc.on("error", (err) => {
        if (!packageSearcher.live) return;
        this.send({
          event: "serviceFailed",
          service: "packageSearch",
          error: `${err}`,
        });
      });
      this.packageSearcher = packageSearcher;
    } catch (err) {
      this.log(`Error while running package search`);
      console.log(err);
      this.sendError(err);
    }
  };

  teardown = async () => {
    try {
      if (this.tearingDown) {
        return;
      }
      this.log(`Tearing down session`);
      this.tearingDown = true;
      allSessions.delete(this);
      if (this.uidInfo) {
        await this.run(this.privilegedTeardown());
        await this.returnUID();
      }
      this.ws.terminate();
    } catch (err) {
      this.log(`Error during teardown`);
      console.log(err);
    }
  };
}