import type { MutexAcquireOptions, MutexOptions, MutexPermit } from "./mutex.ts"
import { Mutex } from "./mutex.ts"
import type { KeyedPermitDetails } from "./permit.ts"
import { Permit } from "./permit.ts"

/**
 * @description 表示 `KeyedLock` 成功获取后的释放句柄。
 *
 * 句柄中会保留当前命中的 key，便于调试或上层记录持锁对象。
 */
export type KeyedLockPermit<Key> = Permit<KeyedPermitDetails<"keyed-lock", Key>>

/**
 * @description 表示单个 key 对应的内部状态。
 *
 * 每个 key 都维护一个独立的 `Mutex`，并通过引用计数判断是否可以从状态表中清理。
 */
interface InternalKeyedLockState<Key> {
  key: Key
  mutex: Mutex
  referenceCount: number
}

/**
 * @description 表示构造 `KeyedLock` 时支持的配置。
 *
 * 当前直接复用 `Mutex` 的配置项，以保证每个 key 下的锁行为一致。
 */
export interface KeyedLockOptions extends MutexOptions {}

/**
 * @description 表示按 key 维度拆分互斥域的锁。
 *
 * 相同 key 的任务会串行执行，不同 key 之间则互不影响，适合做资源 ID、租户 ID、文件路径等粒度的独占控制。
 */
export class KeyedLock<Key> {
  private readonly stateMap: Map<Key, InternalKeyedLockState<Key>>

  private readonly mutexOptions: MutexOptions

  /**
   * @description 创建一个按 key 分桶的互斥锁管理器。
   */
  constructor(options: KeyedLockOptions = {}) {
    this.stateMap = new Map<Key, InternalKeyedLockState<Key>>()
    this.mutexOptions = options
  }

  /**
   * @description 判断当前是否已经为指定 key 建立内部状态。
   */
  hasKey(key: Key): boolean {
    return this.stateMap.has(key)
  }

  /**
   * @description 读取当前仍被追踪的 key 数量。
   */
  getKeyCount(): number {
    return this.stateMap.size
  }

  /**
   * @description 取得指定 key 的状态；若不存在则惰性创建。
   */
  private getOrCreateState(key: Key): InternalKeyedLockState<Key> {
    const existingState = this.stateMap.get(key)
    if (existingState !== undefined) {
      return existingState
    }

    const state: InternalKeyedLockState<Key> = {
      key,
      mutex: new Mutex(this.mutexOptions),
      referenceCount: 0,
    }
    this.stateMap.set(key, state)
    return state
  }

  /**
   * @description 在某个 key 不再被持有且没有等待者时回收对应状态。
   */
  private cleanupState(state: InternalKeyedLockState<Key>): void {
    if (state.referenceCount !== 0) {
      return
    }

    if (state.mutex.isLocked() === true || state.mutex.getPendingCount() !== 0) {
      return
    }

    const currentState = this.stateMap.get(state.key)
    if (currentState === state) {
      this.stateMap.delete(state.key)
    }
  }

  /**
   * @description 基于底层 `Mutex` permit 构造带 key 元信息的释放句柄。
   */
  private buildPermit(
    state: InternalKeyedLockState<Key>,
    permit: MutexPermit,
  ): KeyedLockPermit<Key> {
    return new Permit<KeyedPermitDetails<"keyed-lock", Key>>({
      details: {
        coordination: "keyed-lock",
        key: state.key,
      },
      duplicateReleaseMessage: "KeyedLock permit release was called more than once.",
      onDuplicateRelease: this.mutexOptions.onDuplicateRelease,
      onRelease: (): void => {
        permit.release()
        state.referenceCount = state.referenceCount - 1
        this.cleanupState(state)
      },
    })
  }

  /**
   * @description 尝试立即获取指定 key 的独占锁。
   */
  tryAcquire(key: Key): KeyedLockPermit<Key> | undefined {
    const state = this.getOrCreateState(key)
    const permit = state.mutex.tryAcquire()

    if (permit === undefined) {
      this.cleanupState(state)
      return undefined
    }

    state.referenceCount = state.referenceCount + 1
    return this.buildPermit(state, permit)
  }

  /**
   * @description 等待直到成功获取指定 key 的独占锁。
   */
  async acquire(key: Key, options: MutexAcquireOptions = {}): Promise<KeyedLockPermit<Key>> {
    const state = this.getOrCreateState(key)
    state.referenceCount = state.referenceCount + 1

    try {
      const permit = await state.mutex.acquire(options)
      return this.buildPermit(state, permit)
    } catch (error) {
      state.referenceCount = state.referenceCount - 1
      this.cleanupState(state)
      throw error
    }
  }

  /**
   * @description 在指定 key 的独占锁保护下执行回调。
   */
  async runExclusive<T>(
    key: Key,
    callback: () => Promise<T>,
    options?: MutexAcquireOptions,
  ): Promise<T>
  async runExclusive<T>(key: Key, callback: () => T, options?: MutexAcquireOptions): Promise<T>
  async runExclusive<T>(
    key: Key,
    callback: () => Promise<T> | T,
    options: MutexAcquireOptions = {},
  ): Promise<T> {
    const permit = await this.acquire(key, options)
    try {
      return await callback()
    } finally {
      permit.release()
    }
  }
}