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

import { useState, useCallback, useMemo, useEffect } from 'react';
import { useMessage } from '@/hooks/useSnackbar';
import userService from '@/services/user_service';
import logger from '@/utils/logger';
import type {
  ApiKeyFormData,
  AzureOpenAIFormData,
  BedrockFormData,
  OllamaFormData,
  // SystemModelFormData,
} from '../components/ModelDialogs';
import type { ITenantInfo } from '@/interfaces/database/knowledge';
import { useLlmList } from '@/hooks/llm-hooks';
import type { LlmModelType } from '@/constants/knowledge';
import { useUserData } from '@/hooks/useUserData';
import type { ISetApiKeyRequestBody, IAddLlmRequestBody } from '@/interfaces/request/llm';
import type { ConfigFormItem, ConfigurationFormData, DocLinkConfig } from '../components/Dialog/ConfigurationDialog';
import { getLLMConfig } from '../components/Dialog/llmConfigs';

// 对话框状态管理 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) => {
    if (data != null) {
      setInitialData(data);
    }
    setEditMode(isEdit);
    setOpen(true);
  }, []);

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

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

// 通用配置对话框管理
export const useConfigurationDialog = (onSuccess?: () => void) => {
  const dialogState = useDialogState();
  const showMessage = useMessage();
  const [llmFactory, setLlmFactory] = useState('');
  const [config, setConfig] = useState<{
    formItems: ConfigFormItem[];
    docLink?: DocLinkConfig;
    title: string;
    defaultValues: Record<string, any>;
  } | null>(null);

  const openConfigurationDialog = useCallback((factory: string, data?: Partial<ConfigurationFormData>, isEdit = false) => {
    setLlmFactory(factory);
    const llmConfig = getLLMConfig(factory);
    setConfig(llmConfig);
    
    // 合并默认值和传入的数据
    const mergedData = {
      ...llmConfig.defaultValues,
      ...data,
    };
    
    dialogState.openDialog(mergedData, isEdit);
  }, [dialogState]);

  const submitConfiguration = useCallback(async (data: Partial<ConfigurationFormData>) => {
    dialogState.setLoading(true);
    logger.info('提交配置:', data);
    try {
      // 构建请求参数
      const params: Partial<IAddLlmRequestBody> = {
        ...data,
        llm_factory: llmFactory,
      };

      await userService.add_llm(params);
      showMessage.success(`${config?.title || llmFactory} 配置成功`);
      dialogState.closeDialog();
      
      // 调用成功回调
      if (onSuccess) {
        onSuccess();
      }
    } catch (error) {
      logger.error(`${config?.title || llmFactory} 配置失败:`, error);
      showMessage.error(`${config?.title || llmFactory} 配置失败`);
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [llmFactory, config, dialogState, showMessage, onSuccess]);

  return {
    ...dialogState,
    llmFactory,
    config,
    openConfigurationDialog,
    submitConfiguration,
  };
};

// API Key 对话框管理
export const useApiKeyDialog = (onSuccess?: () => void) => {
  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);
    logger.info('提交 API Key:', data);
    try {
      const params: ISetApiKeyRequestBody = {
        llm_factory: factoryName,
        api_key: data.api_key,
      };

      if (data.base_url && data.base_url.trim() !== '') {
        params.base_url = data.base_url;
      }

      if (data.group_id && data.group_id.trim() !== '') {
        // params.group_id = data.group_id;
      }

      await userService.set_api_key(params);
      showMessage.success('API Key 配置成功');
      dialogState.closeDialog();
      
      // 调用成功回调
      if (onSuccess) {
        onSuccess();
      }
    } catch (error) {
      logger.error('API Key 配置失败:', error);
    } finally {
      dialogState.setLoading(false);
    }
  }, [factoryName, dialogState, onSuccess]);

  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.add_llm({
        llm_factory: data.llm_factory,
        llm_name: data.llm_name,
        model_type: data.model_type,
        api_base: data.api_base,
        api_key: data.api_key,
        // @ts-ignore
        api_version: data.api_version,
        max_tokens: data.max_tokens,
      });
      showMessage.success('Azure OpenAI 配置成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('Azure OpenAI 配置失败:', error);
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState, showMessage]);

  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.add_llm({
        llm_factory: data.llm_factory,
        llm_name: data.llm_name,
        model_type: data.model_type,
        // @ts-ignore
        bedrock_ak: data.bedrock_ak,
        bedrock_sk: data.bedrock_sk,
        bedrock_region: data.bedrock_region,
        max_tokens: data.max_tokens,
      });
      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({
        llm_factory: data.llm_factory,
        llm_name: data.llm_name,
        model_type: data.model_type,
        api_base: data.api_base,
        api_key: data.api_key || '',
        max_tokens: data.max_tokens,
      });
      showMessage.success('Ollama 模型添加成功');
      dialogState.closeDialog();
    } catch (error) {
      logger.error('Ollama 模型添加失败:', error);
      showMessage.error('Ollama 模型添加失败');
      throw error;
    } finally {
      dialogState.setLoading(false);
    }
  }, [dialogState, showMessage]);

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

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

  const deleteLlm = useCallback(async (factoryName: string, modelName: string) => {
    setLoading(true);
    try {
      await userService.delete_llm({
        llm_factory: factoryName,
        llm_name: modelName,
      });
      showMessage.success('模型删除成功');
      
      // 调用成功回调
      if (onSuccess) {
        onSuccess();
      }
    } catch (error) {
      logger.error('模型删除失败:', error);
    } finally {
      setLoading(false);
    }
  }, [onSuccess]);

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

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

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

  const { data: llmList } = useLlmList();

  const { tenantInfo, fetchTenantInfo } = useUserData();

  useEffect(() => {
    fetchTenantInfo();
  }, []);

  const getOptionsByModelType = useCallback((modelType: LlmModelType) => {
    return Object.entries(llmList)
      .filter(([, value]) =>
        modelType
          ? value.some((x) => x.model_type.includes(modelType))
          : true,
      )
      .map(([key, value]) => {
        return {
          label: key,
          options: value
            .filter(
              (x) =>
                (modelType ? x.model_type.includes(modelType) : true) &&
                x.available,
            )
            .map((x) => ({
              label: x.llm_name,
              value: `${x.llm_name}@${x.fid}`,
              disabled: !x.available,
              model: x,
            })),
        };
      })
      .filter((x) => x.options.length > 0);
  }, [llmList]);

  const allModelOptions = useMemo(() => {
    const llmOptions = getOptionsByModelType('chat');
    const image2textOptions = getOptionsByModelType('image2text');
    const embeddingOptions = getOptionsByModelType('embedding');
    const speech2textOptions = getOptionsByModelType('speech2text');
    const rerankOptions = getOptionsByModelType('rerank');
    const ttsOptions = getOptionsByModelType('tts');

    return {
      llmOptions,
      image2textOptions,
      embeddingOptions,
      speech2textOptions,
      rerankOptions,
      ttsOptions,
    }
  }, [llmList, getOptionsByModelType]);

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

  return {
    ...dialogState,
    submitSystemModelSetting,
    allModelOptions,
    initialData: tenantInfo,
  };
};

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

  // 根据工厂类型打开对应的对话框
  const openFactoryDialog = useCallback((factoryName: string, data?: any, isEdit = false) => {
    // 使用通用的 ConfigurationDialog 替代特定的 Dialog
    configurationDialog.openConfigurationDialog(factoryName, data, isEdit);
  }, [configurationDialog]);

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