// BlobStore — the byte layer. R2 backs it.
interface BlobStore {
  put(
    key: string,
    body: BlobBody,
    options?: BlobPutOptions,
  ): Promise<BlobRecord>
  // metadata + byte accessors; `options.range` reads one slice (for `206`s)
  get(key: string, options?: BlobGetOptions): Promise<BlobObject | null>
  head(key: string): Promise<BlobRecord | null> // metadata only
  delete(key: string): Promise<void> // no-op if absent
  list(options?: BlobListOptions): Promise<BlobListPage>
}

// ArtifactStore — the metadata layer. D1 (or KV) backs it.
interface ArtifactStore {
  save(record: ArtifactRecord): Promise<void> // insert or overwrite
  get(artifactId: string): Promise<ArtifactRecord | null>
  list(runId: string): Promise<Array<ArtifactRecord>> // [] when none
  delete(artifactId: string): Promise<void>
  deleteForRun(runId: string): Promise<void>
}

import { defineBlobStore, resolveBlobRange } from '@tanstack/ai-persistence'
import type { BlobObject, BlobRecord } from '@tanstack/ai-persistence'

// R2 multipart parts must be ≥ 5 MiB and — except for the last — all exactly
// the SAME size, so a part reader has to cut on an exact boundary and carry the
// remainder. 8 MiB × the 10,000-part ceiling puts the multipart path's limit at
// ~80 GiB; raise this for larger objects, and check R2's current object-size
// limits before promising more.
const MULTIPART_PART_SIZE = 8 * 1024 * 1024

/**
 * Cut exactly `limit` bytes off the stream (fewer only at EOF), carrying any
 * overshoot into the next part.
 *
 * `carry` is the leftover from the previous call. Chunks are collected by
 * reference and copied once per part: growing a `Uint8Array` chunk-by-chunk
 * instead would re-copy the whole part on every read — quadratic work for a
 * part built from hundreds of small chunks.
 */
async function readPart(
  reader: ReadableStreamDefaultReader<Uint8Array>,
  limit: number,
  carry: Uint8Array,
): Promise<{ bytes: Uint8Array; carry: Uint8Array; eof: boolean }> {
  const chunks: Array<Uint8Array> = carry.byteLength > 0 ? [carry] : []
  let total = carry.byteLength
  let eof = false
  while (total < limit) {
    const { value, done } = await reader.read()
    if (done) {
      eof = true
      break
    }
    chunks.push(value)
    total += value.byteLength
  }
  const joined = new Uint8Array(total)
  let offset = 0
  for (const chunk of chunks) {
    joined.set(chunk, offset)
    offset += chunk.byteLength
  }
  // Equal-sized parts: hand back exactly `limit` and keep the rest for next
  // time. At EOF the last part is whatever is left, which R2 allows.
  if (!eof && total > limit) {
    return {
      bytes: joined.subarray(0, limit),
      carry: joined.subarray(limit),
      eof,
    }
  }
  return { bytes: joined, carry: new Uint8Array(0), eof }
}

// R2 takes a single-shot put up to 5 GiB. Above that the upload has to be
// multipart even when the length is known.
const MAX_SINGLE_SHOT_BYTES = 5 * 1024 * 1024 * 1024

/**
 * `R2Bucket.put` requires a stream with a known length; a pass-through
 * TransformStream (what the middleware sends when it had to cap the body) has
 * none. Re-declare the length when `expectedLength` provides it — the
 * middleware only sets it when it is the exact decoded byte count — and
 * otherwise stream through a multipart upload, one buffered part at a time.
 */
async function putStream(
  bucket: R2Bucket,
  key: string,
  body: ReadableStream<Uint8Array>,
  options: R2PutOptions,
  expectedLength: number | undefined,
): Promise<R2Object | null> {
  if (expectedLength !== undefined && expectedLength <= MAX_SINGLE_SHOT_BYTES) {
    return bucket.put(
      key,
      body.pipeThrough(new FixedLengthStream(expectedLength)),
      options,
    )
  }
  // Unknown length, or too big for one shot: multipart, one part at a time.
  const reader = body.getReader()
  const first = await readPart(reader, MULTIPART_PART_SIZE, new Uint8Array(0))
  if (first.eof) {
    // The whole body fit in one part — a plain put is enough.
    return bucket.put(key, first.bytes, options)
  }
  const upload = await bucket.createMultipartUpload(key, options)
  try {
    const parts: Array<R2UploadedPart> = [
      await upload.uploadPart(1, first.bytes),
    ]
    let carry = first.carry
    let partNumber = 2
    for (;;) {
      const part = await readPart(reader, MULTIPART_PART_SIZE, carry)
      carry = part.carry
      if (part.bytes.byteLength > 0) {
        // 10,000 parts is R2's ceiling; failing here beats a confusing error
        // from `complete` after uploading gigabytes.
        if (partNumber > 10_000) {
          throw new Error(
            `Artifact at ${key} exceeds the multipart part limit — raise MULTIPART_PART_SIZE.`,
          )
        }
        parts.push(await upload.uploadPart(partNumber, part.bytes))
        partNumber += 1
      }
      if (part.eof) break
    }
    return await upload.complete(parts)
  } catch (error) {
    await upload.abort().catch(() => undefined)
    throw error
  }
}

function toRecord(obj: R2Object): BlobRecord {
  const uploaded = obj.uploaded.getTime()
  return {
    key: obj.key,
    size: obj.size,
    etag: obj.etag,
    ...(obj.httpMetadata?.contentType
      ? { contentType: obj.httpMetadata.contentType }
      : {}),
    ...(obj.customMetadata ? { customMetadata: obj.customMetadata } : {}),
    createdAt: uploaded,
    updatedAt: uploaded,
  }
}

export function r2BlobStore(bucket: R2Bucket) {
  return defineBlobStore({
    async put(key, body, options) {
      const r2Options = {
        ...(options?.contentType
          ? { httpMetadata: { contentType: options.contentType } }
          : {}),
        ...(options?.customMetadata
          ? { customMetadata: options.customMetadata }
          : {}),
      }
      const obj =
        body instanceof ReadableStream
          ? await putStream(
              bucket,
              key,
              body,
              r2Options,
              options?.expectedLength,
            )
          : await bucket.put(key, body, r2Options)
      // R2.put returns null only when an onlyIf precondition fails — not used here.
      if (!obj) throw new Error(`R2 put failed for ${key}`)
      return toRecord(obj)
    },

    async get(key, options): Promise<BlobObject | null> {
      if (!options?.range) {
        const whole = await bucket.get(key)
        if (!whole) return null
        return {
          ...toRecord(whole),
          body: whole.body,
          arrayBuffer: () => whole.arrayBuffer(),
          text: () => whole.text(),
        }
      }
      // A ranged read is what a serve route turns into `206` — R2 slices in
      // the bucket, so a seek never streams the whole object. `resolveBlobRange`
      // clamps a too-long length and throws on an offset past the end (the
      // route answers `416` from `record.size` before ever calling in). The
      // `head` buys that clamp deterministically — one class-B op against
      // bytes you did not need to send.
      //
      // Retry a bounded number of times: each attempt measures the object and
      // reads a slice of THAT version, and only an overwrite landing between
      // the two calls costs another lap.
      for (let attempt = 0; attempt < 3; attempt++) {
        const head = await bucket.head(key)
        // A ranged read of a key that is not there is a miss, not a
        // whole-object read: falling through to an un-ranged `get` would
        // answer a `206` carrying the entire file.
        if (!head) return null
        const served = resolveBlobRange(head.size, options.range)
        const obj = await bucket.get(key, {
          range: { offset: served.offset, length: served.length },
          // Tie the slice to the version `head` measured. Without this, a
          // `put` landing between the two calls yields a `Content-Range`
          // computed from one object and bytes from another.
          onlyIf: { etagMatches: head.etag },
        })
        // No body means the precondition failed — the object changed under us.
        if (!obj) return null
        if (!('body' in obj)) continue
        return {
          // `toRecord(obj)` reports the WHOLE object's size even on a ranged
          // read — R2's `R2Object.size` is the object, not the slice — which
          // is exactly the contract, and what `Content-Range`'s total needs.
          ...toRecord(obj),
          range: served,
          body: obj.body,
          arrayBuffer: () => obj.arrayBuffer(),
          text: () => obj.text(),
        }
      }
      throw new Error(`R2 object ${key} changed under every ranged read.`)
    },

    async head(key) {
      const obj = await bucket.head(key)
      return obj ? toRecord(obj) : null
    },

    async delete(key) {
      await bucket.delete(key)
    },

    async list(options) {
      // R2 reads `limit: 0` as "use the default", so short-circuit it.
      if (options?.limit === 0) {
        return { objects: [] }
      }
      const page = await bucket.list({
        ...(options?.prefix !== undefined ? { prefix: options.prefix } : {}),
        ...(options?.cursor !== undefined ? { cursor: options.cursor } : {}),
        ...(options?.limit !== undefined ? { limit: options.limit } : {}),
        // R2 omits httpMetadata/customMetadata from list rows unless asked.
        include: ['httpMetadata', 'customMetadata'],
      })
      return {
        objects: page.objects.map(toRecord),
        ...(page.truncated ? { cursor: page.cursor, truncated: true } : {}),
      }
    },
  })
}

CREATE TABLE IF NOT EXISTS generation_artifacts (
  artifact_id  text PRIMARY KEY NOT NULL,
  run_id       text NOT NULL,
  thread_id    text NOT NULL,
  blob_key     text,
  name         text NOT NULL,
  mime_type    text NOT NULL,
  size         integer NOT NULL,
  source_url   text,
  created_at   integer NOT NULL
);
CREATE INDEX IF NOT EXISTS generation_artifacts_run ON generation_artifacts (run_id);

import { defineArtifactStore } from '@tanstack/ai-persistence'
import type { ArtifactRecord } from '@tanstack/ai-persistence'

interface ArtifactRow {
  artifact_id: string
  run_id: string
  thread_id: string
  blob_key: string | null
  name: string
  mime_type: string
  size: number
  source_url: string | null
  created_at: number
}

function fromRow(row: ArtifactRow): ArtifactRecord {
  return {
    artifactId: row.artifact_id,
    runId: row.run_id,
    threadId: row.thread_id,
    ...(row.blob_key != null ? { blobKey: row.blob_key } : {}),
    name: row.name,
    mimeType: row.mime_type,
    size: row.size,
    ...(row.source_url != null ? { sourceUrl: row.source_url } : {}),
    createdAt: row.created_at,
  }
}

export function d1ArtifactStore(db: D1Database) {
  return defineArtifactStore({
    async save(record) {
      // Insert or overwrite (artifact ids are unique).
      await db
        .prepare(
          `INSERT INTO generation_artifacts
             (artifact_id, run_id, thread_id, blob_key, name, mime_type, size, source_url, created_at)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
           ON CONFLICT(artifact_id) DO UPDATE SET
             run_id = excluded.run_id, thread_id = excluded.thread_id,
             blob_key = excluded.blob_key, name = excluded.name,
             mime_type = excluded.mime_type, size = excluded.size,
             source_url = excluded.source_url, created_at = excluded.created_at`,
        )
        .bind(
          record.artifactId,
          record.runId,
          record.threadId,
          record.blobKey ?? null,
          record.name,
          record.mimeType,
          record.size,
          record.sourceUrl ?? null,
          record.createdAt,
        )
        .run()
    },

    async get(artifactId) {
      const row = await db
        .prepare(`SELECT * FROM generation_artifacts WHERE artifact_id = ?`)
        .bind(artifactId)
        .first<ArtifactRow>()
      return row ? fromRow(row) : null
    },

    async list(runId) {
      const { results } = await db
        .prepare(`SELECT * FROM generation_artifacts WHERE run_id = ?`)
        .bind(runId)
        .all<ArtifactRow>()
      return results.map(fromRow)
    },

    async delete(artifactId) {
      await db
        .prepare(`DELETE FROM generation_artifacts WHERE artifact_id = ?`)
        .bind(artifactId)
        .run()
    },

    async deleteForRun(runId) {
      await db
        .prepare(`DELETE FROM generation_artifacts WHERE run_id = ?`)
        .bind(runId)
        .run()
    },
  })
}

import {
  defineAIPersistence,
  composePersistence,
  withGenerationPersistence,
} from '@tanstack/ai-persistence'
import { r2BlobStore } from './r2-blob-store'
import { d1ArtifactStore } from './d1-artifact-store'
import { d1GenerationRunStore } from './d1-job-store' // your GenerationRunStore

/** Call inside a request handler — bindings are not available at module scope. */
export function generationPersistence(env: Env) {
  return defineAIPersistence({
    stores: {
      generationRuns: d1GenerationRunStore(env.DB),
      artifacts: d1ArtifactStore(env.DB),
      blobs: r2BlobStore(env.ARTIFACTS_BUCKET),
    },
  })
}

// …or add bytes to a persistence that already has chat + generation runs:
// composePersistence(chatAndJobsPersistence(env), {
//   overrides: {
//     artifacts: d1ArtifactStore(env.DB),
//     blobs: r2BlobStore(env.ARTIFACTS_BUCKET),
//   },
// })

import { generateImage, toServerSentEventsResponse } from '@tanstack/ai'
import { openaiImage } from '@tanstack/ai-openai'
import { withGenerationPersistence } from '@tanstack/ai-persistence'
import { generationPersistence } from './lib/generation-persistence'

export default {
  async fetch(request: Request, env: Env) {
    const { prompt, threadId } = await request.json()
    const stream = generateImage({
      adapter: openaiImage('gpt-image-1'),
      prompt,
      threadId, // the slot recorded on the job + artifacts
      stream: true,
      middleware: [
        // Nothing is buffered: the artifact streams from the provider CDN
        // into R2. Add `maxArtifactBytes: false` if you want no ceiling at
        // all on what an origin can stream into your bucket (the default is
        // 1 GiB); it is not needed for the streaming itself.
        withGenerationPersistence(generationPersistence(env), { threadId }),
      ],
    })
    return toServerSentEventsResponse(stream)
  },
}

import {
  parseRangeHeader,
  retrieveArtifact,
  retrieveBlob,
} from '@tanstack/ai-persistence'
import { generationPersistence } from './lib/generation-persistence'

export async function GET(request: Request, env: Env) {
  const artifactId = new URL(request.url).searchParams.get('id') ?? ''
  const persistence = generationPersistence(env)

  // Authorize before serving — derive the owner from the session, never trust
  // a client-supplied id. (The record carries runId/threadId to check against.)
  const record = await retrieveArtifact(persistence, artifactId)
  if (!record) return new Response('Not found', { status: 404 })

  // `parseRangeHeader` resolves the header against the size on the record —
  // suffix ranges included — so an unsatisfiable range is a 416 here and the
  // store only ever sees a range it can serve.
  const range = parseRangeHeader(request.headers.get('range'), record.size)
  if (range === 'unsatisfiable') {
    return new Response('Range not satisfiable', {
      status: 416,
      headers: { 'content-range': `bytes */${record.size}` },
    })
  }

  // Pass the record, not the id: no second metadata lookup.
  const blob = await retrieveBlob(
    persistence,
    record,
    range ? { range } : undefined,
  )
  if (!blob?.body) return new Response('Not found', { status: 404 })

  // `accept-ranges` on every response, including the whole-file one: it is how
  // a player learns it may seek at all.
  const headers = {
    'content-type': record.mimeType,
    'accept-ranges': 'bytes',
  }
  if (!blob.range) {
    return new Response(blob.body, {
      headers: { ...headers, 'content-length': String(record.size) },
    })
  }
  const { offset, length } = blob.range
  return new Response(blob.body, {
    status: 206,
    headers: {
      ...headers,
      'content-length': String(length),
      'content-range': `bytes ${offset}-${offset + length - 1}/${record.size}`,
    },
  })
}

import { runPersistenceConformance } from '@tanstack/ai-persistence/testkit'
import { env } from 'cloudflare:test'
import { generationPersistence } from '../src/lib/generation-persistence'

// A generation-only Worker declares the chat state stores it does not provide.
runPersistenceConformance('app-r2', () => generationPersistence(env), {
  skip: ['messages', 'runs', 'interrupts', 'metadata'],
})