import { TLogType } from "./../query/logger/logger.types";
import * as ivm from "isolated-vm";
import { VModule } from "./module.class";
import { Plugin } from "./plugin.class";
import { QueryResponse } from "./vm.constants";
import { TLog } from "src/query/logger/logger.types";
import { LoggerService } from "src/query/logger/logger.service";

export class Vm {
  private memoryLimit: number;
  private modules: VModule[];
  private context: any;
  private jail: any;
  private plugins: Plugin[];
  private functions: string[];
  private isolate: ivm.Isolate;
  private timeLimit?: bigint;
  private cpuTimeLimit?: bigint;
  private log: TLog;
  private headers?: Record<string, string>;
  private cookies?: Record<string, string>;
  private callStack = 0;

  constructor(
    configs: {
      memoryLimit: number;
      timeLimit?: bigint;
      cpuTimeLimit?: bigint;
      modules: VModule[];
      plugins: Plugin[];
      functions: string[];
      log: TLog;
      headers?: Record<string, string>;
      cookies?: Record<string, string>;
    },
    callStack = 0
  ) {
    this.callStack = callStack;

    if (this.callStack > 5) {
      throw new Error("Maximum call stack size exceeded");
    }

    this.memoryLimit = configs.memoryLimit;
    this.modules = configs.modules;
    this.plugins = configs.plugins;
    this.timeLimit = configs.timeLimit;
    this.cpuTimeLimit = configs.cpuTimeLimit;
    this.functions = configs.functions;
    this.log = configs.log;
    this.headers = configs.headers;
    this.cookies = configs.cookies;
  }

  async init(): Promise<Vm> {
    this.isolate = new ivm.Isolate({ memoryLimit: this.memoryLimit });
    this.context = await this.isolate.createContext();
    this.jail = this.context.global;

    this.jail.set("global", this.jail.derefInto());

    for (const mod of this.modules) {
      await this.context.eval(mod.getSource());
    }

    for (const fn of this.functions) {
      await this.context.eval(fn);
    }

    for (const plugin of this.plugins) {
      const pluginName = plugin.getName();

      await this.context.evalClosure(
        "globalThis[$0] = globalThis[$0] || Object.create(null);",
        [pluginName],
        { arguments: { copy: true } }
      );

      for (const method of plugin.getMethods()) {
        const fnRef = new ivm.Reference(async (...args) => {
          plugin.setLog(this.log);
          plugin.setHeaders(this.headers);
          plugin.setCookies(this.cookies);
          plugin.setCallStack(this.callStack);
          const result = await plugin[method](...args);

          if (result && result.log) {
            this.log = result.log;
            delete result.log;
          }

          return result;
        });

        await this.context.evalClosure(
          `
            globalThis[$0][$1] = (...args) => $2.apply(undefined, args, {
              arguments: { copy: true },
              result:    { promise: true, copy: true }
            });
        `,
          [pluginName, method, fnRef],
          { arguments: { copy: true } }
        );
      }
    }

    return this;
  }

  setFunction(name: string, func: (...args) => any) {
    this.jail.setSync(name, func, { arguments: { copy: true } });
  }

  async runScript(
    script: string,
    args: Record<string, any>,
    headers: Record<string, any>
  ): Promise<QueryResponse> {
    let resolvePromise: (value: any) => void;
    let rejectPromise: (reason?: any) => void;

    const resultPromise = new Promise<any>((resolve, reject) => {
      resolvePromise = resolve;
      rejectPromise = reject;
    });

    this.setFunction("returnResult", (res) => {
      resolvePromise({ ...res, log: this.log });
    });

    this.setFunction("log", (...args) => {
      const logType = args.find(
        (arg) =>
          arg &&
          typeof arg === "object" &&
          arg.type &&
          Object.values(TLogType).includes(arg.type)
      );

      this.log = LoggerService.log(this.log, {
        content: args
          .map((arg) => (typeof arg === "string" ? arg : JSON.stringify(arg)))
          .join(" "),
        type: logType?.type || TLogType.info,
        timeStamp: new Date().getTime(),
      });
    });

    this.setFunction("error", (error: any) => {
      LoggerService.log(this.log, {
        content: error?.stack || error?.toString() || "Unknown error",
        type: TLogType.error,
        timeStamp: new Date().getTime(),
      });

      rejectPromise(error);
    });

    const scriptWithResult = `
        (async () => {
            ${script}
            try {
                const result = await main(${JSON.stringify(
                  args
                )}, ${JSON.stringify(headers)});
                returnResult(result)
            } catch (e) {
                error(e)
            }
        })()
    `;

    const compiledScript = await this.isolate.compileScript(scriptWithResult);
    let timer = 0n;

    const interval = setInterval(() => {
      if (
        this.isolate.cpuTime > this.cpuTimeLimit ||
        this.isolate.wallTime > this.timeLimit ||
        timer > this.timeLimit
      ) {
        this.isolate.dispose();

        rejectPromise(new Error("Script execution timed out"));
      }

      timer += 500000n;
    }, 500);

    compiledScript.run(this.context);

    try {
      return await resultPromise;
    } finally {
      clearInterval(interval);
      this.onFinish();
    }
  }

  onFinish() {
    this.plugins.forEach((plugin) => {
      plugin.onFinish();
    });
  }
}