unstable_onSetAIState: async ({ state, done }) => {
    'use server'

    const session = await auth()

    if (session && session.user) {

      // only save it once it's done. No need to save it after each streamed result
      if (done && !state.saved) {
        await addOrUpdateChat(state, session)
        state.saved = true
      }
    } else {
      return
    }
  }
})

let state = aiState.get()
        state.messages = [
          ...state.messages,
          {
            id: nanoid(),
            role: 'assistant',
            content
          }
        ]

        await saveChat(state)

        reply.done()
        aiState.done({ ...state, done: isFinal, saved: true })

async function saveChat(state: AIState, session: Session | null = null) {
  if (!session) {
    session = await auth()
    if (!session?.user?.id) {
      throw new Error('User not Authorized to save')
    }
  }

  const { chatId, messages } = state

  const createdAt = Date.now()
  const userId = session.user?.id as string
  const path = `/chat/${chatId}`
  const title = messages[0].content.substring(0, 100)

  const chat: Chat = {
    id: chatId,
    title,
    userId,
    createdAt,
    path,
    messages
  }

  await dynamoDbService.saveChat(chat)
}

useEffect(() => {
    setNewChatId(id)
  })

  useEffect(() => {
    // This is here, because it resolve a bug when you go to 'new chat' --> clear history, and then generate a new chat, the IDs were out of sync.
    updateAIState({ ...aiState, chatId: id! })
    setNewChatId(id)
  })

async function updateAIState(currentAIState: any) {
  'use server'
  const aiState = getMutableAIState<typeof AI>()
  aiState.done({ ...currentAIState })
}

export type AIState = {
  chatId: string
  messages: Message[]
  saved?: boolean
}

export type UIState = {
  id: string
  display: React.ReactNode
}[]


async function addOrUpdateChat(state: AIState, session: Session | null = null) {
  if (!session) {
    session = await auth()
    if (!session?.user?.id) {
      throw new Error('User not Authorized to save')
    }
  }

  const { chatId, messages } = state

  const createdAt = new Date()
  const userId = session?.user?.id as string
  const path = `/chat/${chatId}`

  const firstMessageContent = messages[0].content as string
  const title = firstMessageContent.substring(0, 100)

  const flowchart = await getFlowchart(chatId)

  const chat: Chat = {
    id: chatId,
    title,
    userId,
    createdAt,
    path,
    messages,
    flowchart: flowchart || { nodes: [], edges: [] } // Use existing flowchart data if available, otherwise initialize with empty arrays
  }

  await saveChat(chat)
}

export const AI = createAI<AIState, UIState>({
  actions: {
    submitUserMessage
  },
  initialUIState: [],
  initialAIState: { chatId: nanoid(), messages: [] },
  onGetUIState: async () => {
    'use server'

    const session = await auth()

    if (session && session.user) {
      const aiState = getAIState()

      if (aiState) {
        const uiState = getUIStateFromAIState(aiState)
        return uiState
      }
    } else {
      return
    }
  },
  onSetAIState: async ({ state, done }) => {
    'use server'

    const session = await auth()

    if (session && session.user) {
      // only save it once it's done. No need to save it after each streamed result
      if (done && !state.saved) {
        await addOrUpdateChat(state, session)
        state.saved = true
      }
    } else {
      return
    }
  }
})

import { ObjectId } from 'bson'
import { db } from '@/lib/db'
import { auth } from '@/auth'
import { getMissingKeys } from '@/app/actions'
import { AI } from '@/lib/chat/actions'
import { Chat } from '@/components/chat'
import { Session } from '@/lib/types'
import { redirect } from 'next/navigation'

export const metadata = {
  title: 'Sudeep Kudari'
}

const getChatsCustom = async (userId: string) => {
  return await db.chat.findMany({
    where: { userId },
    include: {
      messages: true
    },
    orderBy: { createdAt: 'desc' }
  })

}

const IndexPage = async () => {
  const session = (await auth()) as Session
  const missingKeys = await getMissingKeys()

  if (session?.user) {
    const newData = await getChatsCustom(session.user.id)
    console.log(newData)
    if (newData[0]?.messages?.length > 0) {
      return redirect(`/chat/${newData[0].id}`)
    } else {
      return redirect('/new')
    }
  }

  const chatId = new ObjectId().toHexString()
  return (
    <AI initialAIState={{ chatId, messages: [] }}>
      <Chat id={chatId} session={session} missingKeys={missingKeys} />
    </AI>
  )
}

export default IndexPage

import { ObjectId } from 'bson'
import { db } from '@/lib/db'
import { auth } from '@/auth'
import { getMissingKeys } from '@/app/actions'
import { Session } from '@/lib/types'
import { redirect } from 'next/navigation'

const createChatWithDynamicPath = async (userId: string) => {
  const chatId = new ObjectId().toHexString()
  const path = `/chat/${chatId}`
  return await db.chat.create({
    data: {
      id: chatId,
      userId,
      title: 'New Chat',
      path
    }
  })
}

export default async function NewPage() {

  const session = (await auth()) as Session

  if (session?.user) {
    const newData = await createChatWithDynamicPath(session.user.id)
    return redirect(`/chat/${newData.id}`)
  }
}

// This is used to fix Vercel ai/rsc bug where the state is not correctly maintained/updated
 //    causing the aiState to get out of sync.
 // It shows up when you navigate away from the current route and then then back to the same route quickly.
 // This invalidates the cache and refreshes the AI states
 useEffect(() => {
   router.refresh()
 }, [router])

 useEffect(() => {
   // Do not update the aiState and local key chat Id if its the same. If this is removed, it was creating several new chatIds, getting out of sync when saving.
   if (aiState.chat.chatId === chatId) return
   // Update the aiState chatId when going to a new chat. If not, it will save with the wrong chatId.
   aiState.chat.chatId = chatId!
   setAIState(aiState)

   // Update the local storage chat Id.
   setNewChatId(chatId)
 }, [aiState, chatId, setAIState, setNewChatId])

const aiStateDone = async (aiState: MutableAIState<typeof AI>, newMessage: Message, properties?: Partial<AIState>) => {
  const messages = []
  if (newMessage) {
    messages.push(newMessage)
  }
  let state = {
    ...aiState.get(),
    ...properties,
    chat: {
      ...aiState.get().chat,
      messages: [...aiState.get().chat.messages, ...messages]
    }
  }

  // Doing it here to resolve a bug in aiStateDone where if addOrUpdate took a while, the UI refreshed...
  //    before the db was updated.
  await addOrUpdateChat(state)

  await aiState.done(state)
}