example.ts

import { z } from 'zod'
import {
  s,
  defineCollection,
  createMockClients,
  StorageConfig,
} from './index'

// === Primitives ===

const gen3Phase = z.enum(['context', 'outline', 'visuals'])

const jsonContent = z.any() // TipTap JSONContent

const gen3ContextOptions = z.object({
  topics: z.array(z.string()),
  audiences: z.array(z.string()),
  goals: z.array(z.string()),
  lengths: z.array(
    z.object({
      label: z.string(),
      numCards: z.number(),
      selected: z.boolean(),
    })
  ),
})

const gen3Source = z.object({
  id: z.string(),
  type: z.enum(['pdf', 'image', 'webpage', 'text', 'pptx', 'docx']),
  name: z.string(),
  addedAt: z.number(),
  status: z.enum(['uploading', 'uploaded', 'processing', 'done', 'error']),
  size: z.number().optional(),
  url: z.string().optional(),
  error: z.string().optional(),
  extraction: z
    .object({
      text: z.string().optional(),
      images: z
        .array(z.object({ url: z.string(), caption: z.string().optional() }))
        .optional(),
      pageCount: z.number().optional(),
    })
    .optional(),
})

const gen3Checkpoint = z.object({
  id: z.string(),
  messageId: z.string(),
  timestamp: z.string(),
  outlineContent: jsonContent,
})

const gen3MessageResult = z.object({
  messageId: z.string(),
  type: z.enum(['generate-outline', 'edit-outline']),
  beforeCheckpointId: z.string().nullable(),
  afterCheckpointId: z.string().nullable(),
})

// Simplified - would use actual UIMessage type from 'ai' SDK
const gen3UIMessage = z.object({
  id: z.string(),
  role: z.enum(['user', 'assistant', 'system']),
  content: z.string(),
  createdAt: z.date().optional(),
  metadata: z.any().optional(),
})

// === Schema Versions ===

const gen3ChatSchemaV1 = z.object({
  // Keys
  workspaceId: z.string(),
  userId: z.string(),
  chatId: z.string(),

  // Inline (normalized)
  messagesVersion: z.literal('v1'),
  totalTokensUsed: z.number(),
  currentPhase: gen3Phase,
  outline: z
    .object({
      state: z.enum(['in-progress', 'complete']),
      content: jsonContent,
    })
    .nullable(),
  contextOptions: gen3ContextOptions.nullable(),
  results: z.record(gen3MessageResult),

  // Paged in DynamoDB - strong consistency, auto-paginated
  messages: s.paged(z.array(gen3UIMessage)),
  checkpoints: s.paged(z.record(gen3Checkpoint)),

  // S3 storage - large extractions, ok with eventual consistency
  sources: s.s3(z.array(gen3Source)),
})

// V2 adds createdAt/updatedAt
const gen3ChatSchemaV2 = gen3ChatSchemaV1.extend({
  createdAt: z.date(),
  updatedAt: z.date(),
})

// === Collection Definition ===

const gen3ChatCollection = defineCollection({
  name: 'gen3-chats',
  schema: gen3ChatSchemaV2,
  keys: {
    pk: ['workspaceId', 'userId', 'chatId'],
  },
  versions: {
    1: gen3ChatSchemaV1,
    2: gen3ChatSchemaV2,
  },
  migrations: {
    '1->2': (v1) => ({
      ...v1,
      createdAt: new Date(),
      updatedAt: new Date(),
    }),
  },
  minVersion: 1,
})

// === Usage Example ===

async function example() {
  // Setup storage (use real AWS SDK clients in production)
  const { dynamo, s3 } = createMockClients()
  const storage: StorageConfig = {
    dynamoTable: 'gen3-chats',
    s3Bucket: 'gamma-chat-data',
    dynamo,
    s3,
  }

  // Create collection instance
  const chats = gen3ChatCollection(storage)

  // Create a new chat
  await chats.put({
    workspaceId: 'ws-123',
    userId: 'usr-456',
    chatId: 'chat-789',
    messagesVersion: 'v1',
    totalTokensUsed: 0,
    currentPhase: 'context',
    outline: null,
    contextOptions: null,
    results: {},
    messages: [],
    sources: [],
    checkpoints: {},
    createdAt: new Date(),
    updatedAt: new Date(),
  })

  // Fetch it back
  const chat = await chats.get({
    workspaceId: 'ws-123',
    userId: 'usr-456',
    chatId: 'chat-789',
  })

  console.log('Retrieved chat:', chat)

  // Update with messages
  if (chat) {
    chat.messages.push({
      id: 'msg-1',
      role: 'user',
      content: 'Create a presentation about AI',
    })
    chat.totalTokensUsed = 150
    chat.updatedAt = new Date()
    await chats.put(chat)
  }

  // Delete
  await chats.delete({
    workspaceId: 'ws-123',
    userId: 'usr-456',
    chatId: 'chat-789',
  })
}

example().catch(console.error)