supabase/apps/studio/components/interfaces/Explorer/hooks.ts

import { useRouter } from 'next/router'
import { useEffect, useEffectEvent, useState } from 'react'

import { useNotebookQuery } from '@/data/content/notebooks/notebook-query'
import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { generateUuid } from '@/lib/api/snippets.browser'
import { useProfile } from '@/lib/profile'
import type { AssistantModel } from '@/state/ai-assistant-state'
import { useAiAssistantState, whenAiAssistantInitialized } from '@/state/ai-assistant-state'
import { useExplorerQueryStateSnapshot } from '@/state/explorer-query'
import { readNotebookDraft } from '@/state/notebooks/notebook-drafts'
import { useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state'
import { type Notebook } from '@/state/notebooks/types'
import { Notebooks } from '@/types'

/**
 * Fetches a notebook's content by id and merges it into the valtio store, so landing on
 * a notebook any way other than creating it in this session (direct link, hard refresh,
 * clicking it from the nav list) still hydrates `notebooksState`.
 *
 * A notebook that isn't in the store yet is also checked for a locally-persisted draft of unsaved edits
 * If the notebook loaded from the server, the draft is restored on top of it, otherwise if
 * the notebook was never saved at all, the draft — if present — is the only copy
 * that ever existed, so it's restored as a new local-only notebook instead.
 */
export const useLoadNotebook = ({ id, projectRef }: { id?: string; projectRef?: string }) => {
  const notebooksSnap = useNotebooksStateSnapshot()
  const { profile } = useProfile()
  const { data: project } = useSelectedProjectQuery()
  const currentNotebook = id ? notebooksSnap.notebooks[id] : undefined

  const isCurrentProjectNotebook = currentNotebook?.projectRef === projectRef
  const isNewLocalNotebook = isCurrentProjectNotebook && currentNotebook?.status === 'new'
  const hasLoadedNotebook =
    isCurrentProjectNotebook && currentNotebook?.notebook.content !== undefined

  const { data, error, isError, isLoading } = useNotebookQuery(
    { projectRef, id },
    {
      retry: false,
      enabled: !isNewLocalNotebook && !hasLoadedNotebook,
    }
  )

  const isNotFound = isError && error.code === 404

  const mergeNotebook = useEffectEvent(() => {
    if (!projectRef || !id) return

    if (data) {
      const isFreshLoad = !isCurrentProjectNotebook
      notebooksSnap.setNotebook({ projectRef, notebook: data })
      if (isFreshLoad) {
        notebooksSnap.restoreDraft({ projectRef, id, baseUpdatedAt: data.updated_at })
      }
      return
    }

    if (isNotFound && !isCurrentProjectNotebook && profile && project) {
      const draft = readNotebookDraft({ projectRef, id })
      if (!draft) return

      notebooksSnap.addNotebook({
        projectRef,
        notebook: {
          id,
          type: 'notebook',
          name: draft.name,
          description: '',
          visibility: 'project',
          favorite: false,
          content: draft.content,
          owner_id: profile.id,
          project_id: project.id,
        },
      })
    }
  })

  useEffect(() => {
    mergeNotebook()
  }, [projectRef, id, data, isNotFound, !!profile, !!project])

  return { isNotFound: isNotFound && !isCurrentProjectNotebook, isLoading }
}

export const useCreateNotebook = () => {
  const router = useRouter()
  const { profile } = useProfile()
  const { data: project } = useSelectedProjectQuery()
  const notebooksSnap = useNotebooksStateSnapshot()

  const createNotebook = ({
    id: idOverride,
    name,
    cells,
  }: { id?: string; name?: string; cells?: Notebooks.Content['cells'] } = {}) => {
    if (!profile) return console.error('Profile is required')
    if (!project) return console.error('Project is required')

    const id = idOverride ?? generateUuid()

    const notebook: Notebook = {
      id,
      type: 'notebook',
      name: name ?? 'New Notebook',
      description: '',
      visibility: 'project',
      favorite: false,
      content: {
        schema_version: 1,
        cells: cells ?? [],
      },
      owner_id: profile.id,
      project_id: project.id,
    }

    notebooksSnap.addNotebook({ projectRef: project.ref, notebook })
    notebooksSnap.addNeedsSaving(notebook.id)

    router.push(`/project/${project.ref}/explorer/notebook/${notebook.id}`)
  }

  return { createNotebook }
}

export const useCreateChat = () => {
  const router = useRouter()
  const { data: project } = useSelectedProjectQuery()
  const { data: organization } = useSelectedOrganizationQuery()
  const aiAssistantState = useAiAssistantState()

  const [isCreating, setIsCreating] = useState(false)

  const openChat = (id: string) => {
    if (!project) {
      console.error('Project is required')
      return undefined
    }

    router.push(`/project/${project.ref}/explorer/chat/${id}`)
  }

  const createChat = async ({
    name,
    initialMessage,
    model,
  }: {
    name?: string
    initialMessage?: string
    model?: AssistantModel
  } = {}) => {
    if (!project) {
      console.error('Project is required')
      return undefined
    }

    setIsCreating(true)

    try {
      // Hydration replaces the chat map and the selected model wholesale, so wait it out before
      // creating anything — otherwise the new chat is dropped as soon as the persisted state lands.
      await whenAiAssistantInitialized(aiAssistantState)

      aiAssistantState.setContext({
        projectRef: project.ref,
        orgSlug: organization?.slug,
        connectionString: project.connectionString ?? '',
      })
      if (model) aiAssistantState.setModel(model)

      const id = aiAssistantState.createChat({ name, initialMessage })

      router.push(`/project/${project.ref}/explorer/chat/${id}`)
      return id
    } finally {
      setIsCreating(false)
    }
  }

  return { createChat, openChat, isCreating }
}

/** Opens a new Explorer chat that asks the assistant to run a saved notebook and summarize its results */
export const useAnalyzeNotebook = () => {
  const { createChat, isCreating } = useCreateChat()

  const analyzeNotebook = ({ id, name }: { id?: string; name?: string }) =>
    createChat({
      name: `Analyze ${name} notebook`,
      initialMessage: `Run the notebook "${name}" (id: ${id}) and analyze the results. Summarize the key findings per cell, calling out anomalies or trends, and use any markdown cells for context. Skip or flag any cell that would mutate data rather than running it.`,
    })

  return { analyzeNotebook, isCreating }
}

export const useCreateQuery = () => {
  const router = useRouter()
  const { data: project } = useSelectedProjectQuery()
  const querySnap = useExplorerQueryStateSnapshot()

  const createQuery = ({
    sql,
    name,
    autoRun,
    replace = false,
  }: { sql?: string; name?: string; autoRun?: boolean; replace?: boolean } = {}) => {
    if (!project) return console.error('Project is required')

    const id = generateUuid()
    querySnap.createDraft({ id, projectRef: project.ref, sql, name, autoRun })

    const url = `/project/${project.ref}/explorer/query/${id}`
    if (replace) router.replace(url)
    else router.push(url)

    return id
  }

  return { createQuery, projectRef: project?.ref }
}