/**
 * In-memory `CheckpointStore` adapter. Useful in tests, REPL
 * sessions, and small examples where SQLite would be overkill. The
 * production-grade adapter lives in `@graphorin/store-sqlite`.
 *
 * @packageDocumentation
 */

import type {
  Checkpoint,
  CheckpointId,
  CheckpointMetadata,
  CheckpointPutOptions,
  CheckpointStoreExt,
  CheckpointTuple,
  ListOptions,
  PendingWrite,
  PruneThreadsOptions,
} from '@graphorin/core';
import { CheckpointConflictError } from '@graphorin/core';

interface StoredCheckpoint {
  readonly checkpoint: Checkpoint;
  readonly metadata: CheckpointMetadata;
  pendingWrites: PendingWrite[];
}

/**
 * Pure in-memory `CheckpointStore` implementation. Thread-safe
 * within a single Node.js event loop because every mutation is
 * synchronous; concurrent runs that share the same instance will see
 * a consistent view.
 *
 * @stable
 */
export class InMemoryCheckpointStore implements CheckpointStoreExt {
  #checkpoints = new Map<string, StoredCheckpoint>();
  #threadIndex = new Map<string, string[]>();

  async put(
    threadId: string,
    namespace: string,
    checkpoint: Checkpoint,
    metadata: CheckpointMetadata,
    opts?: CheckpointPutOptions,
  ): Promise<CheckpointId> {
    // D1 / workflow-01: honour the atomic compare-and-set contract. The
    // whole method body is synchronous, so the check + insert cannot be
    // interleaved by another writer on the same event loop.
    if (opts?.expectedLatestId !== undefined) {
      const ids = this.#threadIndex.get(`${threadId}::${namespace}`) ?? [];
      const latest = this.#latestStored(threadId, namespace, ids);
      const latestId = latest?.checkpoint.id ?? null;
      if (latestId !== opts.expectedLatestId) {
        throw new CheckpointConflictError(threadId, opts.expectedLatestId, latestId);
      }
    }
    const key = makeKey(threadId, namespace, checkpoint.id);
    const existing = this.#checkpoints.get(key);
    this.#checkpoints.set(key, {
      checkpoint,
      metadata,
      pendingWrites: existing ? [...existing.pendingWrites] : [],
    });
    const indexKey = `${threadId}::${namespace}`;
    const ids = this.#threadIndex.get(indexKey) ?? [];
    if (!ids.includes(checkpoint.id)) {
      ids.push(checkpoint.id);
      this.#threadIndex.set(indexKey, ids);
    }
    return checkpoint.id;
  }

  async putWrites(
    threadId: string,
    namespace: string,
    checkpointId: CheckpointId,
    writes: ReadonlyArray<PendingWrite>,
    taskId: string,
  ): Promise<void> {
    if (writes.length === 0) return;
    const key = makeKey(threadId, namespace, checkpointId);
    const stored = this.#checkpoints.get(key);
    if (!stored) return;
    for (const write of writes) {
      const annotated: PendingWrite = { ...write, taskId };
      stored.pendingWrites = [
        ...stored.pendingWrites.filter(
          (w) => !(w.taskId === annotated.taskId && w.index === annotated.index),
        ),
        annotated,
      ];
    }
  }

  async getTuple(
    threadId: string,
    namespace: string,
    checkpointId?: CheckpointId,
  ): Promise<CheckpointTuple | null> {
    const indexKey = `${threadId}::${namespace}`;
    const ids = this.#threadIndex.get(indexKey);
    if (!ids || ids.length === 0) return null;

    const stored =
      checkpointId !== undefined
        ? this.#checkpoints.get(makeKey(threadId, namespace, checkpointId))
        : this.#latestStored(threadId, namespace, ids);
    if (!stored) return null;

    return stored.pendingWrites.length > 0
      ? {
          checkpoint: stored.checkpoint,
          metadata: stored.metadata,
          pendingWrites: [...stored.pendingWrites],
        }
      : { checkpoint: stored.checkpoint, metadata: stored.metadata };
  }

  async *list(
    threadId: string,
    namespace: string,
    opts?: ListOptions,
  ): AsyncIterable<CheckpointTuple> {
    const indexKey = `${threadId}::${namespace}`;
    const ids = this.#threadIndex.get(indexKey) ?? [];
    const ordered = [...ids]
      .map((id) => this.#checkpoints.get(makeKey(threadId, namespace, id)))
      .filter((s): s is StoredCheckpoint => s !== undefined)
      .sort((a, b) => b.checkpoint.stepNumber - a.checkpoint.stepNumber);

    let beforeStep: number | null = null;
    if (opts?.before !== undefined) {
      const cursor = this.#checkpoints.get(makeKey(threadId, namespace, opts.before));
      beforeStep = cursor ? cursor.checkpoint.stepNumber : null;
    }

    let yielded = 0;
    const limit = opts?.limit ?? Number.POSITIVE_INFINITY;

    for (const stored of ordered) {
      if (yielded >= limit) return;
      if (beforeStep !== null && stored.checkpoint.stepNumber >= beforeStep) continue;
      if (opts?.status !== undefined && stored.metadata.status !== opts.status) continue;
      yielded += 1;
      yield stored.pendingWrites.length > 0
        ? {
            checkpoint: stored.checkpoint,
            metadata: stored.metadata,
            pendingWrites: [...stored.pendingWrites],
          }
        : { checkpoint: stored.checkpoint, metadata: stored.metadata };
    }
  }

  async deleteThread(threadId: string): Promise<void> {
    for (const key of [...this.#threadIndex.keys()]) {
      if (key.startsWith(`${threadId}::`)) {
        const ids = this.#threadIndex.get(key) ?? [];
        const namespace = key.slice(threadId.length + 2);
        for (const id of ids) {
          this.#checkpoints.delete(makeKey(threadId, namespace, id));
        }
        this.#threadIndex.delete(key);
      }
    }
  }

  /**
   * Enumerate threads whose LATEST checkpoint in `namespace` is
   * suspended with a due `wakeAt` - parity with the SQLite adapter.
   */
  async listSuspended(
    namespace: string,
    opts?: { readonly dueBefore?: number; readonly limit?: number },
  ): Promise<ReadonlyArray<{ readonly threadId: string; readonly wakeAt: number }>> {
    const out: Array<{ threadId: string; wakeAt: number }> = [];
    for (const key of this.#threadIndex.keys()) {
      const sep = key.indexOf('::');
      const threadId = key.slice(0, sep);
      const ns = key.slice(sep + 2);
      if (ns !== namespace) continue;
      const ids = this.#threadIndex.get(key) ?? [];
      const latest = this.#latestStored(threadId, ns, ids);
      if (latest === undefined || latest.metadata.status !== 'suspended') continue;
      const wakeAt = latest.metadata.wakeAt;
      if (typeof wakeAt !== 'number') continue;
      if (opts?.dueBefore !== undefined && wakeAt > opts.dueBefore) continue;
      out.push({ threadId, wakeAt });
    }
    out.sort((a, b) => a.wakeAt - b.wakeAt);
    return opts?.limit !== undefined ? out.slice(0, opts.limit) : out;
  }

  /**
   * Retention sweep - parity with the SQLite implementation:
   * namespace-SCOPED (entries key as `threadId::namespace`), latest
   * checkpoint decides age + status, suspended pairs survive unless
   * `onlyTerminal: false`.
   */
  async pruneThreads(opts: PruneThreadsOptions): Promise<number> {
    const onlyTerminal = opts.onlyTerminal !== false;
    let pruned = 0;
    for (const key of [...this.#threadIndex.keys()]) {
      const sep = key.indexOf('::');
      const threadId = key.slice(0, sep);
      const namespace = key.slice(sep + 2);
      const ids = this.#threadIndex.get(key) ?? [];
      const latest = this.#latestStored(threadId, namespace, ids);
      if (latest === undefined) continue;
      if (Date.parse(latest.checkpoint.createdAt) >= opts.beforeEpochMs) continue;
      if (
        onlyTerminal &&
        latest.metadata.status !== 'completed' &&
        latest.metadata.status !== 'failed' &&
        latest.metadata.status !== 'aborted'
      ) {
        continue;
      }
      for (const id of ids) {
        this.#checkpoints.delete(makeKey(threadId, namespace, id));
      }
      this.#threadIndex.delete(key);
      pruned += 1;
    }
    return pruned;
  }

  /** Compaction - keep the `keepLast` newest checkpoints of one pair. */
  async compactThread(threadId: string, namespace: string, keepLast: number): Promise<number> {
    const keep = Math.max(1, Math.floor(keepLast));
    const indexKey = `${threadId}::${namespace}`;
    const ids = this.#threadIndex.get(indexKey) ?? [];
    const ordered = [...ids]
      .map((id) => this.#checkpoints.get(makeKey(threadId, namespace, id)))
      .filter((s): s is StoredCheckpoint => s !== undefined)
      .sort((a, b) => b.checkpoint.stepNumber - a.checkpoint.stepNumber);
    const victims = ordered.slice(keep);
    for (const victim of victims) {
      this.#checkpoints.delete(makeKey(threadId, namespace, victim.checkpoint.id));
    }
    this.#threadIndex.set(
      indexKey,
      ordered.slice(0, keep).map((s) => s.checkpoint.id),
    );
    return victims.length;
  }

  /**
   * Test-only helper that exposes the raw count of stored checkpoints
   * - handy for assertions like "the runtime wrote exactly N
   * checkpoints across the run".
   */
  size(): number {
    return this.#checkpoints.size;
  }

  #latestStored(
    threadId: string,
    namespace: string,
    ids: ReadonlyArray<string>,
  ): StoredCheckpoint | undefined {
    let latest: StoredCheckpoint | undefined;
    for (const id of ids) {
      const stored = this.#checkpoints.get(makeKey(threadId, namespace, id));
      if (!stored) continue;
      if (!latest || stored.checkpoint.stepNumber > latest.checkpoint.stepNumber) {
        latest = stored;
      }
    }
    return latest;
  }
}

function makeKey(threadId: string, namespace: string, id: string): string {
  return `${threadId}::${namespace}::${id}`;
}