import { registerSync } from './reconstruct.ts';

const UNLOCKED = 0;
const LOCKED = 1;

/** Disposes by calling `unlock()` on the associated mutex. */
export class MutexGuard {
  private readonly mutex: Mutex;

  constructor(mutex: Mutex) {
    this.mutex = mutex;
  }

  [Symbol.dispose](): void {
    this.mutex.unlock();
  }
}

/** An async mutual exclusion lock backed by SharedArrayBuffer. Works across threads. */
export class Mutex {
  static readonly byteSize = 4;
  static readonly byteAlignment = 4;
  private readonly view: Int32Array;

  constructor(buffer?: SharedArrayBuffer, byteOffset?: number) {
    const buf = buffer ?? new SharedArrayBuffer(4);
    const offset = byteOffset ?? 0;
    this.view = new Int32Array(buf, offset, 1);
  }

  /**
   * Acquires the lock, waiting asynchronously if held by another thread.
   * @returns A disposable {@link MutexGuard} that releases the lock on dispose.
   */
  async lock(): Promise<MutexGuard> {
    while (true) {
      // Try to acquire: CAS from UNLOCKED to LOCKED
      if (Atomics.compareExchange(this.view, 0, UNLOCKED, LOCKED) === UNLOCKED) {
        return new MutexGuard(this);
      }
      // Wait until notified that the lock might be free
      const result = Atomics.waitAsync(this.view, 0, LOCKED);
      if (result.async) {
        await result.value;
      }
    }
  }

  /** Releases the lock and wakes one waiting thread. */
  unlock(): void {
    Atomics.store(this.view, 0, UNLOCKED);
    Atomics.notify(this.view, 0, 1);
  }

  [Symbol.for('moroutine.shared')](): { tag: string; buffer: SharedArrayBuffer; byteOffset: number } {
    return { tag: 'Mutex', buffer: this.view.buffer as SharedArrayBuffer, byteOffset: this.view.byteOffset };
  }
}

registerSync('Mutex', Mutex);