supabase/apps/studio/hooks/analytics/useLegacyLogsRewrite.ts

import { useReducer } from 'react'

import { constructHeaders } from '@/data/fetchers'
import { rewriteLogsSqlWithAI } from '@/data/logs/logs-sql-rewrite'
import { useLogsAttributeKeys } from '@/hooks/analytics/useLogsAttributeKeys'
import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { getErrorMessage } from '@/lib/get-error-message'

export type LegacyLogsRewriteState =
  | { status: 'offered' }
  | { status: 'rewriting' }
  | { status: 'failed'; message: string }
  | { status: 'noRewriteNeeded' }
  | { status: 'dismissed' }

export type LegacyLogsRewriteEvent =
  | { type: 'rewriteRequested' }
  | { type: 'rewriteProposed' }
  | { type: 'rewriteFailed'; message: string }
  | { type: 'rewriteNoop' }
  | { type: 'dismissed' }

export const INITIAL_LEGACY_LOGS_REWRITE_STATE: LegacyLogsRewriteState = { status: 'offered' }

/**
 * The events each state accepts. Anything absent is an invalid transition and
 * leaves the state untouched — notably `dismissed` is terminal, and the offer
 * can't be dismissed mid-rewrite.
 */
const VALID_EVENTS: {
  [S in LegacyLogsRewriteState['status']]: readonly LegacyLogsRewriteEvent['type'][]
} = {
  offered: ['rewriteRequested', 'dismissed'],
  rewriting: ['rewriteProposed', 'rewriteFailed', 'rewriteNoop'],
  // A failure is recoverable: the same Rewrite control retries it.
  failed: ['rewriteRequested', 'dismissed'],
  noRewriteNeeded: ['dismissed'],
  dismissed: [],
}

function targetState(event: LegacyLogsRewriteEvent): LegacyLogsRewriteState {
  switch (event.type) {
    case 'rewriteRequested':
      return { status: 'rewriting' }
    // The proposal is handed to the caller; the offer returns to idle behind it so
    // it's ready again if the user discards the proposal.
    case 'rewriteProposed':
      return { status: 'offered' }
    case 'rewriteFailed':
      return { status: 'failed', message: event.message }
    case 'rewriteNoop':
      return { status: 'noRewriteNeeded' }
    case 'dismissed':
      return { status: 'dismissed' }
  }
}

export function legacyLogsRewriteReducer(
  state: LegacyLogsRewriteState,
  event: LegacyLogsRewriteEvent
): LegacyLogsRewriteState {
  if (!VALID_EVENTS[state.status].includes(event.type)) return state
  return targetState(event)
}

const CHANGED_WHILE_REWRITING_MESSAGE =
  'The query changed while the Assistant was working, so the rewrite no longer matches it.'

const NO_RESPONSE_MESSAGE = 'The Assistant did not respond. Try again.'

export type LegacyLogsRewriteProposal = { original: string; modified: string }

type UseLegacyLogsRewriteArgs = {
  /**
   * Reads the query to rewrite at the moment the user asks. A callback rather than
   * a value so the rewrite operates on exactly what the user sees, not on whatever
   * a surface last rendered.
   */
  readSql: () => string
  /** Receives a rewrite worth reviewing. Each surface routes this to its own diff. */
  onProposal: (proposal: LegacyLogsRewriteProposal) => void
  /**
   * Called when the offer is dismissed, for surfaces that persist that. The
   * machine covers the current session only — a surface that remembers dismissals
   * across sessions layers that on top of its own visibility check, since a value
   * read from storage isn't available in time to seed the machine.
   */
  onDismissed?: () => void
}

/**
 * Owns the BigQuery → ClickHouse rewrite request end to end: key discovery, the
 * completion call, the stale-edit guard, no-op detection, and the resulting state.
 */
export function useLegacyLogsRewrite({
  readSql,
  onProposal,
  onDismissed,
}: UseLegacyLogsRewriteArgs) {
  const { data: project } = useSelectedProjectQuery()
  const { data: organization } = useSelectedOrganizationQuery()
  const projectRef = project?.ref

  const [state, dispatch] = useReducer(legacyLogsRewriteReducer, INITIAL_LEGACY_LOGS_REWRITE_STATE)

  const { fetchAttributeKeys } = useLogsAttributeKeys()

  const requestRewrite = async () => {
    if (!projectRef) return console.error('[useLegacyLogsRewrite] Project ref is required')

    const currentSql = readSql()
    if (currentSql.trim().length === 0) return

    dispatch({ type: 'rewriteRequested' })
    try {
      const [headerData, availableKeys] = await Promise.all([
        constructHeaders(),
        fetchAttributeKeys(currentSql),
      ])
      const rewritten = await rewriteLogsSqlWithAI({
        sql: currentSql,
        projectRef,
        connectionString: project?.connectionString,
        orgSlug: organization?.slug,
        authorizationHeader: headerData.get('Authorization'),
        availableKeys,
      })

      // The user may have kept typing while the model worked; a proposal built from
      // stale text would clobber those edits when accepted.
      if (readSql() !== currentSql) {
        dispatch({ type: 'rewriteFailed', message: CHANGED_WHILE_REWRITING_MESSAGE })
        return
      }

      // An unchanged response means the query already runs on ClickHouse and the
      // dialect heuristic was over-eager. Proposing it would show an empty diff.
      if (rewritten.trim() === currentSql.trim()) {
        dispatch({ type: 'rewriteNoop' })
        return
      }

      onProposal({ original: currentSql, modified: rewritten })
      dispatch({ type: 'rewriteProposed' })
    } catch (error) {
      dispatch({ type: 'rewriteFailed', message: getErrorMessage(error, NO_RESPONSE_MESSAGE) })
    }
  }

  const dismiss = () => {
    const dismissed: LegacyLogsRewriteEvent = { type: 'dismissed' }
    // The transition table is the contract, not the UI that happens to disable the
    // control: never report a dismissal the machine rejected (mid-rewrite, say),
    // or a surface that persists it would suppress an offer that's still live.
    if (legacyLogsRewriteReducer(state, dismissed) === state) return
    dispatch(dismissed)
    onDismissed?.()
  }

  return { state, requestRewrite, dismiss }
}