simstudioai--sim/apps/sim/hooks/queries/workspace-files.ts

import { toast } from '@sim/emcn'
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import { sleep } from '@sim/utils/helpers'
import { backoffWithJitter } from '@sim/utils/retry'
import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { ApiClientError, isApiClientError } from '@/lib/api/client/errors'
import { requestJson } from '@/lib/api/client/request'
import { getUsageLimitsContract } from '@/lib/api/contracts/usage-limits'
import {
  deleteWorkspaceFileContract,
  listWorkspaceFilesContract,
  registerWorkspaceFileContract,
  renameWorkspaceFileContract,
  restoreWorkspaceFileContract,
  updateWorkspaceFileContentContract,
} from '@/lib/api/contracts/workspace-files'
import {
  DirectUploadError,
  runUploadStrategy,
  type UploadProgressEvent,
} from '@/lib/uploads/client/direct-upload'
import type { WorkspaceFileRecord } from '@/lib/uploads/contexts/workspace'
import type { UserFile } from '@/executor/types'
import { useFileContentSource } from '@/hooks/use-file-content-source'

const logger = createLogger('WorkspaceFilesQuery')

type WorkspaceFileQueryScope = 'active' | 'archived' | 'all'

/**
 * Query key factories for workspace files
 */
export const workspaceFilesKeys = {
  all: ['workspaceFiles'] as const,
  lists: () => [...workspaceFilesKeys.all, 'list'] as const,
  workspaceLists: (workspaceId: string) => [...workspaceFilesKeys.lists(), workspaceId] as const,
  list: (workspaceId: string, scope: WorkspaceFileQueryScope = 'active') =>
    [...workspaceFilesKeys.workspaceLists(workspaceId), scope] as const,
  contents: () => [...workspaceFilesKeys.all, 'content'] as const,
  contentFile: (workspaceId: string, fileId: string) =>
    [...workspaceFilesKeys.contents(), workspaceId, fileId] as const,
  content: (
    workspaceId: string,
    fileId: string,
    mode: 'text' | 'raw' | 'binary' = 'text',
    storageKey?: string
  ) =>
    [
      ...workspaceFilesKeys.contentFile(workspaceId, fileId),
      mode,
      ...(storageKey ? [storageKey] : []),
    ] as const,
  storageInfo: () => [...workspaceFilesKeys.all, 'storageInfo'] as const,
}

export const WORKSPACE_FILES_LIST_STALE_TIME = 30 * 1000
export const WORKSPACE_FILE_CONTENT_STALE_TIME = 30 * 1000
export const WORKSPACE_FILE_BINARY_STALE_TIME = 30 * 1000
export const WORKSPACE_STORAGE_INFO_STALE_TIME = 60 * 1000

/**
 * Storage info type
 */
interface StorageInfo {
  usedBytes: number
  limitBytes: number
  percentUsed: number
  plan?: string
}

/**
 * Hook to fetch a single workspace file record by ID.
 * Shares the `list(workspaceId, 'active')` query key with {@link useWorkspaceFiles} so no extra
 * network request is made when the list is already cached (warm path).
 * On a cold path (e.g. direct navigation to a file URL), this fetches the full active file list
 * for the workspace and selects the matching record via `select`.
 */
export function useWorkspaceFileRecord(workspaceId: string, fileId: string) {
  return useQuery({
    queryKey: workspaceFilesKeys.list(workspaceId, 'active'),
    queryFn: ({ signal }) => fetchWorkspaceFiles(workspaceId, 'active', signal),
    enabled: !!workspaceId && !!fileId,
    staleTime: WORKSPACE_FILES_LIST_STALE_TIME,
    select: (files) => files.find((f) => f.id === fileId) ?? null,
  })
}

/**
 * Fetch workspace files from API
 */
async function fetchWorkspaceFiles(
  workspaceId: string,
  scope: WorkspaceFileQueryScope = 'active',
  signal?: AbortSignal
): Promise<WorkspaceFileRecord[]> {
  const data = await requestJson(listWorkspaceFilesContract, {
    params: { id: workspaceId },
    query: { scope },
    signal,
  })
  return data.success ? data.files : []
}

/**
 * Hook to fetch workspace files
 */
export function useWorkspaceFiles(
  workspaceId: string,
  scope: WorkspaceFileQueryScope = 'active',
  options?: { enabled?: boolean }
) {
  return useQuery({
    queryKey: workspaceFilesKeys.list(workspaceId, scope),
    queryFn: ({ signal }) => fetchWorkspaceFiles(workspaceId, scope, signal),
    enabled: !!workspaceId && (options?.enabled ?? true),
    staleTime: WORKSPACE_FILES_LIST_STALE_TIME, // 30 seconds - files can change frequently
    placeholderData: keepPreviousData, // Show cached data immediately
  })
}

/**
 * Fetch file content as text via a content-source URL
 */
async function fetchWorkspaceFileContent(url: string, signal?: AbortSignal): Promise<string> {
  // boundary-raw-fetch: binary/text download, response is not JSON
  const response = await fetch(url, { signal, cache: 'no-store' })

  if (!response.ok) {
    throw new Error('Failed to fetch file content')
  }

  return response.text()
}

/**
 * Hook to fetch workspace file content as text.
 * `key` (the storage object key) is forwarded into the query key factory so that a new
 * storage key (e.g. after a file is re-uploaded) correctly busts the cache.
 *
 * `refetchInterval` lets a caller poll while waiting for the server content to advance — the
 * editor's post-stream reconcile (see `use-editable-file-content.ts`) exits only when a fetch
 * returns content that moved past its baseline, and would otherwise wedge read-only forever if
 * its single refetch raced the agent's write. The function form is re-evaluated by react-query
 * after every fetch and options pass, so a condition read through a ref stops the polling as soon
 * as it flips — no re-render required.
 */
export function useWorkspaceFileContent(
  workspaceId: string,
  fileId: string,
  key: string,
  raw?: boolean,
  options?: { refetchInterval?: number | false | (() => number | false) }
) {
  const source = useFileContentSource()
  return useQuery({
    queryKey: workspaceFilesKeys.content(workspaceId, fileId, raw ? 'raw' : 'text', key),
    queryFn: ({ signal }) =>
      fetchWorkspaceFileContent(source.buildUrl(key, { raw, bust: true }), signal),
    enabled: !!workspaceId && !!fileId && !!key,
    staleTime: WORKSPACE_FILE_CONTENT_STALE_TIME,
    refetchOnWindowFocus: 'always',
    refetchInterval: options?.refetchInterval ?? false,
  })
}

/**
 * Thrown when the serve route returns 409 — a generated document (pptx/docx/pdf/
 * xlsx) whose source is still being written/compiled. Distinct from a real fetch
 * failure so the binary query can keep retrying (and the preview keeps showing
 * its loading state) until the compiled artifact is ready.
 */
export class DocNotReadyError extends Error {
  constructor() {
    super('Document is still being generated')
    this.name = 'DocNotReadyError'
  }
}

/**
 * Fetch compiled/binary file content via the serve URL.
 *
 * A `version` (the file record's `updatedAt`) makes the URL content-immutable: the
 * serve route marks versioned responses `immutable`, so the browser HTTP cache
 * resolves re-opens and focus refetches with no round trip. Generated docs are
 * edited in place (same storage key), so an unversioned caller cannot assume
 * immutability and instead busts + bypasses the cache to always read fresh. A 409
 * means a generated doc is still compiling — surfaced as {@link DocNotReadyError}
 * so the query keeps polling.
 */
async function fetchWorkspaceFileBinary(
  url: string,
  version: string | number | undefined,
  signal?: AbortSignal
): Promise<ArrayBuffer> {
  const init: RequestInit = version != null ? { signal } : { signal, cache: 'no-store' }
  // boundary-raw-fetch: binary download consumed as ArrayBuffer
  const response = await fetch(url, init)
  if (response.status === 409) throw new DocNotReadyError()
  if (!response.ok) throw new Error('Failed to fetch file content')
  return response.arrayBuffer()
}

/**
 * Hook to fetch workspace file content as binary (ArrayBuffer).
 * `key` (the storage object key) is forwarded into the query key factory so that a new
 * storage key (e.g. after a file is re-uploaded) correctly busts the cache.
 *
 * `options.version` is a content version (the record's `updatedAt`) folded into the
 * query key. Generated docs are edited IN PLACE — `edit_content` keeps the SAME
 * storage key — so without a version the cache is never busted and the open
 * preview keeps showing the stale binary after a regenerate. Versioning the key
 * makes the preview refetch whenever the file's content changes (and on first
 * open, keyed to the current content rather than a stale cached entry).
 */
export function useWorkspaceFileBinary(
  workspaceId: string,
  fileId: string,
  key: string,
  options?: { enabled?: boolean; version?: string | number }
) {
  const source = useFileContentSource()
  return useQuery({
    queryKey:
      options?.version != null
        ? [...workspaceFilesKeys.content(workspaceId, fileId, 'binary', key), options.version]
        : workspaceFilesKeys.content(workspaceId, fileId, 'binary', key),
    queryFn: ({ signal }) =>
      fetchWorkspaceFileBinary(
        source.buildUrl(key, { version: options?.version, bust: true }),
        options?.version,
        signal
      ),
    // Callers gate this on a readiness signal (e.g. the file has committed
    // content) so we don't 409-poll the serve route for a generated doc whose
    // compiled artifact hasn't been written yet — the doc is fetched once, when
    // it's actually ready, instead of hammering the serve URL through generation.
    enabled: !!workspaceId && !!fileId && !!key && (options?.enabled ?? true),
    staleTime: WORKSPACE_FILE_BINARY_STALE_TIME,
    refetchOnWindowFocus: 'always',
    placeholderData: keepPreviousData,
    // While a generated doc is still compiling, serve returns 409. Poll (stay in
    // the loading state) until the artifact is ready instead of surfacing an
    // error. The artifact is written before the source commits, so a fresh serve
    // normally hits immediately; this only bridges S3 read-after-write lag and the
    // brief mid-generation window. Poll on a short jittered backoff (~0.6s rising
    // to ~2.5s, ~30s budget) so the common case recovers fast without hammering the
    // serve URL on the long tail. SSE content invalidation also re-fetches when the
    // file actually updates.
    retry: (failureCount, error) =>
      error instanceof DocNotReadyError ? failureCount < 14 : failureCount < 2,
    retryDelay: (failureCount, error) =>
      error instanceof DocNotReadyError
        ? backoffWithJitter(failureCount, null, { baseMs: 600, maxMs: 2500 })
        : Math.min(1000 * 2 ** failureCount, 5000),
  })
}

/**
 * Fetch storage info from API
 */
async function fetchStorageInfo(signal?: AbortSignal): Promise<StorageInfo | null> {
  try {
    const data = await requestJson(getUsageLimitsContract, { signal })

    if (data.success && data.storage) {
      return {
        usedBytes: data.storage.usedBytes,
        limitBytes: data.storage.limitBytes,
        percentUsed: data.storage.percentUsed,
        plan: data.usage?.plan || 'free',
      }
    }

    return null
  } catch (error) {
    if (isApiClientError(error) && error.status === 404) {
      return null
    }
    throw error
  }
}

/**
 * Hook to fetch storage info
 */
export function useStorageInfo(enabled = true) {
  return useQuery({
    queryKey: workspaceFilesKeys.storageInfo(),
    queryFn: ({ signal }) => fetchStorageInfo(signal),
    enabled,
    retry: false, // Don't retry on 404
    staleTime: WORKSPACE_STORAGE_INFO_STALE_TIME, // 1 minute - storage info doesn't change often
  })
}

/**
 * Upload workspace file mutation
 */
interface UploadFileParams {
  workspaceId: string
  file: File
  folderId?: string | null
  onProgress?: (event: UploadProgressEvent) => void
  signal?: AbortSignal
  skipToast?: boolean
  skipInvalidation?: boolean
}

interface UploadFileResponse {
  success: boolean
  file: UserFile
}

async function uploadViaApiFallback(
  workspaceId: string,
  file: File,
  folderId?: string | null,
  signal?: AbortSignal
): Promise<UploadFileResponse> {
  const formData = new FormData()
  formData.append('file', file)
  if (folderId) formData.append('folderId', folderId)

  // boundary-raw-fetch: multipart/form-data fallback upload, requestJson only supports JSON bodies
  const response = await fetch(`/api/workspaces/${workspaceId}/files`, {
    method: 'POST',
    body: formData,
    signal,
  })

  return parseUploadResponse(response, 'Upload failed')
}

async function parseUploadResponse(
  response: Response,
  fallbackMessage: string
): Promise<UploadFileResponse> {
  let data: { success?: boolean; error?: string; file?: UserFile } | null = null
  try {
    data = await response.json()
  } catch {}

  if (!response.ok || !data?.success) {
    throw new Error(data?.error || `${fallbackMessage} (${response.status})`)
  }
  return data as UploadFileResponse
}

async function uploadWorkspaceFile(
  workspaceId: string,
  file: File,
  folderId?: string | null,
  onProgress?: (event: UploadProgressEvent) => void,
  signal?: AbortSignal
): Promise<UploadFileResponse> {
  let result
  try {
    result = await runUploadStrategy({
      file,
      presignedEndpoint: `/api/workspaces/${workspaceId}/files/presigned`,
      presignedBody: { folderId },
      workspaceId,
      context: 'workspace',
      onProgress,
      signal,
    })
  } catch (error) {
    if (error instanceof DirectUploadError && error.code === 'FALLBACK_REQUIRED') {
      return uploadViaApiFallback(workspaceId, file, folderId, signal)
    }
    throw error
  }

  const data = await registerWithRetry(workspaceId, result, folderId, signal)

  if (!data.success || !data.file) {
    throw new Error(data.error || 'Failed to register file')
  }
  return { success: true, file: data.file }
}

const REGISTER_MAX_ATTEMPTS = 3
const REGISTER_RETRY_DELAY_MS = 500

/**
 * Register the uploaded object with bounded retries. The server-side handler
 * is idempotent (existing-record short-circuit), so safely retrying handles
 * dropped responses that would otherwise orphan the object in storage.
 */
async function registerWithRetry(
  workspaceId: string,
  result: { key: string; name: string; contentType: string },
  folderId?: string | null,
  signal?: AbortSignal
) {
  let lastError: unknown
  for (let attempt = 1; attempt <= REGISTER_MAX_ATTEMPTS; attempt++) {
    try {
      return await requestJson(registerWorkspaceFileContract, {
        params: { id: workspaceId },
        body: {
          key: result.key,
          name: result.name,
          contentType: result.contentType,
          folderId,
        },
        signal,
      })
    } catch (error) {
      lastError = error
      if (signal?.aborted) throw error
      const isTransient =
        !(error instanceof ApiClientError) || (error.status >= 500 && error.status < 600)
      if (!isTransient || attempt === REGISTER_MAX_ATTEMPTS) throw error
      await sleep(REGISTER_RETRY_DELAY_MS * attempt)
    }
  }
  throw lastError
}

export function useUploadWorkspaceFile() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: ({ workspaceId, file, folderId, onProgress, signal }: UploadFileParams) =>
      uploadWorkspaceFile(workspaceId, file, folderId, onProgress, signal),
    onSettled: (_data, _error, variables) => {
      if (variables.skipInvalidation) return
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.workspaceLists(variables.workspaceId),
      })
      queryClient.invalidateQueries({ queryKey: workspaceFilesKeys.storageInfo() })
    },
    onSuccess: (_data, variables) => {
      if (!variables.skipToast) {
        toast.success(`Uploaded "${variables.file.name}"`)
      }
    },
    onError: (error, variables) => {
      logger.error('Failed to upload file:', error)
      if (!variables.skipToast) {
        toast.error(`Failed to upload "${variables.file.name}": ${error.message}`, {
          duration: 5000,
        })
      }
    },
  })
}

/**
 * Update workspace file content mutation
 */
interface UpdateFileContentParams {
  workspaceId: string
  fileId: string
  content: string
  encoding?: 'base64' | 'utf-8'
}

export function useUpdateWorkspaceFileContent() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: async ({ workspaceId, fileId, content, encoding }: UpdateFileContentParams) => {
      return requestJson(updateWorkspaceFileContentContract, {
        params: { id: workspaceId, fileId },
        body: encoding ? { content, encoding } : { content },
      })
    },
    onSettled: (_data, _error, variables) => {
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.contentFile(variables.workspaceId, variables.fileId),
      })
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.workspaceLists(variables.workspaceId),
      })
      queryClient.invalidateQueries({ queryKey: workspaceFilesKeys.storageInfo() })
    },
    onError: (error) => {
      logger.error('Failed to update file content:', error)
    },
  })
}

/**
 * Rename a workspace file
 */
interface RenameFileParams {
  workspaceId: string
  fileId: string
  name: string
}

export function useRenameWorkspaceFile() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: async ({ workspaceId, fileId, name }: RenameFileParams) =>
      requestJson(renameWorkspaceFileContract, {
        params: { id: workspaceId, fileId },
        body: { name },
      }),
    onMutate: async ({ workspaceId, fileId, name }) => {
      await queryClient.cancelQueries({ queryKey: workspaceFilesKeys.workspaceLists(workspaceId) })
      const previous = queryClient.getQueryData<WorkspaceFileRecord[]>(
        workspaceFilesKeys.list(workspaceId, 'active')
      )
      if (previous) {
        queryClient.setQueryData<WorkspaceFileRecord[]>(
          workspaceFilesKeys.list(workspaceId, 'active'),
          previous.map((f) => (f.id === fileId ? { ...f, name } : f))
        )
      }
      return { previous }
    },
    onError: (error, variables, context) => {
      if (context?.previous) {
        queryClient.setQueryData(
          workspaceFilesKeys.list(variables.workspaceId, 'active'),
          context.previous
        )
      }
      toast.error(error.message, { duration: 5000 })
    },
    onSettled: (_data, _error, variables) => {
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.workspaceLists(variables.workspaceId),
      })
    },
  })
}

/**
 * Delete workspace file mutation
 */
interface DeleteFileParams {
  workspaceId: string
  fileId: string
}

export function useDeleteWorkspaceFile() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: async ({ workspaceId, fileId }: DeleteFileParams) =>
      requestJson(deleteWorkspaceFileContract, {
        params: { id: workspaceId, fileId },
      }),
    onMutate: async ({ workspaceId, fileId }) => {
      await queryClient.cancelQueries({ queryKey: workspaceFilesKeys.workspaceLists(workspaceId) })

      const previousFiles = queryClient.getQueryData<WorkspaceFileRecord[]>(
        workspaceFilesKeys.list(workspaceId, 'active')
      )

      if (previousFiles) {
        queryClient.setQueryData<WorkspaceFileRecord[]>(
          workspaceFilesKeys.list(workspaceId, 'active'),
          previousFiles.filter((f) => f.id !== fileId)
        )
      }

      return { previousFiles }
    },
    onError: (_err, variables, context) => {
      if (context?.previousFiles) {
        queryClient.setQueryData(
          workspaceFilesKeys.list(variables.workspaceId, 'active'),
          context.previousFiles
        )
      }
      logger.error('Failed to delete file')
      toast.error(toError(_err).message)
    },
    onSuccess: () => {
      toast.success('File moved to trash')
    },
    onSettled: (_data, _error, variables) => {
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.workspaceLists(variables.workspaceId),
      })
      queryClient.removeQueries({
        queryKey: workspaceFilesKeys.contentFile(variables.workspaceId, variables.fileId),
      })
      queryClient.invalidateQueries({ queryKey: workspaceFilesKeys.storageInfo() })
    },
  })
}

export function useRestoreWorkspaceFile() {
  const queryClient = useQueryClient()

  return useMutation({
    mutationFn: async ({ workspaceId, fileId }: { workspaceId: string; fileId: string }) =>
      requestJson(restoreWorkspaceFileContract, {
        params: { id: workspaceId, fileId },
      }),
    onSuccess: () => {
      toast.success('File restored')
    },
    onError: (err) => {
      toast.error(toError(err).message)
    },
    onSettled: (_data, _error, variables) => {
      queryClient.invalidateQueries({
        queryKey: workspaceFilesKeys.workspaceLists(variables.workspaceId),
      })
      queryClient.invalidateQueries({ queryKey: workspaceFilesKeys.storageInfo() })
    },
  })
}