template-saas/apps/frontend/src/hooks/useAI.ts

import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query';
import toast from 'react-hot-toast';
import { aiApi, ChatRequest, ChatResponse, AIConfig, AIModel, AIUsageStats } from '@/services/api';
import { AxiosError } from 'axios';

interface ApiError {
  message: string;
  statusCode?: number;
}

// ==================== Query Keys ====================

export const aiQueryKeys = {
  all: ['ai'] as const,
  config: () => [...aiQueryKeys.all, 'config'] as const,
  models: () => [...aiQueryKeys.all, 'models'] as const,
  usage: (page?: number, limit?: number) => [...aiQueryKeys.all, 'usage', { page, limit }] as const,
  currentUsage: () => [...aiQueryKeys.all, 'current-usage'] as const,
  health: () => [...aiQueryKeys.all, 'health'] as const,
};

// ==================== Config Hooks ====================

export function useAIConfig() {
  return useQuery({
    queryKey: aiQueryKeys.config(),
    queryFn: () => aiApi.getConfig(),
    retry: 1,
    staleTime: 5 * 60 * 1000, // 5 minutes
  });
}

export function useUpdateAIConfig() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (data: Partial<AIConfig>) => aiApi.updateConfig(data),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: aiQueryKeys.config() });
      toast.success('AI configuration updated');
    },
    onError: (error: AxiosError<ApiError>) => {
      toast.error(error.response?.data?.message || 'Failed to update AI configuration');
    },
  });
}

// ==================== Models Hook ====================

export function useAIModels() {
  return useQuery({
    queryKey: aiQueryKeys.models(),
    queryFn: () => aiApi.getModels(),
    staleTime: 30 * 60 * 1000, // 30 minutes - models don't change often
  });
}

// ==================== Chat Hook ====================

export function useAIChat() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (data: ChatRequest) => aiApi.chat(data),
    onSuccess: () => {
      // Invalidate usage queries after successful chat
      queryClient.invalidateQueries({ queryKey: aiQueryKeys.currentUsage() });
      queryClient.invalidateQueries({ queryKey: aiQueryKeys.usage() });
    },
    onError: (error: AxiosError<ApiError>) => {
      const message = error.response?.data?.message || 'Failed to get AI response';
      toast.error(message);
    },
  });
}

// ==================== Usage Hooks ====================

export interface AIUsageRecord {
  id: string;
  model: string;
  input_tokens: number;
  output_tokens: number;
  total_tokens: number;
  cost: number;
  latency_ms: number;
  status: string;
  created_at: string;
}

export interface AIUsageResponse {
  data: AIUsageRecord[];
  total: number;
  page: number;
  limit: number;
  totalPages: number;
}

export function useAIUsage(page = 1, limit = 10) {
  return useQuery({
    queryKey: aiQueryKeys.usage(page, limit),
    queryFn: () => aiApi.getUsage({ page, limit }) as Promise<AIUsageResponse>,
  });
}

export function useCurrentAIUsage() {
  return useQuery({
    queryKey: aiQueryKeys.currentUsage(),
    queryFn: () => aiApi.getCurrentUsage(),
    refetchInterval: 60000, // Refetch every minute
  });
}

// ==================== Health Hook ====================

export interface AIHealthStatus {
  status: 'healthy' | 'degraded' | 'unhealthy';
  provider: string;
  latency_ms: number;
  models_available: number;
}

export function useAIHealth() {
  return useQuery({
    queryKey: aiQueryKeys.health(),
    queryFn: () => aiApi.getHealth() as Promise<AIHealthStatus>,
    refetchInterval: 5 * 60 * 1000, // Check every 5 minutes
    retry: 1,
  });
}

// ==================== Re-export types ====================

export type { ChatRequest, ChatResponse, AIConfig, AIModel, AIUsageStats };