bookwiz.io / lib / hooks / useUsageLimit.ts
useUsageLimit.ts
Raw
import { useState, useCallback, useEffect } from 'react'
import { getModelTier } from '@/lib/config/models'

export interface UsageInfo {
  smartUsage: number
  fastUsage: number
  smartLimit: number
  fastLimit: number
}

export interface UseUsageLimitReturn {
  // Usage states
  usageInfo: UsageInfo | null
  usageLimitError: string | null

  // Functions
  loadUsageInfo: () => Promise<void>
  checkUsageLimit: (modelName: string) => Promise<boolean>
  setUsageLimitError: (error: string | null) => void
  refreshUsage: () => Promise<void>
}

export function useUsageLimit(
  userId: string | null,
  refreshUsageStats?: () => void
): UseUsageLimitReturn {
  // Usage tracking state
  const [usageInfo, setUsageInfo] = useState<UsageInfo | null>(null)
  const [usageLimitError, setUsageLimitError] = useState<string | null>(null)

  const loadUsageInfo = useCallback(async () => {
    if (!userId) return
    
    try {
      const response = await fetch('/api/usage/check', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ userId })
      })
      
      if (response.ok) {
        const data = await response.json()
        setUsageInfo(data)
      }
    } catch (error) {
      console.error('Failed to load usage info:', error)
    }
  }, [userId])

  // Check if user can make request for the selected model
  const checkUsageLimit = useCallback(async (modelName: string): Promise<boolean> => {
    if (!userId) return true // Allow if no user (shouldn't happen)
    
    const modelTier = getModelTier(modelName)
    
    try {
      const response = await fetch('/api/usage/check-limit', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ 
          userId, 
          tier: modelTier
        })
      })
      
      if (!response.ok) {
        setUsageLimitError('Failed to check usage limits')
        return false
      }
      
      const data = await response.json()
      
      if (!data.can_proceed) {
        const tierDisplay = modelTier === 'smart' ? '🧠 Smart AI' : '⚡ Fast AI'
        
        // Calculate when the limit resets (first day of next month)
        const nextReset = new Date()
        nextReset.setMonth(nextReset.getMonth() + 1)
        nextReset.setDate(1)
        nextReset.setHours(0, 0, 0, 0)
        
        const daysUntilReset = Math.ceil((nextReset.getTime() - Date.now()) / (1000 * 60 * 60 * 24))
        
        // Enhanced error message with all details
        const enhancedMessage = 
          `🚫 You've reached your ${tierDisplay} limit!\n\n` +
          `Usage: ${data.current_usage}/${data.limit === 999999 ? '∞' : data.limit} requests this month\n` +
          `Resets in ${daysUntilReset} day${daysUntilReset !== 1 ? 's' : ''} (${nextReset.toLocaleDateString()})\n\n` +
          `💡 Upgrade your plan for more ${tierDisplay.toLowerCase()} requests or try a different model.`
        
        setUsageLimitError(enhancedMessage)
        return false
      }
      
      // Update usage info
      await loadUsageInfo()
      setUsageLimitError(null)
      return true
    } catch (error) {
      console.error('Failed to check usage limit:', error)
      // Allow request to proceed if check fails (fail open)
      return true
    }
  }, [userId, loadUsageInfo])

  // Helper function to refresh usage stats and reload usage info
  const refreshUsage = useCallback(async () => {
    try {
      // Refresh dashboard usage stats
      if (refreshUsageStats) {
        refreshUsageStats()
      }
      // Also refresh local usage info
      await loadUsageInfo()
    } catch (error) {
      console.error('Failed to refresh usage:', error)
    }
  }, [refreshUsageStats, loadUsageInfo])

  // Load usage info when userId changes
  useEffect(() => {
    if (userId) {
      loadUsageInfo()
    }
  }, [userId, loadUsageInfo])

  return {
    // States
    usageInfo,
    usageLimitError,

    // Functions
    loadUsageInfo,
    checkUsageLimit,
    setUsageLimitError,
    refreshUsage
  }
}