//
// Copyright 2022 DXOS.org
//

import { Context, ContextDisposedError } from '@dxos/context';
import { StackTrace } from '@dxos/debug';
import { type MaybePromise } from '@dxos/util';

import { trackResource } from './track-leaks.ts';
import { Trigger } from './trigger.ts';

/**
 * A task that can be scheduled to run in the next event loop iteration.
 * Could be triggered multiple times, but only runs once.
 * If a new task is triggered while a previous one is running, the next run would occur immediately after the current run has finished.
 */
// TODO(dmaretskyi): Consider calling `join` on context dispose.
// TODO(burdon): Add throttling support.
export class DeferredTask {
  private _scheduled = false;
  private _currentTask: Promise<void> | null = null; // Can't be rejected.
  private _nextTask = new Trigger();

  constructor(
    private readonly _ctx: Context,
    private readonly _callback: () => Promise<void>,
  ) {}

  get scheduled() {
    return this._scheduled;
  }

  /**
   * Schedule the task to run asynchronously.
   */
  schedule(): void {
    if (this._scheduled) {
      return; // Already scheduled.
    }

    scheduleTask(this._ctx, async () => {
      // The previous task might still be running, so we need to wait for it to finish.
      //
      // A single await suffices only because this is the one claim site and `_scheduled` collapses
      // schedules into one waiter. A second path claiming `_currentTask` directly would let two
      // callbacks run at once — see dxos/edge#758.
      await this._currentTask; // Can't be rejected.

      // Reset the flag. New tasks can now be scheduled. They would wait for the callback to finish.
      this._scheduled = false;
      const completionTrigger = this._nextTask;
      this._nextTask = new Trigger(); // Re-create the trigger as opposed to resetting it since there might be listeners waiting for it.

      // Store the promise so that new tasks could wait for this one to finish.
      this._currentTask = runInContextAsync(this._ctx, () => this._callback()).then(() => {
        completionTrigger.wake();
      });
    });

    this._scheduled = true;
  }

  /**
   * Schedule the task to run and wait for it to finish.
   */
  async runBlocking(): Promise<void> {
    if (this._ctx.disposed) {
      throw new ContextDisposedError();
    }

    this.schedule();
    await this._nextTask.wait();
  }

  /**
   * Waits for the current task to finish if it is running.
   * Does not schedule a new task.
   */
  async join(): Promise<void> {
    await this._currentTask;
  }
}

// TODO(dmaretskyi): Protyping this an alternative API to DeferredTask.
export class AsyncTask {
  #callback: () => Promise<void>;
  #ctx?: Context = undefined;

  #scheduled = false;
  #currentTask: Promise<void> | null = null; // Can't be rejected.
  #nextTask = new Trigger();

  constructor(callback: () => Promise<void>) {
    this.#callback = callback;
  }

  get scheduled() {
    return this.#scheduled;
  }

  /**
   * Context of the resource that owns the task.
   * When the context is disposed, the task is cancelled and cannot be scheduled again.
   */
  open(): void {
    this.#ctx = new Context();
  }

  /**
   * Closes the task and waits for it to finish if it is running.
   */
  async close(): Promise<void> {
    await this.#ctx?.dispose();
    await this.join();
    this.#ctx = undefined;
  }

  [Symbol.asyncDispose](): Promise<void> {
    return this.close();
  }

  /**
   * Schedule the task to run asynchronously.
   */
  // TODO(dmaretskyi): Add scheduleAt. Where the earlier time will override the later one.
  schedule(): void {
    if (!this.#ctx || this.#ctx.disposed) {
      throw new Error('AsyncTask not open');
    }

    if (this.#scheduled) {
      return; // Already scheduled.
    }

    scheduleTask(this.#ctx, async () => {
      // The previous task might still be running, so we need to wait for it to finish.
      await this.#currentTask; // Can't be rejected.

      if (!this.#ctx || this.#ctx.disposed) {
        return;
      }

      // Reset the flag. New tasks can now be scheduled. They would wait for the callback to finish.
      this.#scheduled = false;
      const completionTrigger = this.#nextTask;
      this.#nextTask = new Trigger(); // Re-create the trigger as opposed to resetting it since there might be listeners waiting for it.

      // Store the promise so that new tasks could wait for this one to finish.
      this.#currentTask = runInContextAsync(this.#ctx, () => this.#callback()).then(() => {
        completionTrigger.wake();
      });
    });

    this.#scheduled = true;
  }

  /**
   * Schedule the task to run and wait for it to finish.
   */
  async runBlocking(): Promise<void> {
    if (this.#ctx?.disposed) {
      throw new ContextDisposedError();
    }

    this.schedule();
    await this.#nextTask.wait();
  }

  /**
   * Waits for the current task to finish if it is running.
   * Does not schedule a new task.
   */
  async join(): Promise<void> {
    await this.#currentTask;
  }
}

export const runInContext = (ctx: Context, fn: () => void) => {
  try {
    fn();
  } catch (err: any) {
    ctx.raise(err);
  }
};

export const runInContextAsync = async (ctx: Context, fn: () => MaybePromise<void>) => {
  try {
    await fn();
  } catch (err: any) {
    ctx.raise(err);
  }
};

export const scheduleMicroTask = (ctx: Context, fn: () => MaybePromise<void>) => {
  queueMicrotask(async () => {
    if (ctx.disposed) {
      return;
    }
    await runInContextAsync(ctx, fn);
  });
};

export const scheduleTask = (ctx: Context, fn: () => MaybePromise<void>, afterMs?: number) => {
  const clearTracking = trackResource(() => ({
    name: `task (${fn.name || 'anonymous'})`,
    openStack: new StackTrace(),
  }));

  const timeout = setTimeout(async () => {
    clearDispose();
    await runInContextAsync(ctx, fn);
    clearTracking();
  }, afterMs);

  const clearDispose = ctx.onDispose(() => {
    clearTracking();
    clearTimeout(timeout);
  });
};

/**
 * Run the task in the next event loop iteration, and then repeat in `interval` ms after the previous iteration completes.
 */
export const scheduleTaskInterval = (ctx: Context, task: () => Promise<void>, interval: number) => {
  const clearTracking = trackResource(() => ({
    name: `repeating task (${task.name || 'anonymous'})`,
    openStack: new StackTrace(),
  }));

  let timeoutId: NodeJS.Timeout;

  const run = async () => {
    await runInContextAsync(ctx, task);
    if (ctx.disposed) {
      return;
    }
    timeoutId = setTimeout(run, interval);
  };

  timeoutId = setTimeout(run, interval);
  ctx.onDispose(() => {
    clearTracking();
    clearTimeout(timeoutId);
  });
};

export const scheduleExponentialBackoffTaskInterval = (
  ctx: Context,
  task: () => Promise<void>,
  initialInterval: number,
) => {
  const clearTracking = trackResource(() => ({
    name: `repeating task (${task.name || 'anonymous'})`,
    openStack: new StackTrace(),
  }));

  let timeoutId: NodeJS.Timeout;

  let interval = initialInterval;
  const repeat = async () => {
    await runInContextAsync(ctx, task);
    if (ctx.disposed) {
      return;
    }
    interval *= 2;
    timeoutId = setTimeout(repeat, interval);
  };

  timeoutId = setTimeout(repeat, interval);
  ctx.onDispose(() => {
    clearTracking();
    clearTimeout(timeoutId);
  });
};