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

const UNLOCKED = 0;
const WRITE_LOCKED = -1;

/** Disposes by calling `readUnlock()` on the associated rwlock. */
export class ReadGuard {
  private readonly rwlock: RwLock;

  constructor(rwlock: RwLock) {
    this.rwlock = rwlock;
  }

  [Symbol.dispose](): void {
    this.rwlock.readUnlock();
  }
}

/** Disposes by calling `writeUnlock()` on the associated rwlock. */
export class WriteGuard {
  private readonly rwlock: RwLock;

  constructor(rwlock: RwLock) {
    this.rwlock = rwlock;
  }

  [Symbol.dispose](): void {
    this.rwlock.writeUnlock();
  }
}

/** An async reader-writer lock backed by SharedArrayBuffer. Multiple readers or one exclusive writer. */
export class RwLock {
  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 a read lock. Multiple readers can hold the lock concurrently.
   * @returns A disposable {@link ReadGuard} that releases the read lock on dispose.
   */
  async readLock(): Promise<ReadGuard> {
    while (true) {
      const state = Atomics.load(this.view, 0);
      // Can acquire if not write-locked
      if (state >= UNLOCKED) {
        if (Atomics.compareExchange(this.view, 0, state, state + 1) === state) {
          return new ReadGuard(this);
        }
        continue; // CAS failed, retry immediately
      }
      // Write-locked, wait for change
      const result = Atomics.waitAsync(this.view, 0, state);
      if (result.async) {
        await result.value;
      }
    }
  }

  /**
   * Attempts to acquire a read lock without waiting.
   * Makes a single atomic attempt; returns `null` on contention even from other readers.
   * @returns A disposable {@link ReadGuard} if acquired, otherwise `null`.
   */
  tryReadLock(): ReadGuard | null {
    const state = Atomics.load(this.view, 0);
    if (state >= UNLOCKED && Atomics.compareExchange(this.view, 0, state, state + 1) === state) {
      return new ReadGuard(this);
    }
    return null;
  }

  /** Releases a read lock. Wakes a waiting writer if this was the last reader. */
  readUnlock(): void {
    const prev = Atomics.sub(this.view, 0, 1);
    if (prev === 1) {
      // Last reader, wake a waiting writer
      Atomics.notify(this.view, 0);
    }
  }

  /**
   * Acquires an exclusive write lock. Waits for all readers and writers to release.
   * @returns A disposable {@link WriteGuard} that releases the write lock on dispose.
   */
  async writeLock(): Promise<WriteGuard> {
    while (true) {
      // Can only acquire from unlocked state
      if (Atomics.compareExchange(this.view, 0, UNLOCKED, WRITE_LOCKED) === UNLOCKED) {
        return new WriteGuard(this);
      }
      // Someone holds the lock, wait for change
      const state = Atomics.load(this.view, 0);
      if (state !== UNLOCKED) {
        const result = Atomics.waitAsync(this.view, 0, state);
        if (result.async) {
          await result.value;
        }
      }
    }
  }

  /**
   * Attempts to acquire a write lock without waiting.
   * @returns A disposable {@link WriteGuard} if acquired, otherwise `null`.
   */
  tryWriteLock(): WriteGuard | null {
    if (Atomics.compareExchange(this.view, 0, UNLOCKED, WRITE_LOCKED) === UNLOCKED) {
      return new WriteGuard(this);
    }
    return null;
  }

  /** Releases the write lock and wakes all waiting threads. */
  writeUnlock(): void {
    Atomics.store(this.view, 0, UNLOCKED);
    // Wake all waiters — both readers and writers may be waiting
    Atomics.notify(this.view, 0, +Infinity);
  }

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

registerSync('RwLock', RwLock);