TERES_web_frontend/src/pages/setting/hooks/useModelDialogs.ts

import { useState, useCallback } from 'react';
import { useMessage } from '@/hooks/useSnackbar';
import userService from '@/services/user_service';
import logger from '@/utils/logger';
import type { 
  ApiKeyFormData,
  AzureOpenAIFormData,
  BedrockFormData,
  OllamaFormData,
} from '../components/ModelDialogs';

// 对话框状态管理 hook
export const useDialogState = () => {
  const [open, setOpen] = useState(false);
  const [loading, setLoading] = useState(false);
  const [editMode, setEditMode] = useState(false);
  const [initialData, setInitialData] = useState<any>(null);

  const openDialog = useCallback((data?: any, isEdit = false) => {
    setInitialData(data);
    setEditMode(isEdit);
    setOpen(true);
  }, []);

  const closeDialog = useCallback(() => {
    setOpen(false);
    setInitialData(null);
    setEditMode(false);
  }, []);

  return {
    open,
    loading,
    editMode,
    initialData,
    setLoading,
    openDialog,
    closeDialog,
  };
};

// API Key 对话框管理
export const useApiKeyDialog = () => {
  const dialogState = useDialogState();
  const showMessage  = useMessage();
  const [factoryName, setFactoryName] = useState('');

  const openApiKeyDialog = useCallback((factory: string, data?: Partial<ApiKeyFormData>, isEdit = false) => {
    setFactoryName(factory);
    dialogState.openDialog(data, isEdit);
  }, [dialogState]);

  const submitApiKey = useCallback(async (data: ApiKeyFormData) => {
    dialogState.setLoading(true);
    try {
      await userService.set_api_key({
        factory_name: factoryName,
        model_name: '', // 根据实际需求调整
        // api_key: data.api_key,
        ...data
      });
      showMessage.success('API Key 配置成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('API Key 配置失败:', error);
      showMessage.error('API Key 配置失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [factoryName, dialogState]);

  return {
    ...dialogState,
    factoryName,
    openApiKeyDialog,
    submitApiKey,
  };
};

// Azure OpenAI 对话框管理
export const useAzureOpenAIDialog = () => {
  const dialogState = useDialogState();
  const showMessage  = useMessage();

  const submitAzureOpenAI = useCallback(async (data: AzureOpenAIFormData) => {
    dialogState.setLoading(true);
    try {
      // 调用 Azure OpenAI 特定的 API
      await userService.set_api_key({
        factory_name: 'AzureOpenAI',
        model_name: data.deployment_name,
        api_key: data.api_key,
        // azure_endpoint: data.azure_endpoint,
        // api_version: data.api_version,
      });
      showMessage.success('Azure OpenAI 配置成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('Azure OpenAI 配置失败:', error);
      showMessage.error('Azure OpenAI 配置失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState]);

  return {
    ...dialogState,
    submitAzureOpenAI,
  };
};

// AWS Bedrock 对话框管理
export const useBedrockDialog = () => {
  const dialogState = useDialogState();
  const showMessage  = useMessage();

  const submitBedrock = useCallback(async (data: BedrockFormData) => {
    dialogState.setLoading(true);
    try {
      // 调用 Bedrock 特定的 API
      await userService.set_api_key({
        factory_name: 'Bedrock',
        model_name: '',
        api_key: '', // Bedrock 使用 access key
        // access_key_id: data.access_key_id,
        // secret_access_key: data.secret_access_key,
        // region: data.region,
      });
      showMessage.success('AWS Bedrock 配置成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('AWS Bedrock 配置失败:', error);
      showMessage.error('AWS Bedrock 配置失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState]);

  return {
    ...dialogState,
    submitBedrock,
  };
};

// Ollama 对话框管理
export const useOllamaDialog = () => {
  const dialogState = useDialogState();
  const showMessage  = useMessage();

  const submitOllama = useCallback(async (data: OllamaFormData) => {
    dialogState.setLoading(true);
    try {
      // 调用添加 LLM 的 API
      await userService.add_llm({
        factory_name: 'Ollama',
        model_name: data.model_name,
        // base_url: data.base_url,
      });
      showMessage.success('Ollama 模型添加成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('Ollama 模型添加失败:', error);
      showMessage.error('Ollama 模型添加失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState]);

  return {
    ...dialogState,
    submitOllama,
  };
};

// 删除操作管理
export const useDeleteOperations = () => {
  const showMessage  = useMessage();
  const [loading, setLoading] = useState(false);

  const deleteLlm = useCallback(async (factoryName: string, modelName: string) => {
    setLoading(true);
    try {
      await userService.delete_llm({
        factory_name: factoryName,
        model_name: modelName,
      });
      showMessage.success('模型删除成功');
    } catch (error) {
      logger.error('模型删除失败:', error);
      showMessage.error('模型删除失败');
      throw error;
    } finally {
      setLoading(false);
    }
  }, []);

  const deleteFactory = useCallback(async (factoryName: string) => {
    setLoading(true);
    try {
      await userService.deleteFactory({
        factory_name: factoryName,
      });
      showMessage.success('模型工厂删除成功');
    } catch (error) {
      logger.error('模型工厂删除失败:', error);
      showMessage.error('模型工厂删除失败');
      throw error;
    } finally {
      setLoading(false);
    }
  }, []);

  return {
    loading,
    deleteLlm,
    deleteFactory,
  };
};

// 系统默认模型设置
export const useSystemModelSetting = () => {
  const dialogState = useDialogState();
  const showMessage  = useMessage();

  const submitSystemModelSetting = useCallback(async (data: { defaultModel: string }) => {
    dialogState.setLoading(true);
    try {
      // 这里需要根据实际的 API 接口调整
      // await userService.setSystemDefaultModel(data);
      showMessage.success('系统默认模型设置成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('系统默认模型设置失败:', error);
      showMessage.error('系统默认模型设置失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState]);

  return {
    ...dialogState,
    submitSystemModelSetting,
  };
};

// 统一的模型对话框管理器
export const useModelDialogs = () => {
  const apiKeyDialog = useApiKeyDialog();
  const azureDialog = useAzureOpenAIDialog();
  const bedrockDialog = useBedrockDialog();
  const ollamaDialog = useOllamaDialog();
  const systemDialog = useSystemModelSetting();
  const deleteOps = useDeleteOperations();

  // 根据工厂类型打开对应的对话框
  const openFactoryDialog = useCallback((factoryName: string, data?: any, isEdit = false) => {
    switch (factoryName.toLowerCase()) {
      case 'azureopenai':
        azureDialog.openDialog(data, isEdit);
        break;
      case 'bedrock':
        bedrockDialog.openDialog(data, isEdit);
        break;
      case 'ollama':
        ollamaDialog.openDialog(data, isEdit);
        break;
      default:
        // 默认使用 API Key 对话框
        apiKeyDialog.openApiKeyDialog(factoryName, data, isEdit);
        break;
    }
  }, [apiKeyDialog, azureDialog, bedrockDialog, ollamaDialog]);

  return {
    apiKeyDialog,
    azureDialog,
    bedrockDialog,
    ollamaDialog,
    systemDialog,
    deleteOps,
    openFactoryDialog,
  };
};