// Copyright 2021 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only

import { join } from 'node:path';
import { Worker } from 'node:worker_threads';
import { format } from 'node:util';
import { app } from 'electron';

import { strictAssert } from '../util/assert.std.ts';
import { explodePromise } from '../util/explodePromise.std.ts';
import { getAppRootDir } from '../util/appRootDir.main.ts';
import type { LoggerType } from '../types/Logging.std.ts';
import * as Errors from '../types/errors.std.ts';
import { SqliteErrorKind } from './errors.std.ts';
import type {
  ServerReadableDirectInterface,
  ServerWritableDirectInterface,
} from './Interface.std.ts';

const MIN_TRACE_DURATION = 40;

const WORKER_COUNT = 4;

const PAGING_QUERIES = new Set<keyof ServerReadableDirectInterface>([
  'pageMessages',
  'finishPageMessages',
  'getKnownMessageAttachments',
  'finishGetKnownMessageAttachments',
]);

export type InitializeOptions = Readonly<{
  appVersion: string;
  configDir: string;
  key: string;
  logger: LoggerType;
}>;

export type WorkerRequest = Readonly<
  | {
      type: 'init';
      options: Omit<InitializeOptions, 'logger'>;
      isPrimary: boolean;
    }
  | {
      type: 'close' | 'removeDB';
    }
  | {
      type: 'walCheckpoint';
      reason: string;
    }
  | {
      type: 'sqlCall:read';
      encoding: 'js';
      method: keyof ServerReadableDirectInterface;
      args: ReadonlyArray<unknown>;
    }
  | {
      type: 'sqlCall:read';
      encoding: 'serialized';
      method: keyof ServerReadableDirectInterface;
      data: Uint8Array<ArrayBuffer>;
    }
  | {
      type: 'sqlCall:write';
      encoding: 'js';
      method: keyof ServerWritableDirectInterface;
      args: ReadonlyArray<unknown>;
    }
  | {
      type: 'sqlCall:write';
      encoding: 'serialized';
      method: keyof ServerWritableDirectInterface;
      data: Uint8Array<ArrayBuffer>;
    }
>;

export type WrappedWorkerRequest = Readonly<{
  seq: number;
  request: WorkerRequest;
}>;

export type WrappedWorkerLogEntry = Readonly<{
  type: 'log';
  level: 'fatal' | 'error' | 'warn' | 'info' | 'debug' | 'trace';
  args: ReadonlyArray<unknown>;
}>;

export type WrappedWorkerResponse =
  | Readonly<{
      type: 'response';
      seq: number;
      error:
        | Readonly<{
            name: string;
            message: string;
            stack: string | undefined;
          }>
        | undefined;
      errorKind: SqliteErrorKind | undefined;
      // oxlint-disable-next-line typescript/no-explicit-any
      response: any;
    }>
  | Readonly<{
      type: 'walCheckpointNeeded';
      reason: string;
    }>
  | WrappedWorkerLogEntry;

type ResponseEntry<T> = {
  errorId: string;
  resolve: (response: T) => void;
  reject: (error: Error) => void;
};

type KnownErrorResolverType = Readonly<{
  kind: SqliteErrorKind;
  resolve: (err: Error) => void;
  once?: boolean;
}>;

type CreateWorkerResultType = Readonly<{
  worker: Worker;
  onExit: Promise<void>;
}>;

type PoolEntry = {
  readonly worker: Worker;
  load: number;
};

type QueryStatsType = {
  queryName: string;
  count: number;
  cumulative: number;
  max: number;
};

export type QueryStatsOptions = {
  maxQueriesToLog?: number;
  epochName?: string;
};

export class MainSQL {
  readonly #pool = new Array<PoolEntry>();
  #pauseWaiters: Array<() => void> | undefined;
  #isReady = false;
  #onReady: Promise<void> | undefined;
  readonly #onExit: Promise<unknown>;

  // Promise resolve callbacks for corruption and readonly errors.
  #errorResolvers = new Array<KnownErrorResolverType>();

  #seq = 0;
  #logger?: LoggerType;

  // oxlint-disable-next-line typescript/no-explicit-any
  readonly #onResponse = new Map<number, ResponseEntry<any>>();

  #checkpointPendingReason: string | null = null;

  readonly #shouldLogQueryTime: (queryName: string) => boolean;
  #shouldTrackQueryStats = false;

  #queryStats?: {
    start: number;
    statsByQuery: Map<string, QueryStatsType>;
  };

  constructor() {
    const exitPromises = new Array<Promise<void>>();
    for (let i = 0; i < WORKER_COUNT; i += 1) {
      const { worker, onExit } = this.#createWorker();
      this.#pool.push({ worker, load: 0 });

      exitPromises.push(onExit);
    }
    this.#onExit = Promise.all(exitPromises);

    const timeQueryEnvVar = process.env.TIME_QUERIES;
    this.#shouldLogQueryTime = (queryName: string) => {
      return timeQueryEnvVar === queryName || timeQueryEnvVar === '1';
    };
  }

  public async initialize({
    appVersion,
    configDir,
    key,
    logger,
  }: InitializeOptions): Promise<void> {
    if (this.#isReady || this.#onReady) {
      throw new Error('Already initialized');
    }

    this.#logger = logger;

    this.#onReady = (async () => {
      const primary = this.#pool[0];
      strictAssert(primary, 'Missing primary');
      const rest = this.#pool.slice(1);

      await this.#send(primary, {
        type: 'init',
        options: { appVersion, configDir, key },
        isPrimary: true,
      });

      await Promise.all(
        rest.map(worker =>
          this.#send(worker, {
            type: 'init',
            options: { appVersion, configDir, key },
            isPrimary: false,
          })
        )
      );
    })();

    await this.#onReady;

    this.#onReady = undefined;
    this.#isReady = true;
  }

  public pauseWriteAccess(): void {
    strictAssert(this.#pauseWaiters == null, 'Already paused');

    this.#pauseWaiters = [];
  }

  public resumeWriteAccess(): void {
    const pauseWaiters = this.#pauseWaiters;
    if (pauseWaiters == null) {
      return;
    }

    this.#pauseWaiters = undefined;

    for (const waiter of pauseWaiters) {
      waiter();
    }
  }

  public whenCorrupted(): Promise<Error> {
    const { promise, resolve } = explodePromise<Error>();
    this.#errorResolvers.push({
      kind: SqliteErrorKind.Corrupted,
      resolve,
      once: true,
    });
    return promise;
  }

  public whenReadonly(): Promise<Error> {
    const { promise, resolve } = explodePromise<Error>();
    this.#errorResolvers.push({
      kind: SqliteErrorKind.Readonly,
      resolve,
      once: true,
    });
    return promise;
  }

  public onUnknownSqlError(callback: (error: Error) => void): void {
    this.#errorResolvers.push({
      kind: SqliteErrorKind.Unknown,
      resolve: callback,
    });
  }

  public async close(): Promise<void> {
    if (this.#onReady) {
      try {
        await this.#onReady;
      } catch (err) {
        this.#logger?.error(
          `MainSQL close, failed: ${Errors.toLogFormat(err)}`
        );
        // Init failed
        return;
      }
    }

    if (!this.#isReady) {
      throw new Error('Not initialized');
    }

    await this.#terminate({ type: 'close' });
    await this.#onExit;
  }

  public async removeDB(): Promise<void> {
    await this.#terminate({ type: 'removeDB' });
  }

  public async sqlRead<Method extends keyof ServerReadableDirectInterface>(
    method: Method,
    ...args: Parameters<ServerReadableDirectInterface[Method]>
  ): Promise<ReturnType<ServerReadableDirectInterface[Method]>> {
    return this.#sqlRead({
      type: 'sqlCall:read',
      method,
      encoding: 'js',
      args,
    }) as Promise<ReturnType<ServerReadableDirectInterface[Method]>>;
  }

  public async sqlReadSerialized<
    Method extends keyof ServerReadableDirectInterface,
  >(
    method: Method,
    data: Uint8Array<ArrayBuffer>
  ): Promise<Uint8Array<ArrayBuffer>> {
    return this.#sqlRead({
      type: 'sqlCall:read',
      method,
      encoding: 'serialized',
      data,
    }) as Promise<Uint8Array<ArrayBuffer>>;
  }

  async #sqlRead(
    request: Extract<WorkerRequest, { type: 'sqlCall:read' }>
  ): Promise<unknown> {
    type SqlCallResult = Readonly<{
      result: unknown;
      duration: number;
    }>;

    if (request.method === 'pageBackupMessages' && this.#pauseWaiters == null) {
      throw new Error(
        'pageBackupMessages can only run after pauseWriteAccess()'
      );
    }

    // pageMessages runs over several queries and needs to have access to
    // the same temporary table, it also creates temporary insert/update
    // triggers so it has to run on the same connection that updates the tables
    const isPaging = PAGING_QUERIES.has(request.method);

    const entry = isPaging ? this.#pool[0] : this.#getWorker();
    strictAssert(entry != null, 'Must have a pool entry');

    const { result, duration } = await this.#send<SqlCallResult>(
      entry,
      request
    );

    this.#traceDuration(request.method, duration);

    return result;
  }

  public async sqlWrite<Method extends keyof ServerWritableDirectInterface>(
    method: Method,
    ...args: Parameters<ServerWritableDirectInterface[Method]>
  ): Promise<ReturnType<ServerWritableDirectInterface[Method]>> {
    return this.#sqlWrite({
      type: 'sqlCall:write',
      method,
      encoding: 'js',
      args,
    }) as Promise<ReturnType<ServerWritableDirectInterface[Method]>>;
  }

  public async sqlWriteSerialized<
    Method extends keyof ServerWritableDirectInterface,
  >(
    method: Method,
    data: Uint8Array<ArrayBuffer>
  ): Promise<Uint8Array<ArrayBuffer>> {
    return this.#sqlWrite({
      type: 'sqlCall:write',
      method,
      encoding: 'serialized',
      data,
    }) as Promise<Uint8Array<ArrayBuffer>>;
  }

  async #sqlWrite(
    request: Extract<WorkerRequest, { type: 'sqlCall:write' }>
  ): Promise<unknown> {
    type SqlCallResult = Readonly<{
      result: unknown;
      duration: number;
    }>;

    while (this.#pauseWaiters != null) {
      const { promise, resolve } = explodePromise<void>();
      this.#pauseWaiters.push(resolve);
      // oxlint-disable-next-line no-await-in-loop
      await promise;
    }

    const primary = this.#pool[0];
    strictAssert(primary, 'Missing primary');

    const { result, duration } = await this.#send<SqlCallResult>(
      primary,
      request
    );

    this.#traceDuration(request.method, duration);

    return result;
  }

  public startTrackingQueryStats(): void {
    if (this.#shouldTrackQueryStats) {
      this.#logQueryStats({});
      this.#logger?.info('Resetting query stats');
    }
    this.#resetQueryStats();
    this.#shouldTrackQueryStats = true;
  }

  public stopTrackingQueryStats(options: QueryStatsOptions): void {
    if (this.#shouldTrackQueryStats) {
      this.#logQueryStats(options);
    }
    this.#queryStats = undefined;
    this.#shouldTrackQueryStats = false;
  }

  async #send<Response>(
    entry: PoolEntry,
    request: WorkerRequest
  ): Promise<Response> {
    let errorId: string = request.type;
    if (request.type === 'sqlCall:read' || request.type === 'sqlCall:write') {
      errorId = `${request.type}(${request.method})`;
      if (this.#onReady) {
        await this.#onReady;
      }

      if (!this.#isReady) {
        throw new Error('Not initialized');
      }
    }

    const seq = this.#seq;
    // oxlint-disable-next-line no-bitwise
    this.#seq = (this.#seq + 1) >>> 0;

    const { promise: result, resolve, reject } = explodePromise<Response>();
    this.#onResponse.set(seq, {
      errorId,
      resolve,
      reject,
    });

    const wrappedRequest: WrappedWorkerRequest = {
      seq,
      request,
    };
    entry.worker.postMessage(wrappedRequest);

    try {
      // oxlint-disable-next-line no-param-reassign
      entry.load += 1;
      return await result;
    } finally {
      // oxlint-disable-next-line no-param-reassign
      entry.load -= 1;
      this.#maybeRunCheckpoint();
    }
  }

  #maybeRunCheckpoint(): void {
    if (!this.#isReady || this.#checkpointPendingReason == null) {
      return;
    }

    for (const entry of this.#pool) {
      if (entry.load !== 0) {
        return;
      }
    }

    const reason = this.#checkpointPendingReason;
    this.#checkpointPendingReason = null;
    const primary = this.#pool[0];
    strictAssert(primary, 'Missing primary');
    void this.#send(primary, { type: 'walCheckpoint', reason });
  }

  async #terminate(request: WorkerRequest): Promise<void> {
    const primary = this.#pool[0];
    strictAssert(primary, 'Missing primary');

    const rest = this.#pool.slice(1);

    // Terminate non-primary workers first
    await Promise.all(rest.map(worker => this.#send(worker, request)));

    // Primary last
    await this.#send(primary, request);
  }

  #onError(errorKind: SqliteErrorKind, error: Error): void {
    const resolvers = new Array<(error: Error) => void>();
    this.#errorResolvers = this.#errorResolvers.filter(entry => {
      if (entry.kind === errorKind) {
        resolvers.push(entry.resolve);
        if (entry.once) {
          return false;
        }
      }
      return true;
    });

    for (const resolve of resolvers) {
      resolve(error);
    }
  }

  #resetQueryStats() {
    this.#queryStats = { start: Date.now(), statsByQuery: new Map() };
  }

  #roundDuration(duration: number): number {
    return Math.round(100 * duration) / 100;
  }

  #logQueryStats({ maxQueriesToLog = 10, epochName }: QueryStatsOptions) {
    if (!this.#queryStats) {
      return;
    }
    const epochDuration = Date.now() - this.#queryStats.start;
    const sortedByCumulativeDuration = [
      ...this.#queryStats.statsByQuery.values(),
    ].sort((a, b) => (b.cumulative ?? 0) - (a.cumulative ?? 0));
    const cumulativeDuration = sortedByCumulativeDuration.reduce(
      (sum, stats) => sum + stats.cumulative,
      0
    );
    this.#logger?.info(
      `Top ${maxQueriesToLog} queries by cumulative duration (ms) over last ${epochDuration}ms` +
        `${epochName ? ` during '${epochName}'` : ''}: ` +
        sortedByCumulativeDuration
          .slice(0, maxQueriesToLog)
          .map(stats => {
            return (
              `${stats.queryName}: cumulative ${this.#roundDuration(stats.cumulative)} | ` +
              `average: ${this.#roundDuration(stats.cumulative / (stats.count || 1))} | ` +
              `max: ${this.#roundDuration(stats.max)} | ` +
              `count: ${stats.count}`
            );
          })
          .join(' ||| ') +
        `; Total cumulative duration of all SQL queries during this epoch: ${this.#roundDuration(cumulativeDuration)}ms`
    );
  }

  #traceDuration(method: string, duration: number): void {
    if (this.#shouldTrackQueryStats) {
      if (!this.#queryStats) {
        this.#resetQueryStats();
      }
      strictAssert(this.#queryStats, 'has been initialized');
      let currentStats = this.#queryStats.statsByQuery.get(method);
      if (!currentStats) {
        currentStats = { count: 0, cumulative: 0, queryName: method, max: 0 };
        this.#queryStats.statsByQuery.set(method, currentStats);
      }
      currentStats.count += 1;
      currentStats.cumulative += duration;
      currentStats.max = Math.max(currentStats.max, duration);
    }

    if (this.#shouldLogQueryTime(method) && !app.isPackaged) {
      const twoDecimals = this.#roundDuration(duration);
      this.#logger?.info(`MainSQL query: ${method}, duration=${twoDecimals}ms`);
    }
    if (duration > MIN_TRACE_DURATION) {
      strictAssert(this.#logger !== undefined, 'Logger not initialized');
      this.#logger.info(
        `MainSQL: slow query ${method} duration=${Math.round(duration)}ms`
      );
    }
  }

  #createWorker(): CreateWorkerResultType {
    const scriptPath = join(getAppRootDir(), 'bundles', 'workers', 'sql.js');

    const worker = new Worker(scriptPath);

    worker.on('message', (wrappedResponse: WrappedWorkerResponse) => {
      if (wrappedResponse.type === 'log') {
        const { level, args } = wrappedResponse;
        strictAssert(this.#logger !== undefined, 'Logger not initialized');
        this.#logger[level](`MainSQL: ${format(...args)}`);
        return;
      }

      if (wrappedResponse.type === 'walCheckpointNeeded') {
        this.#checkpointPendingReason = wrappedResponse.reason;
        this.#maybeRunCheckpoint();
        return;
      }

      const { seq, error, errorKind, response } = wrappedResponse;

      const entry = this.#onResponse.get(seq);
      this.#onResponse.delete(seq);
      if (!entry) {
        throw new Error(`Unexpected worker response with seq: ${seq}`);
      }

      if (error) {
        const errorObj = new Error(`${entry.errorId}: ${error.message}`);
        errorObj.stack = `${entry.errorId}: ${error.stack}`;
        errorObj.name = error.name;
        this.#onError(errorKind ?? SqliteErrorKind.Unknown, errorObj);

        entry.reject(errorObj);
      } else {
        entry.resolve(response);
      }
    });

    const { promise: onExit, resolve: resolveOnExit } = explodePromise<void>();
    worker.once('exit', resolveOnExit);

    return { worker, onExit };
  }

  // Find first pool entry with minimal load
  #getWorker(): PoolEntry {
    let min = this.#pool[0];
    strictAssert(min, 'Missing min');

    for (const entry of this.#pool) {
      if (min && min.load < entry.load) {
        continue;
      }

      min = entry;
    }
    return min;
  }
}
