bookwiz.io / lib / hooks / useImageGeneration.ts
useImageGeneration.ts
Raw
import { useState, useCallback, useEffect } from 'react'
import { supabase } from '@/lib/supabase'

// Helper function to get auth headers for API requests
const getAuthHeaders = async () => {
  const { data: { session } } = await supabase.auth.getSession()
  const headers: Record<string, string> = { 'Content-Type': 'application/json' }
  
  if (session?.access_token) {
    headers['Authorization'] = `Bearer ${session.access_token}`
  }
  
  return headers
}

export interface GeneratedImage {
  id: string
  prompt: string
  revised_prompt?: string
  image_url: string
  size: string
  quality: string
  style: string
  created_at: string
  book_id?: string
}

export interface ImageGenerationSettings {
  size: '1024x1024' | '1536x1024' | '1024x1536'
  quality: 'low' | 'medium' | 'high' | 'auto'
}

export interface UseImageGenerationReturn {
  // State
  images: GeneratedImage[]
  isGenerating: boolean
  isLoading: boolean
  error: string | null
  settings: ImageGenerationSettings
  
  // Actions
  generateImage: (prompt: string, bookId?: string) => Promise<GeneratedImage | null>
  refreshImages: () => Promise<void>
  deleteImage: (imageId: string) => Promise<void>
  updateSettings: (newSettings: Partial<ImageGenerationSettings>) => void
  clearError: () => void
}

export function useImageGeneration(): UseImageGenerationReturn {
  const [images, setImages] = useState<GeneratedImage[]>([])
  const [isGenerating, setIsGenerating] = useState(false)
  const [isLoading, setIsLoading] = useState(false)
  const [error, setError] = useState<string | null>(null)
  const [settings, setSettings] = useState<ImageGenerationSettings>({
    size: '1024x1536',
    quality: 'auto'
  })

  // Fetch user's generated images
  const refreshImages = useCallback(async () => {
    setIsLoading(true)
    setError(null)
    
    try {
      const headers = await getAuthHeaders()
      const response = await fetch('/api/images', { headers })
      
      if (!response.ok) {
        setError('Failed to fetch images')
        return
      }
      
      const data = await response.json()
      setImages(data.images || [])
    } catch (err) {
      setError('Failed to fetch images')
    }
  }, [])

  // Automatically manage loading state based on images data
  useEffect(() => {
    if (images.length > 0 || error) {
      setIsLoading(false)
    }
  }, [images, error])

  // Generate a new image
  const generateImage = useCallback(async (prompt: string, bookId?: string): Promise<GeneratedImage | null> => {
    if (!prompt.trim()) {
      setError('Please enter a prompt')
      return null
    }

    setIsGenerating(true)
    setError(null)

    try {
      const headers = await getAuthHeaders()
      const response = await fetch('/api/images/generate', {
        method: 'POST',
        headers,
        body: JSON.stringify({
          prompt: prompt.trim(),
          bookId,
          ...settings
        })
      })

      if (!response.ok) {
        const errorData = await response.json()
        throw new Error(errorData.error || 'Failed to generate image')
      }

      const newImage = await response.json()
      
      // Add the new image to the beginning of the list
      setImages(prev => [newImage, ...prev])
      
      // Reset generating state after successful generation
      setIsGenerating(false)
      
      return newImage
    } catch (err) {
      setError('Failed to generate image')
      setIsGenerating(false)
      return null
    }
  }, [settings])

  // Delete an image
  const deleteImage = useCallback(async (imageId: string) => {
    try {
      const headers = await getAuthHeaders()
      const response = await fetch('/api/images', {
        method: 'DELETE',
        headers,
        body: JSON.stringify({ imageId })
      })

      if (!response.ok) {
        setError('Failed to delete image')
      }

      // Remove the image from the list
      setImages(prev => prev.filter(img => img.id !== imageId))
    } catch (err) {
      setError('Failed to delete image')
    }
  }, [])

  // Update generation settings
  const updateSettings = useCallback((newSettings: Partial<ImageGenerationSettings>) => {
    setSettings(prev => ({ ...prev, ...newSettings }))
  }, [])

  // Clear error
  const clearError = useCallback(() => {
    setError(null)
  }, [])

  // Load images on mount
  useEffect(() => {
    refreshImages()
  }, [refreshImages])

  return {
    images,
    isGenerating,
    isLoading,
    error,
    settings,
    generateImage,
    refreshImages,
    deleteImage,
    updateSettings,
    clearError
  }
}