import type {
  ActivityRecord,
  ModelMessage,
  MetadataStore,
  PersistedArtifactRef,
  RunStatus,
  RunStore,
  Scope,
  TokenUsage,
} from '@tanstack/ai'

// Re-export the shared identity type so app code can import Scope from either
// `@tanstack/ai` or `@tanstack/ai-persistence`. See {@link Scope} security notes:
// pair a client-visible `threadId` with a server-trusted `userId`/`tenantId`
// before authorizing load/save (e.g. via `reconstructChat({ authorize })`).
export type { ActivityRecord, MetadataStore, Scope }

// ===========================================================================
// Store contracts
// ===========================================================================
//
// EVOLUTION POLICY
// ----------------
// These store interfaces are the compatibility surface between the core
// middleware and every backend — the in-memory reference store and every
// adapter an application writes against its own database.
//
//   - Store METHODS are REQUIRED. A new method is a breaking contract change:
//     every adapter gets a compile error and implements it. Do NOT add methods
//     as optional-and-feature-detected (`store.method?.(...)`) — an adapter
//     that has not implemented one is then indistinguishable from one whose
//     answer is legitimately empty, so the feature silently does nothing in
//     production instead of failing at build time. `findActiveRun` was optional
//     for exactly one release cycle and cost us precisely that: reconnect
//     degraded to "no active run" on every backend that had not caught up.
//   - Capability tiers belong at the STORE level, not the method level. A
//     backend that only stores a transcript declares `ChatTranscriptStores`
//     (no `runs`); it does not declare a half-implemented `RunStore`.
//   - Never tighten an existing method's required arguments or widen its
//     required return shape in a breaking way.
//
// The shared conformance testkit (`./testkit/conformance.ts`) is the
// authoritative compatibility gate: every invariant documented on the methods
// below is asserted there, and every backend runs the identical suite. If an
// invariant is not encoded in the testkit, adapters cannot discover it — so
// promote new invariants into both the JSDoc here AND the testkit.
//
// TIMESTAMP CONVENTION
// --------------------
// Store *records* (`RunRecord`, `InterruptRecord`, `ArtifactRecord`,
// `BlobRecord`) speak **epoch milliseconds** (`number`), the native unit for
// SQL/`BIGINT` columns and `Date.now()`. Wire/result references that leave the
// persistence layer (e.g. core's `PersistedArtifactRef.createdAt`) speak
// **ISO-8601 strings**. The middleware performs the number→ISO conversion at
// the boundary; do not mix the two on a single field.

/**
 * One page of a thread from {@link MessageStore.loadThread} when the caller
 * passed a paging hint.
 *
 * Middleware omits the hint and always gets a full `Array<ModelMessage>`,
 * never this shape.
 *
 * `truncated: true` requires `cursor`. Without a cursor the client cannot
 * request the next older window, so `reconstructChat` treats that page as
 * complete.
 */
export type MessagePage =
  | {
      messages: Array<ModelMessage>
      truncated: false
      cursor?: never
    }
  | {
      messages: Array<ModelMessage>
      truncated: true
      cursor: string
    }

/**
 * Durable store for a thread's full message transcript.
 *
 * A "thread" is the unit of conversation history. The key is
 * {@link Scope.threadId} (the same conversation id as
 * `ChatMiddlewareContext.threadId`). Store methods take a bare string for
 * adapter simplicity; multi-user isolation is the **host's** job — authorize
 * against `Scope.userId` / `Scope.tenantId` (derived server-side from session)
 * before calling load/save, and never treat a client-supplied thread id alone
 * as an ownership proof (see `Scope` security notes in `@tanstack/ai`).
 *
 * `saveThread` always receives and persists the **complete, authoritative**
 * message list — it is an overwrite, never an append. The middleware snapshots
 * `ctx.messages` (the full running transcript) into it.
 */
export interface MessageStore {
  /**
   * Return the stored transcript for `threadId` ({@link Scope.threadId}),
   * in insertion order.
   *
   * Call with only `threadId` (middleware, `onStart`, `onFinish`) and this
   * MUST return the full transcript as an `Array<ModelMessage>`. Never a
   * {@link MessagePage}.
   *
   * `options.limit` and `options.before` are an optional paging hint for
   * hydrate. Adapters may ignore them and still return the full array. An
   * adapter that pages returns a {@link MessagePage}.
   *
   * INVARIANT: returns an empty array (never `null`/`undefined`) for a thread
   * that was never saved. Callers treat `[]` as "no history".
   */
  loadThread: {
    (threadId: string): Promise<Array<ModelMessage>>
    (
      threadId: string,
      options: { limit?: number; before?: string },
    ): Promise<Array<ModelMessage> | MessagePage>
  }
  /**
   * Overwrite the stored transcript for `threadId` with `messages`.
   *
   * INVARIANT: this is a full replace. `messages` is the complete authoritative
   * history; the previous contents are discarded (not merged or appended).
   */
  saveThread: (threadId: string, messages: Array<ModelMessage>) => Promise<void>
}

/**
 * Durable sidecar for frontend-only AG-UI activity.
 *
 * Activity is never a {@link ModelMessage} and must not live in
 * {@link MessageStore}. This store is optional: backends that omit it keep
 * today's behavior (activity is not durable on the server).
 *
 * `saveActivities` always receives and persists the **complete, authoritative**
 * activity list — it is an overwrite, never an append.
 */
export interface ActivityStore {
  /**
   * Return the stored activity rows for `threadId`, in the order they were
   * last saved.
   *
   * INVARIANT: returns an empty array (never `null`/`undefined`) for a thread
   * that was never saved. Callers treat `[]` as "no activity".
   */
  loadActivities: (threadId: string) => Promise<Array<ActivityRecord>>
  /**
   * Overwrite the stored activity rows for `threadId` with `activities`.
   *
   * INVARIANT: this is a full replace. `activities` is the complete
   * authoritative list; the previous contents are discarded.
   */
  saveActivities: (
    threadId: string,
    activities: Array<ActivityRecord>,
  ) => Promise<void>
}

// Run lifecycle types live in `@tanstack/ai` and are re-exported here: one run,
// one record — shared by this package's `runs` store and `@tanstack/ai-sandbox`'s
// run driver, instead of each package keeping a rival definition that can drift.
export type {
  RunStatus,
  TerminalRunStatus,
  RunRecord,
  RunStore,
} from '@tanstack/ai'
export { isTerminalRunStatus, defineRunStore } from '@tanstack/ai'

/**
 * Lifecycle status of a generation run. Deliberately the same vocabulary as
 * {@link RunStatus}, so an adapter that stores both kinds of run can share one
 * status column and one set of checks.
 */
export type GenerationRunStatus = RunStatus

/**
 * A single generation run (one `generateImage` / `generateVideo` / … call).
 *
 * Its primary identity is `runId`: the run/request id the activity mints, the
 * same AG-UI run id the client sends on the wire. `threadId` is the SLOT the
 * run fills, a stable app-chosen name that groups successive runs of the same
 * thing, and it is what a server-driven client hydrates by. Generation state is
 * kept here, never in the chat {@link RunStore}.
 *
 * `result` holds terminal result METADATA (ids, model, urls, a provider video
 * job id), never the media bytes — those live in a {@link BlobStore}.
 * `artifacts` are the durable {@link PersistedArtifactRef}s, present only when
 * byte storage is on.
 *
 * @property startedAt - Epoch ms when the run was first created.
 * @property finishedAt - Epoch ms when the run reached a terminal status.
 */
export interface GenerationRunRecord {
  runId: string
  /**
   * The scope this run belongs to: a stable, app-chosen name for the slot
   * successive runs fill (`product-123-hero`, `video-9-start-frame`).
   *
   * REQUIRED, per the store-contract rule at the top of this file.
   * {@link GenerationRunStore.findLatestForThread} is the only query that
   * hydrates a run, and it keys on this — so a record without one can be
   * written and then never found again. `withGenerationPersistence` already
   * refuses to start a run without a scope, and a server-driven client
   * discards a snapshot that arrives without one, so an optional field here
   * only described a record no path could produce and no client would accept.
   */
  threadId: string
  /** `'image' | 'audio' | 'tts' | 'video' | 'transcription'`. */
  activity: string
  provider: string
  model: string
  status: GenerationRunStatus
  startedAt: number
  finishedAt?: number
  error?: { message: string; code?: string }
  /** Terminal result metadata (ids, model, urls). Never the media bytes. */
  result?: unknown
  /** Durable artifact references, when an artifacts + blobs backend is used. */
  artifacts?: Array<PersistedArtifactRef>
  usage?: TokenUsage
}

/**
 * Durable store for generation run records, the generation counterpart to
 * {@link RunStore}. Keyed by its own `runId`, with `threadId` the slot
 * {@link GenerationRunStore.findLatestForThread} looks runs up by.
 */
export interface GenerationRunStore {
  /**
   * Create a run record, or return the existing one if `runId` is already
   * present (resume).
   *
   * INVARIANT (idempotency): a second call for a `runId` returns the existing
   * record unchanged; `startedAt`/`activity`/`provider`/`model`/`threadId` are
   * not mutated. `status` defaults to `'running'` on first creation.
   */
  createOrResume: (
    input: Pick<
      GenerationRunRecord,
      'runId' | 'threadId' | 'activity' | 'provider' | 'model' | 'startedAt'
    > & { status?: GenerationRunStatus },
  ) => Promise<GenerationRunRecord>
  /**
   * Patch a run record's mutable fields.
   *
   * INVARIANT: patching a `runId` that does not exist is a **no-op** — it must
   * not throw and must not create a record.
   */
  update: (
    runId: string,
    patch: Partial<
      Pick<
        GenerationRunRecord,
        'status' | 'finishedAt' | 'error' | 'result' | 'artifacts' | 'usage'
      >
    >,
  ) => Promise<void>
  /** Return the run record for `runId`, or `null` if none exists. */
  get: (runId: string) => Promise<GenerationRunRecord | null>
  /**
   * The most recent run linked to `threadId`, or `null`.
   *
   * REQUIRED, per the store-contract rule at the top of this file: a
   * server-authoritative client hydrates by the stable thread id on every
   * mount, so an adapter without this would be indistinguishable from one that
   * legitimately has no run — `persistence: true` would silently restore
   * nothing, forever. `null` is the correct answer only when the thread really
   * has no runs. The chat parallel is {@link RunStore.findActiveRun}.
   */
  findLatestForThread: (threadId: string) => Promise<GenerationRunRecord | null>
}

/** Lifecycle status of a human-in-the-loop interrupt. */
export type InterruptStatus = 'pending' | 'resolved' | 'cancelled'

/**
 * A human-in-the-loop interrupt (tool approval, client-tool input request, …).
 *
 * @property requestedAt - Epoch ms when the interrupt was created.
 * @property resolvedAt - Epoch ms when the interrupt was resolved/cancelled;
 *   absent while pending.
 */
export interface InterruptRecord {
  interruptId: string
  runId: string
  threadId: string
  status: InterruptStatus
  requestedAt: number
  resolvedAt?: number
  payload: Record<string, unknown>
  response?: unknown
}

/** A terminal interrupt write for {@link InterruptStore.commitBatch}. */
export type InterruptCommitEntry =
  | {
      interruptId: string
      status: 'resolved'
      response?: unknown
    }
  | {
      interruptId: string
      status: 'cancelled'
    }

/** Durable store for human-in-the-loop interrupts. */
export interface InterruptStore {
  /**
   * Persist a new interrupt in the `'pending'` state.
   *
   * The record is accepted without `status`/`resolvedAt` so a "born resolved"
   * interrupt is unrepresentable — every interrupt begins pending and only
   * `resolve`/`cancel` may move it to a terminal state.
   *
   * INVARIANT (insert-if-absent): if an interrupt with the same `interruptId`
   * already exists, `create` is a **no-op** — it must NOT overwrite the
   * existing record. This is the canonical behaviour (SQL backends implement it
   * via `ON CONFLICT DO NOTHING` / upsert-with-empty-update), so a duplicate
   * create can never clobber a resolved interrupt back to pending.
   */
  create: (
    record: Omit<InterruptRecord, 'status' | 'resolvedAt'>,
  ) => Promise<void>
  /**
   * Move an interrupt to `'resolved'`, stamping `resolvedAt` and storing
   * `response`. A no-op if `interruptId` does not exist.
   */
  resolve: (interruptId: string, response?: unknown) => Promise<void>
  /**
   * Move an interrupt to `'cancelled'`, stamping `resolvedAt`. A no-op if
   * `interruptId` does not exist.
   */
  cancel: (interruptId: string) => Promise<void>
  /**
   * Commit terminal writes for a validated resume batch.
   *
   * Optional. When present, `withPersistence` calls it once instead of
   * calling `resolve` and `cancel` for each entry. Apply every entry or none.
   *
   * Reject the whole batch (throw, writing nothing) when any entry has a
   * duplicate `interruptId`, references an `interruptId` that does not exist,
   * or references an interrupt whose status is not `'pending'`. This is
   * stricter than `resolve` / `cancel`, which are no-ops for a missing
   * `interruptId`.
   */
  commitBatch?: (entries: ReadonlyArray<InterruptCommitEntry>) => Promise<void>
  /** Return the interrupt for `interruptId`, or `null` if none exists. */
  get: (interruptId: string) => Promise<InterruptRecord | null>
  /**
   * All interrupts for a thread.
   *
   * INVARIANT: ordered by insertion (equivalently `requestedAt` ascending). SQL
   * backends MUST `ORDER BY requested_at` — the middleware and testkit rely on
   * this stable ordering.
   */
  list: (threadId: string) => Promise<Array<InterruptRecord>>
  /** Pending interrupts for a thread, ordered by `requestedAt` ascending. */
  listPending: (threadId: string) => Promise<Array<InterruptRecord>>
  /** All interrupts for a run, ordered by `requestedAt` ascending. */
  listByRun: (runId: string) => Promise<Array<InterruptRecord>>
  /** Pending interrupts for a run, ordered by `requestedAt` ascending. */
  listPendingByRun: (runId: string) => Promise<Array<InterruptRecord>>
}

// ===========================================================================
// Store typers
// ===========================================================================
//
// Identity helpers that type a store implementation inline: pass an object
// literal and get autocomplete + contract checking, with no separate
// `: MessageStore` return annotation. They compose into `defineAIPersistence`,
// which infers **exact presence** — a store you define becomes a defined,
// non-optional, autocompleted key on `persistence.stores`, and accessing a store
// you did not define is a compile error.
//
// ```ts
// const persistence = defineAIPersistence({
//   stores: {
//     messages: defineMessageStore({ loadThread, saveThread }),
//     runs: defineRunStore({ createOrResume, update, get, findActiveRun }),
//   },
// })
// persistence.stores.runs        // RunStore (defined)
// persistence.stores.interrupts  // compile error — not provided
// ```
//
// Presence is per STORE, not per method: every method of a store you define is
// required (see the evolution policy above). Omitting one is a compile error,
// not a partial store.

/** Type a {@link MessageStore} implementation inline. */
export function defineMessageStore(store: MessageStore): MessageStore {
  return store
}
/** Type an {@link ActivityStore} implementation inline. */
export function defineActivityStore(store: ActivityStore): ActivityStore {
  return store
}
/** Type an {@link InterruptStore} implementation inline. */
export function defineInterruptStore(store: InterruptStore): InterruptStore {
  return store
}
/** Type a {@link MetadataStore} implementation inline. */
export function defineMetadataStore(store: MetadataStore): MetadataStore {
  return store
}
/** Type a {@link GenerationRunStore} implementation inline. */
export function defineGenerationRunStore(
  store: GenerationRunStore,
): GenerationRunStore {
  return store
}
/** Type an {@link ArtifactStore} implementation inline. */
export function defineArtifactStore(store: ArtifactStore): ArtifactStore {
  return store
}
/** Type a {@link BlobStore} implementation inline. */
export function defineBlobStore(store: BlobStore): BlobStore {
  return store
}

/**
 * Metadata row describing a persisted artifact (generated media, tool output).
 *
 * The bytes themselves live in a {@link BlobStore}; this record holds the
 * descriptive metadata and an optional `sourceUrl` for reference-only
 * backends.
 *
 * @property createdAt - Epoch ms. (Core's wire-facing `PersistedArtifactRef`
 *   exposes the same instant as an ISO string; see the timestamp convention.)
 */
export interface ArtifactRecord {
  artifactId: string
  runId: string
  threadId: string
  /**
   * The blob-store key these bytes actually live under. `storageKey` can put
   * bytes anywhere, so a reader cannot recompute the path; it has to be
   * remembered.
   */
  blobKey: string
  name: string
  mimeType: string
  size: number
  sourceUrl?: string
  createdAt: number
}

/** Durable store for artifact metadata records. */
export interface ArtifactStore {
  /** Insert or overwrite the artifact metadata record. */
  save: (record: ArtifactRecord) => Promise<void>
  /** Return the artifact for `artifactId`, or `null` if none exists. */
  get: (artifactId: string) => Promise<ArtifactRecord | null>
  /**
   * All artifacts for a run in deterministic snapshot order: `createdAt`
   * ascending, then `artifactId` ascending by the unsigned UTF-8 bytes of
   * each string (compare bytes left-to-right; shorter equal prefixes first).
   * Returns `[]` when the run has none.
   */
  list: (runId: string) => Promise<Array<ArtifactRecord>>
  /**
   * All artifacts for a thread in deterministic snapshot order.
   * Records are ordered by `createdAt` ascending, then by `artifactId` using
   * the unsigned UTF-8 bytes of each string (compare bytes left-to-right; shorter
   * equal prefixes first).
   */
  listForThread: (threadId: string) => Promise<Array<ArtifactRecord>>
  /**
   * Delete a single artifact by id. A no-op if absent, mirroring
   * {@link BlobStore.delete} — the two are written and deleted as a pair, so
   * their contracts match.
   */
  delete: (artifactId: string) => Promise<void>
  /**
   * Delete every artifact belonging to `runId`. A no-op when the run has none.
   *
   * Required rather than feature-detected: retention and erasure are the point
   * of storing media durably, and an adapter silently lacking deletion is
   * indistinguishable from one where there was nothing to delete.
   */
  deleteForRun: (runId: string) => Promise<void>
}

/**
 * Accepted body shapes for {@link BlobStore.put}. `ArrayBufferView` already
 * covers `Uint8Array` and every other typed-array/`DataView`, so no separate
 * `Uint8Array` member is needed.
 */
export type BlobBody =
  | ReadableStream<Uint8Array>
  | ArrayBuffer
  | ArrayBufferView
  | string
  | Blob

/**
 * Metadata for a stored blob.
 *
 * @property size - Byte length, when known.
 * @property createdAt - Epoch ms first written.
 * @property updatedAt - Epoch ms last overwritten.
 */
export interface BlobRecord {
  key: string
  size?: number
  etag?: string
  contentType?: string
  customMetadata?: Record<string, string>
  createdAt?: number
  updatedAt?: number
}

/**
 * A byte range to read, in the shape an HTTP `Range` header resolves to.
 *
 * `offset` is measured from the start of the object and must be inside it;
 * `length` defaults to "everything from `offset` to the end" and is clamped to
 * the end when it overshoots. Suffix ranges (`bytes=-500`) are the caller's to
 * resolve against the known size — a serve route has the size on the artifact
 * record, and has to compare against it anyway to answer `416` before reading.
 */
export interface BlobRange {
  offset: number
  length?: number
}

/** Options for {@link BlobStore.get}. */
export interface BlobGetOptions {
  /**
   * Read only this slice of the object. `body`, `arrayBuffer()` and `text()`
   * then cover the slice, `size` still reports the WHOLE object, and `range`
   * reports the slice actually served — the three numbers a `206` response
   * needs (`Content-Range: bytes <offset>-<offset+length-1>/<size>`).
   */
  range?: BlobRange
}

/** A stored blob's metadata plus lazy accessors for its bytes. */
export interface BlobObject extends BlobRecord {
  arrayBuffer: () => Promise<ArrayBuffer>
  text: () => Promise<string>
  body?: ReadableStream<Uint8Array>
  /**
   * The slice this object exposes, when a {@link BlobGetOptions.range} was
   * requested and honoured: `offset` as asked, `length` as actually served
   * (clamped to the end of the object). Absent on a whole-object read.
   */
  range?: { offset: number; length: number }
}

/**
 * One page of a {@link BlobStore.list} scan.
 *
 * @property cursor - Opaque continuation token; present only when `truncated`.
 * @property truncated - `true` when more objects match beyond this page.
 */
export interface BlobListPage {
  objects: Array<BlobRecord>
  cursor?: string
  truncated?: boolean
}

export interface BlobPutOptions {
  contentType?: string
  customMetadata?: Record<string, string>
  /**
   * The exact byte length of `body`, when the producer knows it up front.
   *
   * Advisory, not a contract the store must honor: it exists so a store can
   * pick an upload strategy knowingly instead of discovering the length by
   * buffering. Most useful to an SDK that wants the length as a separate
   * argument rather than reading it off the stream — S3's `PutObject`
   * (`ContentLength`) is the archetype — and to a runtime that can re-attach
   * one (workerd's `FixedLengthStream` ahead of `R2Bucket.put`).
   *
   * Only ever set when the length is exact — a wrong value is worse than none,
   * since runtimes that enforce declared lengths fail the write. Absent means
   * unknown, and a store must accept a length-less stream regardless:
   * producers hand one over whenever the origin does not declare a length.
   */
  expectedLength?: number
}

export interface BlobListOptions {
  prefix?: string
  cursor?: string
  limit?: number
}

/** Durable object/blob store (byte-storing or reference-only backends). */
export interface BlobStore {
  /** Insert or overwrite the object at `key`, returning its metadata. */
  put: (
    key: string,
    body: BlobBody,
    options?: BlobPutOptions,
  ) => Promise<BlobRecord>
  /**
   * Return the object at `key` (metadata + byte accessors), or `null`.
   *
   * RANGE SEMANTICS: with `options.range`, return only that slice — the bytes
   * a `206` response carries — and report it back as `range`. `size` still
   * reports the whole object, so the caller can build `Content-Range` without
   * a second `head`. The reported `length` is what was actually served: a
   * requested `length` past the end clamps. An `offset` at or past the end is
   * a caller error, not a store one — the size is on the artifact record, so a
   * serve route answers `416` before ever asking the store.
   *
   * Range support is part of the contract for any store that holds bytes (the
   * conformance testkit asserts it): serving a whole file where a slice was
   * asked for is what makes `<video>` seeking, and Safari playback at all,
   * fail. A reference-only backend that stores no bytes skips `blobs`
   * entirely rather than half-implementing it.
   */
  get: (key: string, options?: BlobGetOptions) => Promise<BlobObject | null>
  /** Return only the metadata for `key`, or `null`. */
  head: (key: string) => Promise<BlobRecord | null>
  /** Remove the object at `key`. A no-op if absent. */
  delete: (key: string) => Promise<void>
  /**
   * List objects, optionally filtered by `prefix`, in ascending key order.
   *
   * CURSOR SEMANTICS: `prefix` matches literally and case-sensitively (SQL
   * backends must escape LIKE metacharacters, so `run_` matches only the exact
   * bytes `run_`, not `_` as a wildcard). When `limit` is given and more keys
   * match, the page is `truncated: true` with a `cursor`; passing that `cursor`
   * back returns the strictly-following keys (keys `> cursor`). Cursor ordering
   * is the same byte ordering as the sort, so paging visits every key exactly
   * once with no gaps or repeats. `limit: 0` yields an empty, untruncated page
   * with no cursor.
   */
  list: (options?: BlobListOptions) => Promise<BlobListPage>
}

// ===========================================================================
// Session stores
// ===========================================================================
//
// Opt-in stores for a host that keeps long-lived sessions: an append-only
// log, and a list of sessions. `withPersistence` writes only `sessions`, an entry
// for each subagent child, in the tenant of its `tenantId` option.

/**
 * One record of a session log. `type` names the kind of record. The other
 * fields are JSON. The store does not read records: the writer owns the
 * `type` names.
 */
export interface LogRecord {
  type: string
  [key: string]: unknown
}

/** A {@link LogRecord} and its position in the log. */
export interface LogEntry {
  /** The position of the record. The first record of a log is at 1. */
  seq: number
  record: LogRecord
}

/**
 * The epoch of a log that no host ever acquired. A writer that never calls
 * {@link LogStore.acquireOwnership} writes with this epoch.
 */
export const NO_OWNER_EPOCH = 0

/**
 * The error {@link LogStore.append} rejects with when `seq` is not the next
 * free position of the log. It means that another writer appended first.
 * Nothing of the rejected batch is written.
 */
export class LogConflictError extends Error {
  readonly logId: string
  readonly seq: number

  constructor(logId: string, seq: number) {
    super(
      `Log conflict on log ${JSON.stringify(logId)}: position ${seq} is not the next free position.`,
    )
    this.name = 'LogConflictError'
    this.logId = logId
    this.seq = seq
  }
}

/**
 * The error that a fenced write ({@link LogStore.append},
 * {@link LogStore.putSnapshot}, {@link LogStore.truncate},
 * {@link LogStore.pinToWorkers}) rejects with when `epoch` does not give the
 * writer the log. It means that another host acquired the log, that the
 * owner released it, that the log was deleted, or that the writer did not
 * acquire a log that needs an owner. Nothing is written. Stop the work on the
 * log.
 */
export class OwnershipLostError extends Error {
  readonly logId: string
  readonly epoch: number

  constructor(logId: string, epoch: number) {
    super(
      `Ownership lost on log ${JSON.stringify(logId)}: epoch ${epoch} does not own the log.`,
    )
    this.name = 'OwnershipLostError'
    this.logId = logId
    this.epoch = epoch
  }
}

/**
 * The error {@link LogStore.read} rejects with when the query starts before
 * the first record that the log keeps. A truncate dropped the records before
 * `firstSeq`. Load the snapshot, then read after its `seq`.
 */
export class LogTruncatedError extends Error {
  readonly logId: string
  readonly firstSeq: number

  constructor(logId: string, firstSeq: number) {
    super(
      `Log ${JSON.stringify(logId)} is truncated: the first kept position is ${firstSeq}.`,
    )
    this.name = 'LogTruncatedError'
    this.logId = logId
    this.firstSeq = firstSeq
  }
}

/**
 * The result of {@link LogStore.offerInput}:
 *
 * - `'offered'`: the store keeps the input as a pending offer.
 * - `'duplicate'`: the id was offered before with an equal record.
 * - `'conflict'`: the id was offered before with a different record.
 * - `'deleted'`: `deleteLog` deleted the log, and no host acquired it since.
 *   Nothing is stored.
 */
export type OfferResult = 'offered' | 'duplicate' | 'conflict' | 'deleted'

/** A pending input of a log, from {@link LogStore.listOfferedInputs}. */
export interface OfferedInput {
  inputId: string
  record: LogRecord
  /** Epoch milliseconds, from the store clock. */
  offeredAt: number
}

/** A log that has work and no live owner, from {@link LogStore.listLogsWithoutOwner}. */
export interface OrphanedLog {
  logId: string
  /** The name of the setup that runs the log, so a worker can open it. */
  harness: string
  /**
   * Why the log has work:
   *
   * - `'expired'`: the ownership expired and the owner did not release it.
   * - `'offered'`: the log has pending offers and no live owner.
   * - `'wake'`: the owner released the log with a wake time that is due.
   */
  reason: 'expired' | 'offered' | 'wake'
}

/** The saved state of a log, from {@link LogStore.getSnapshot}. */
export interface LogSnapshot {
  /** The position of the last record that `state` includes. */
  seq: number
  /** The version of the `state` shape. The writer owns the version names. */
  version: string
  /** JSON state. The store does not read it. */
  state: unknown
}

/**
 * Durable append-only log of a session tree. A host keeps the events, the
 * transcript, the inputs, and the tool steps of the tree in it. The store also
 * keeps one ownership record per log, the pending inputs (the input lane), and
 * the latest snapshot.
 *
 * The log has one writer at a time: the owner. A host gets the log with
 * `acquireOwnership`, and each fenced write carries the epoch that `acquireOwnership` returned.
 * When another host acquired the log, the epoch is old and the write fails.
 * `append` is also a compare-and-append: the writer names the position of the
 * first record, and the store refuses the batch when that position is taken.
 *
 * The store clock decides expiry. Callers pass `ttlMs` and `wakeInMs`, never
 * timestamps. Records are JSON: a store keeps copies, never the objects of
 * the caller.
 */
export interface LogStore {
  /**
   * Write `records` at positions `seq`, `seq + 1`, and so on, and consume the
   * pending offers in `consumes` in the same write.
   *
   * The store checks in this order, and the first check that fails rejects:
   *
   * 1. The fence. On a log that a host acquired, `epoch` must be the current
   *    epoch, and the owner must not have released the log. On a log that no
   *    host acquired, `epoch` must be {@link NO_OWNER_EPOCH}. A deleted log
   *    refuses every epoch until the next `acquireOwnership`. Else reject with
   *    {@link OwnershipLostError}.
   * 2. `seq` must be the next free position (the last position plus 1, or 1
   *    for a new log). Else reject with {@link LogConflictError}.
   * 3. Each id in `consumes` must be a pending offer. Else reject with a
   *    `TypeError`.
   * 4. Each record must become JSON. Else reject with the error of the copy.
   *
   * INVARIANT (atomic): the store writes the whole batch and consumes all of
   * `consumes`, or it changes nothing. A crash never leaves a part of a batch.
   *
   * An empty batch with no `consumes` resolves, writes nothing, and checks
   * nothing.
   */
  append: (batch: {
    logId: string
    seq: number
    epoch: number
    records: ReadonlyArray<LogRecord>
    consumes?: ReadonlyArray<string>
  }) => Promise<void>
  /**
   * The entries of `logId` with `seq` greater than `after` (default 0), in
   * ascending order, at most `limit` of them. `limit: 0` gives `[]`.
   *
   * If a truncate dropped the record at position `after + 1`, reject with
   * {@link LogTruncatedError}.
   *
   * INVARIANT: a log that has no records gives `[]`. The returned records are
   * copies: a change to them does not change the log.
   */
  read: (query: {
    logId: string
    after?: number
    limit?: number
  }) => Promise<Array<LogEntry>>
  /**
   * Call `listener` after each successful `append` to `logId`, each `offerInput`
   * that returns `'offered'`, and each `deleteLog` of `logId`. The calls come
   * from every store instance on the same backend (other machines). A store
   * that cannot push must poll. A listener that throws is logged, and the
   * write still succeeds. Returns a function that stops the calls.
   */
  subscribe: (watch: { logId: string; listener: () => void }) => () => void

  /**
   * Get the log for `ownerId` for `ttlMs` milliseconds of the store clock.
   *
   * - Returns `null` when another owner holds the log (not released, and not
   *   expired).
   * - Returns `null` when the log has an affinity that is not
   *   `claim.affinity`, and the host of that affinity is alive (see
   *   `markHostAlive`). The pin of a host with no live heartbeat does not count.
   * - When `ownerId` already holds the log, extends the expiry and returns
   *   the same epoch.
   * - Else increments the epoch, stores `harness`, clears the wake time,
   *   opens a deleted log again, and returns the new epoch. The epoch never
   *   goes down, also after `releaseOwnership` and `deleteLog`.
   * - When the log had another owner whose ownership expired with no
   *   `releaseOwnership`, the result has `expired: true`: that owner stopped.
   *
   * Rejects with a `RangeError` when `ttlMs` is not more than 0.
   */
  acquireOwnership: (claim: {
    logId: string
    harness: string
    ownerId: string
    ttlMs: number
    affinity?: string
  }) => Promise<{ epoch: number; expired?: boolean } | null>
  /**
   * Extend the ownership to `ttlMs` milliseconds from now. Returns `true` only
   * when the log is not released, `epoch` is the current epoch, and `ownerId`
   * is the owner. A renew after the expiry succeeds when no other host
   * acquired the log in between.
   *
   * Rejects with a `RangeError` when `ttlMs` is not more than 0.
   */
  renewOwnership: (claim: {
    logId: string
    ownerId: string
    epoch: number
    ttlMs: number
  }) => Promise<boolean>
  /**
   * Give up the log. With a matching `epoch` and `ownerId`, mark the log
   * released, and set a wake time `wakeInMs` from now when `wakeInMs` is
   * given. Any other call does nothing.
   *
   * Rejects with a `RangeError` when `wakeInMs` is less than 0.
   */
  releaseOwnership: (claim: {
    logId: string
    ownerId: string
    epoch: number
    wakeInMs?: number
  }) => Promise<void>
  /**
   * Pin the log to the workers with `affinity`, or unpin it with `null`.
   * Only the owner pins: `epoch` must be the current epoch of an acquired
   * log that is not released. Else reject with {@link OwnershipLostError},
   * also for {@link NO_OWNER_EPOCH}.
   */
  pinToWorkers: (pin: {
    logId: string
    epoch: number
    affinity: string | null
  }) => Promise<void>
  /**
   * Mark the host `hostId` alive for `ttlMs` milliseconds of the store
   * clock. `hostId` is the affinity of its pins. A pin counts only while its
   * host is alive, so the logs of a host that stopped go to other hosts.
   *
   * Rejects with a `RangeError` when `ttlMs` is not more than 0.
   */
  markHostAlive: (beat: { hostId: string; ttlMs: number }) => Promise<void>
  /**
   * The logs that have work and no live owner, at most `limit` of them. A log
   * is listed for the first reason that applies:
   *
   * 1. `'expired'`: not released, and the expiry is due.
   * 2. `'offered'`: pending offers, and no live owner.
   * 3. `'wake'`: released, and the wake time is due.
   *
   * Only logs of `harnesses` are listed, when `harnesses` is given. A log with
   * an affinity is listed only when `affinity` is equal, or the host of the
   * affinity is not alive (see `markHostAlive`). The log that is due
   * first comes first. `limit: 0` gives `[]`.
   */
  listLogsWithoutOwner: (query: {
    limit: number
    harnesses?: ReadonlyArray<string>
    affinity?: string
  }) => Promise<Array<OrphanedLog>>

  /**
   * Keep `record` as a pending input of the log. Any machine can offer. Only
   * the owner consumes, with `append({ consumes })`.
   *
   * - A deleted log returns `'deleted'` and stores nothing, until the next
   *   `acquireOwnership`.
   * - An `inputId` that was offered before (pending or consumed) returns
   *   `'duplicate'` when the JSON of the record is equal, else `'conflict'`.
   *   Nothing changes.
   * - A new `inputId` is stored with the store clock as `offeredAt`, calls
   *   the listeners, and returns `'offered'`. `harness` names the harness of
   *   the log only when the log has no owner and no earlier offer: the
   *   harness of the owner, else of the first offer, stays.
   *
   * Rejects, and stores nothing, when the record cannot become JSON.
   */
  offerInput: (input: {
    logId: string
    harness: string
    inputId: string
    record: LogRecord
  }) => Promise<OfferResult>
  /** The pending offers of `logId`, in offer order. */
  listOfferedInputs: (query: { logId: string }) => Promise<Array<OfferedInput>>

  /** The latest snapshot of `logId`, or `null` when the log has none. */
  getSnapshot: (query: { logId: string }) => Promise<LogSnapshot | null>
  /**
   * Replace the snapshot of the log. Fenced like `pinToWorkers`: only the
   * owner of an acquired log writes it. Rejects with a `RangeError`
   * unless `firstSeq - 1 <= seq <= head`. `head` is the last position of the
   * log. `firstSeq` is the first kept position after a `truncate`, or 1. A
   * snapshot before `firstSeq - 1` needs records that a truncate dropped.
   */
  putSnapshot: (
    snapshot: { logId: string; epoch: number } & LogSnapshot,
  ) => Promise<void>
  /**
   * Drop the records with `seq` less than `beforeSeq`. The positions keep
   * counting, so the next `append` is still at the last position plus 1.
   * Fenced like `pinToWorkers`: only the owner of an acquired log truncates.
   * Rejects with a `RangeError` unless a snapshot
   * has a `seq` of `beforeSeq - 1` or more.
   */
  truncate: (cut: {
    logId: string
    epoch: number
    beforeSeq: number
  }) => Promise<void>
  /**
   * Remove the records, the offers, the snapshot, the truncation point, the
   * affinity, and the wake time of `logId`, and call the listeners. The epoch
   * counter stays. Every fenced write to the log fails until the next
   * `acquireOwnership`, also with {@link NO_OWNER_EPOCH}, so a late write cannot bring
   * the deleted records back. An `offerInput` returns `'deleted'`.
   */
  deleteLog: (target: { logId: string }) => Promise<void>
}

/** Type a {@link LogStore} implementation inline. */
export function defineLogStore(store: LogStore): LogStore {
  return store
}

/**
 * A user who can open a session, and the role of that user in it. The
 * `owner` is the user who opened the session. The host decides what a
 * `member`, a `viewer`, and an `operator` (a person who takes over from an
 * agent) can do.
 */
export interface SessionMember {
  userId: string
  role: 'owner' | 'member' | 'viewer' | 'operator'
}

/**
 * The key of one session index entry. Without `tenantId`, the key is in the
 * default tenant.
 */
export interface SessionIndexKey {
  tenantId?: string
  threadId: string
}

/**
 * One entry of the session index: what a list of sessions shows for one
 * thread. The index holds no messages. The thread data stays in the stores
 * that hold it.
 *
 * @property createdAt - Epoch ms when the thread was first opened.
 * @property updatedAt - Epoch ms of the last change. {@link SessionIndexStore.list}
 *   sorts by it.
 */
export interface SessionIndexEntry<TMetadata = Record<string, unknown>> {
  /**
   * The tenant of the thread. Without it, the entry is in the default tenant.
   * The tenant and `threadId` together are the key, so one thread id in two
   * tenants is two entries.
   */
  tenantId?: string
  threadId: string
  /**
   * The users who can open the thread. The owner is the member with the role
   * `owner`. Empty when the thread has no owner: a host without access rules
   * opened it without a user. A subagent child keeps it empty: access to a
   * child goes through its top-level thread.
   */
  members: Array<SessionMember>
  /** The name of the setup that runs the thread. */
  harness?: string
  title?: string
  /**
   * The thread that started this one. A subagent child has the thread id
   * `subagent:<runId>` and the thread of its parent here.
   */
  parentThreadId?: string
  /** The tool call in the parent thread that started this thread. */
  parentToolCallId?: string
  createdAt: number
  updatedAt: number
  pinned?: boolean
  /**
   * `true` when the thread is archived. {@link SessionIndexStore.list} leaves
   * archived entries out unless `withArchived` is `true`. The thread and its
   * data stay, and it still opens.
   */
  archived?: boolean
  /** Token totals of the thread. */
  usage?: {
    /** Model calls. */
    turns: number
    promptTokens: number
    completionTokens: number
    totalTokens: number
    /** Input tokens the provider read from its prompt cache. */
    cachedTokens: number
    /** Input tokens the provider wrote to its prompt cache. */
    cacheWriteTokens: number
    /** Cost in USD, when the host counts it. */
    cost?: number
  }
  /** Data the app owns, for example the context items a user pinned. */
  metadata?: TMetadata
}

/** Options for {@link SessionIndexStore.list}. */
export interface SessionIndexListOptions {
  /**
   * The tenant to list. Without it, the list has only the default tenant. A
   * list never has the entries of another tenant.
   */
  tenantId?: string
  /**
   * Only the entries whose `members` have this user id, with any role. An
   * entry with no `members` (a subagent child) never matches.
   */
  memberId?: string
  /** The most entries in one page. Without it, the page has all entries. */
  limit?: number
  /** The `cursor` of the previous page. */
  cursor?: string
  /**
   * Only the entries whose `parentThreadId` is this thread. `null`: only the
   * entries with no parent (the top-level sessions).
   */
  parentThreadId?: string | null
  /**
   * Only the entries whose `title` contains this text, without case. An entry
   * without a title never matches.
   */
  search?: string
  /** Only the entries whose `harness` is exactly this name. */
  harness?: string
  /**
   * Only the entries whose `metadata` has each key with exactly this string
   * value.
   */
  metadata?: Record<string, string>
  /**
   * `true`: the list has the archived entries too. Default: it leaves out the
   * entries with `archived: true`.
   */
  withArchived?: boolean
}

/**
 * One page of a {@link SessionIndexStore.list} scan.
 *
 * @property cursor - The token for the next page. Only when `truncated`.
 * @property truncated - `true` when more entries match after this page.
 */
export interface SessionIndexPage {
  entries: Array<SessionIndexEntry>
  cursor?: string
  truncated?: boolean
}

/**
 * Durable index of sessions, one entry per tenant and thread. The key is
 * `(tenantId ?? '', threadId)`: a key without `tenantId` is in the default
 * tenant, and one thread id in two tenants is two entries.
 */
export interface SessionIndexStore {
  /** The entry for `key`, or `undefined`. */
  get: (key: SessionIndexKey) => Promise<SessionIndexEntry | undefined>
  /**
   * The entry for `key` with its revision, or `undefined` when it is absent.
   * A write that changes the entry gives it a new revision. A revision may
   * name the content (a content-hash ETag, as in S3 or Netlify Blobs), so the
   * same entry written again, also after a `delete`, can get the same
   * revision back.
   */
  getVersioned: (
    key: SessionIndexKey,
  ) => Promise<{ entry: SessionIndexEntry; revision: string } | undefined>
  /**
   * Write the entry for the tenant and thread of `entry`.
   *
   * INVARIANT (full replace): a second upsert for the same key replaces the
   * whole entry. A field that the new entry does not have is gone. To change
   * one field, read the entry and write a changed copy with `upsertIf`.
   */
  upsert: (entry: SessionIndexEntry) => Promise<void>
  /**
   * Write `entry` only when the current revision of its key is `revision`
   * (`null` means "only when absent"). Returns the new revision, or
   * `{ ok: false, reason: 'conflict' }` when someone else wrote first. With
   * `getVersioned`, hosts on several machines do not lose each other's
   * changes to one entry. A full replace, the same as `upsert`.
   */
  upsertIf: (input: {
    entry: SessionIndexEntry
    revision: string | null
  }) => Promise<
    { ok: true; revision: string } | { ok: false; reason: 'conflict' }
  >
  /**
   * The entries of one tenant that match the filters, newest `updatedAt`
   * first. Entries with the same `updatedAt` are in `threadId` order. Pinned
   * entries get no special place. All filters must match. Archived entries
   * are left out unless `withArchived` is `true`.
   *
   * CURSOR SEMANTICS: when `limit` is given and more entries match, the page
   * is `truncated: true` with a `cursor`. Pass that `cursor` back, with the
   * same filters, to get the entries that follow the last entry of the page.
   * An entry that changes between two pages can move in the order.
   */
  list: (options?: SessionIndexListOptions) => Promise<SessionIndexPage>
  /**
   * Remove the entry for `key`. A hard delete: after it, `get` and
   * `getVersioned` return `undefined`, and `upsertIf` with revision `null`
   * writes a new entry. A no-op for an unknown key.
   *
   * This removes the entry only. The messages, the log, and the other data of
   * the thread stay. The stores that hold that data remove it.
   */
  delete: (key: SessionIndexKey) => Promise<void>
}

/** Type a {@link SessionIndexStore} implementation inline. */
export function defineSessionIndexStore(
  store: SessionIndexStore,
): SessionIndexStore {
  return store
}

/**
 * Sparse bag of **state** store keys — composition / validation only.
 *
 * **Not a public product shape.** Prefer the named chat shapes below
 * ({@link ChatTranscriptStores}, {@link ChatPersistenceStores},
 * {@link ChatWithInterruptsStores}). Locks are not included — use
 * `withLocks` from `@tanstack/ai`.
 *
 * @internal Exported from this module for generics; the package root does not
 * re-export this type — use a named shape or `AIPersistence<{ … }>` instead.
 */
export interface AIPersistenceStores {
  messages?: MessageStore
  activities?: ActivityStore
  runs?: RunStore
  interrupts?: InterruptStore
  metadata?: MetadataStore
  generationRuns?: GenerationRunStore
  artifacts?: ArtifactStore
  blobs?: BlobStore
  /** Append-only session log. */
  log?: LogStore
  /** The session index. */
  sessions?: SessionIndexStore
}

/**
 * Chat floor: durable transcript. `messages` is required.
 *
 * `runs` / `interrupts` / `metadata` remain optional. If `interrupts` is set,
 * `runs` is required (enforced by `withPersistence` / validators).
 */
export interface ChatTranscriptStores {
  messages: MessageStore
  activities?: ActivityStore
  runs?: RunStore
  interrupts?: InterruptStore
  metadata?: MetadataStore
}

/**
 * Full chat durability — all four state stores are present. This is what
 * `memoryPersistence()` returns, and the shape most adapters should declare.
 *
 * Backends that only need a transcript should use
 * {@link ChatTranscriptStores} instead.
 */
export interface ChatPersistenceStores {
  messages: MessageStore
  activities?: ActivityStore
  runs: RunStore
  interrupts: InterruptStore
  metadata: MetadataStore
}

/**
 * Chat with durable human-in-the-loop interrupts (and optional metadata).
 * Implies `runs` (interrupt records are run-scoped).
 *
 * Prefer {@link ChatPersistenceStores} when you also have metadata (packaged
 * backends). Use this when interrupts are required but metadata is not.
 */
export interface ChatWithInterruptsStores {
  messages: MessageStore
  runs: RunStore
  interrupts: InterruptStore
  metadata?: MetadataStore
}

/**
 * Persistence aggregate. Parameterize with a named store shape, or a sparse
 * map for composition (`defineAIPersistence` / `composePersistence`).
 *
 * Default is the sparse bag so untyped / dynamic bags still type-check;
 * prefer {@link ChatTranscriptPersistence} or {@link ChatPersistence} at
 * call sites.
 */
export interface AIPersistence<
  TStores extends AIPersistenceStores = AIPersistenceStores,
> {
  stores: ExactStoreKeys<TStores>
}

/** {@link AIPersistence} for {@link ChatTranscriptStores}. */
export type ChatTranscriptPersistence = AIPersistence<ChatTranscriptStores>

/** {@link AIPersistence} for {@link ChatPersistenceStores}. */
export type ChatPersistence = AIPersistence<ChatPersistenceStores>

/** {@link AIPersistence} for {@link ChatWithInterruptsStores}. */
export type ChatWithInterruptsPersistence =
  AIPersistence<ChatWithInterruptsStores>

type StoreKey = keyof AIPersistenceStores
type ExactStoreKeys<TStores> =
  Exclude<keyof TStores, StoreKey> extends never
    ? TStores
    : TStores & Record<Exclude<keyof TStores, StoreKey>, never>

export type AIPersistenceOverrides = {
  [TKey in StoreKey]?: AIPersistenceStores[TKey] | false
}

type BaseStoreValue<
  TBase extends AIPersistenceStores,
  TKey extends StoreKey,
> = TKey extends keyof TBase ? TBase[TKey] : never

type OverrideStoreValue<
  TOverrides extends AIPersistenceOverrides,
  TKey extends StoreKey,
> = TKey extends keyof TOverrides ? TOverrides[TKey] : never

type ResolvedStoreValue<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
  TKey extends StoreKey,
> = TKey extends keyof TOverrides
  ?
      | Exclude<OverrideStoreValue<TOverrides, TKey>, false | undefined>
      | (undefined extends OverrideStoreValue<TOverrides, TKey>
          ? Exclude<BaseStoreValue<TBase, TKey>, undefined>
          : never)
  : Exclude<BaseStoreValue<TBase, TKey>, undefined>

type BaseStoreIsRequired<
  TBase extends AIPersistenceStores,
  TKey extends StoreKey,
> = TKey extends keyof TBase
  ? object extends Pick<TBase, TKey>
    ? false
    : true
  : false

type ResolvedStoreIsRequired<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
  TKey extends StoreKey,
> = TKey extends keyof TOverrides
  ? false extends OverrideStoreValue<TOverrides, TKey>
    ? false
    : undefined extends OverrideStoreValue<TOverrides, TKey>
      ? BaseStoreIsRequired<TBase, TKey>
      : true
  : BaseStoreIsRequired<TBase, TKey>

type ResolvedRequiredKeys<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
> = {
  [TKey in StoreKey]-?: [ResolvedStoreValue<TBase, TOverrides, TKey>] extends [
    never,
  ]
    ? never
    : ResolvedStoreIsRequired<TBase, TOverrides, TKey> extends true
      ? TKey
      : never
}[StoreKey]

type ResolvedOptionalKeys<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
> = {
  [TKey in StoreKey]-?: [ResolvedStoreValue<TBase, TOverrides, TKey>] extends [
    never,
  ]
    ? never
    : ResolvedStoreIsRequired<TBase, TOverrides, TKey> extends true
      ? never
      : TKey
}[StoreKey]

type Simplify<T> = { [TKey in keyof T]: T[TKey] }

export type ComposedAIPersistenceStores<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
> = Simplify<
  {
    [TKey in ResolvedRequiredKeys<TBase, TOverrides>]: ResolvedStoreValue<
      TBase,
      TOverrides,
      TKey
    >
  } & {
    [TKey in ResolvedOptionalKeys<TBase, TOverrides>]?: ResolvedStoreValue<
      TBase,
      TOverrides,
      TKey
    >
  }
>

const storeKeys = [
  'messages',
  'activities',
  'runs',
  'generationRuns',
  'interrupts',
  'metadata',
  'artifacts',
  'blobs',
  'log',
  'sessions',
] satisfies Array<StoreKey>

const storeKeySet = new Set<string>(storeKeys)

function assertKnownStoreKeys(stores: object, location: string): void {
  for (const key of Object.keys(stores)) {
    if (!storeKeySet.has(key)) {
      throw new Error(`Unknown AIPersistence ${location} key: ${key}`)
    }
  }
}

export function validatePersistenceStoreKeys(persistence: AIPersistence): void {
  assertKnownStoreKeys(persistence.stores, 'store')
}

/**
 * Chat middleware entrypoint rules:
 * - `messages` is required (chat persistence means a durable transcript)
 * - `interrupts` requires `runs` (interrupt records are run-scoped)
 */
export function validateChatPersistenceStores(
  persistence: AIPersistence,
): void {
  validatePersistenceStoreKeys(persistence)
  if (!persistence.stores.messages) {
    throw new Error('Chat persistence requires stores.messages.')
  }
  if (persistence.stores.interrupts && !persistence.stores.runs) {
    throw new Error('Chat persistence stores.interrupts requires stores.runs.')
  }
}

/**
 * Generation middleware entrypoint rule: `generationRuns` is required (the
 * generation run lifecycle is keyed on its own `runId`, not a chat conversation
 * `threadId`). When artifact persistence is used, `artifacts` and `blobs` must
 * be provided together.
 */
export function validateGenerationPersistenceStores(
  persistence: AIPersistence,
): void {
  validatePersistenceStoreKeys(persistence)
  const hasArtifacts = persistence.stores.artifacts !== undefined
  const hasBlobs = persistence.stores.blobs !== undefined
  if (hasArtifacts !== hasBlobs) {
    throw new Error(
      'Generation artifact persistence requires both stores.artifacts and stores.blobs.',
    )
  }
  if (!persistence.stores.generationRuns) {
    throw new Error('Generation persistence requires stores.generationRuns.')
  }
}

/**
 * Server hydrate entrypoint rule: `messages` is required.
 */
export function validateReconstructChatStores(
  persistence: AIPersistence,
): void {
  validatePersistenceStoreKeys(persistence)
  if (!persistence.stores.messages) {
    throw new Error('reconstructChat requires stores.messages.')
  }
}

/**
 * Server hydrate entrypoint rule for generation: `generationRuns` is required.
 * The run store resolves the latest generation for a thread (or a specific run
 * id), so a server-authoritative client can hydrate the last generation's
 * status, result, and artifact refs on load.
 */
export function validateReconstructGenerationStores(
  persistence: AIPersistence,
): void {
  validatePersistenceStoreKeys(persistence)
  if (!persistence.stores.generationRuns) {
    throw new Error('reconstructGeneration requires stores.generationRuns.')
  }
}

export function defineAIPersistence<TStores extends AIPersistenceStores>(
  persistence: AIPersistence<ExactStoreKeys<TStores>>,
): AIPersistence<TStores> {
  validatePersistenceStoreKeys(persistence)
  return persistence
}

export function composePersistence<
  TBase extends AIPersistenceStores,
  TOverrides extends AIPersistenceOverrides,
>(
  base: AIPersistence<TBase>,
  config: {
    overrides: ExactStoreKeys<TOverrides>
  },
): AIPersistence<ComposedAIPersistenceStores<TBase, TOverrides>>
export function composePersistence(
  base: AIPersistence,
  config: { overrides: AIPersistenceOverrides },
): AIPersistence {
  validatePersistenceStoreKeys(base)
  assertKnownStoreKeys(config.overrides, 'override')

  const stores: AIPersistenceStores = { ...base.stores }
  for (const key of storeKeys) {
    if (!Object.prototype.hasOwnProperty.call(config.overrides, key)) continue
    const override = config.overrides[key]
    if (override === false) {
      delete stores[key]
    } else if (override !== undefined) {
      setStore(stores, key, override)
    }
  }
  return { stores }
}

function setStore<TKey extends StoreKey>(
  stores: AIPersistenceStores,
  key: TKey,
  value: NonNullable<AIPersistenceStores[TKey]>,
): void {
  stores[key] = value
}