YG-Datasets/easy-dataset-main/app/projects/[projectId]/questions/hooks/useQuestionGeneration.js

'use client';

import { useState } from 'react';
import { useTranslation } from 'react-i18next';
import { toast } from 'sonner';
import axios from 'axios';
import i18n from '@/lib/i18n';
import request from '@/lib/util/request';
import { processInParallel } from '@/lib/util/processInParallel';

export function useQuestionGeneration(projectId, model, taskSettings, getQuestionList) {
  const { t } = useTranslation();

  // 处理状态
  const [processing, setProcessing] = useState(false);

  // 进度状态
  const [progress, setProgress] = useState({
    total: 0, // 总共选择的问题数量
    completed: 0, // 已处理完成的数量
    percentage: 0, // 进度百分比
    datasetCount: 0 // 已生成的数据集数量
  });

  // 批量生成答案
  const handleBatchGenerateAnswers = async selectedQuestions => {
    if (selectedQuestions.length === 0) {
      toast.warning(t('questions.noQuestionsSelected'));
      return;
    }

    if (!model) {
      toast.warning(t('models.configNotFound'));
      return;
    }

    try {
      setProgress({
        total: selectedQuestions.length,
        completed: 0,
        percentage: 0,
        datasetCount: 0
      });

      // 然后设置处理状态为真，确保进度条显示
      setProcessing(true);

      toast.info(t('questions.batchGenerateStart', { count: selectedQuestions.length }));

      // 单个问题处理函数
      const processQuestion = async questionId => {
        try {
          console.log('开始生成数据集:', { questionId });
          const language = i18n.language === 'zh-CN' ? '中文' : 'en';
          // 调用API生成数据集
          const response = await request(`/api/projects/${projectId}/datasets`, {
            method: 'POST',
            headers: {
              'Content-Type': 'application/json'
            },
            body: JSON.stringify({
              questionId,
              model,
              language
            })
          });

          if (!response.ok) {
            const errorData = await response.json();
            console.error(t('datasets.generateError'), errorData.error || t('datasets.generateFailed'));

            // 更新进度状态（即使失败也计入已处理）
            setProgress(prev => {
              const completed = prev.completed + 1;
              const percentage = Math.round((completed / prev.total) * 100);

              return {
                ...prev,
                completed,
                percentage
              };
            });

            return { success: false, questionId, error: errorData.error || t('datasets.generateFailed') };
          }

          const data = await response.json();

          // 更新进度状态
          setProgress(prev => {
            const completed = prev.completed + 1;
            const percentage = Math.round((completed / prev.total) * 100);
            const datasetCount = prev.datasetCount + 1;

            return {
              ...prev,
              completed,
              percentage,
              datasetCount
            };
          });

          console.log(`数据集生成成功: ${questionId}`);
          return { success: true, questionId, data: data.dataset };
        } catch (error) {
          console.error('生成数据集失败:', error);

          // 更新进度状态（即使失败也计入已处理）
          setProgress(prev => {
            const completed = prev.completed + 1;
            const percentage = Math.round((completed / prev.total) * 100);

            return {
              ...prev,
              completed,
              percentage
            };
          });

          return { success: false, questionId, error: error.message };
        }
      };

      // 并行处理所有问题，最多同时处理2个
      const results = await processInParallel(selectedQuestions, processQuestion, taskSettings.concurrencyLimit);

      // 刷新数据
      getQuestionList();

      // 处理完成后设置结果消息
      const successCount = results.filter(r => r.success).length;
      const failCount = results.filter(r => !r.success).length;

      if (failCount > 0) {
        toast.warning(
          t('datasets.partialSuccess', {
            successCount,
            total: selectedQuestions.length,
            failCount
          })
        );
      } else {
        toast.success(t('common.success', { successCount }));
      }
    } catch (error) {
      console.error('生成数据集出错:', error);
      toast.error(error.message || '生成数据集失败');
    } finally {
      // 延迟关闭处理状态，确保用户可以看到完成的进度
      setTimeout(() => {
        setProcessing(false);
        // 再次延迟重置进度状态
        setTimeout(() => {
          setProgress({
            total: 0,
            completed: 0,
            percentage: 0,
            datasetCount: 0
          });
        }, 500);
      }, 2000); // 延迟关闭处理状态，让用户看到完成的进度
    }
  };

  // 自动生成数据集
  const handleAutoGenerateDatasets = async () => {
    try {
      if (!model) {
        toast.error(t('questions.selectModelFirst', { defaultValue: '请先选择模型' }));
        return;
      }

      // 调用创建任务接口
      const response = await axios.post(`/api/projects/${projectId}/tasks`, {
        taskType: 'answer-generation',
        modelInfo: model,
        language: i18n.language
      });

      if (response.data?.code === 0) {
        toast.success(t('tasks.createSuccess', { defaultValue: '后台任务已创建，系统将自动处理未生成答案的问题' }));
      } else {
        toast.error(t('tasks.createFailed', { defaultValue: '创建后台任务失败' }));
      }
    } catch (error) {
      console.error('创建任务失败:', error);
      toast.error(t('tasks.createFailed', { defaultValue: '创建任务失败' }) + ': ' + error.message);
    }
  };

  // 自动生成图像问答数据集
  const handleAutoGenerateImageDatasets = async () => {
    try {
      if (!model) {
        toast.error(t('questions.selectModelFirst', { defaultValue: '请先选择模型' }));
        return;
      }

      if (model.type !== 'vision') {
        toast.error(t('images.visionModelRequired', { defaultValue: '请选择支持视觉的模型' }));
        return;
      }

      // 调用创建任务接口
      const response = await axios.post(`/api/projects/${projectId}/tasks`, {
        taskType: 'image-dataset-generation',
        modelInfo: model,
        language: i18n.language
      });

      if (response.data?.code === 0) {
        toast.success(t('tasks.createSuccess', { defaultValue: '后台任务已创建，系统将自动处理未生成答案的图片问题' }));
      } else {
        toast.error(t('tasks.createFailed', { defaultValue: '创建后台任务失败' }));
      }
    } catch (error) {
      console.error('创建图片数据集任务失败:', error);
      toast.error(t('tasks.createFailed', { defaultValue: '创建任务失败' }) + ': ' + error.message);
    }
  };

  // 自动生成多轮对话数据集
  const handleAutoGenerateMultiTurnDatasets = async () => {
    try {
      if (!model) {
        toast.error(t('questions.selectModelFirst', { defaultValue: '请先选择模型' }));
        return;
      }

      // 首先检查项目是否配置了多轮对话设置
      const configResponse = await axios.get(`/api/projects/${projectId}/tasks`);
      if (configResponse.status !== 200) {
        throw new Error('获取项目配置失败');
      }

      const config = configResponse.data;
      const multiTurnConfig = {
        systemPrompt: config.multiTurnSystemPrompt,
        scenario: config.multiTurnScenario,
        rounds: config.multiTurnRounds,
        roleA: config.multiTurnRoleA,
        roleB: config.multiTurnRoleB
      };

      // 检查是否已配置必要的多轮对话设置
      if (
        !multiTurnConfig.scenario ||
        !multiTurnConfig.roleA ||
        !multiTurnConfig.roleB ||
        !multiTurnConfig.rounds ||
        multiTurnConfig.rounds < 1
      ) {
        toast.error(t('questions.multiTurnNotConfigured', '请先在项目设置中配置多轮对话相关参数'));
        return;
      }

      // 调用创建任务接口
      const response = await axios.post(`/api/projects/${projectId}/tasks`, {
        taskType: 'multi-turn-generation',
        modelInfo: model,
        language: i18n.language,
        config: JSON.stringify(multiTurnConfig)
      });

      if (response.data?.code === 0) {
        toast.success(
          t('tasks.multiTurnCreateSuccess', {
            defaultValue: '多轮对话生成任务已创建，系统将自动处理未生成多轮对话的问题'
          })
        );
      } else {
        toast.error(t('tasks.createFailed', { defaultValue: '创建后台任务失败' }));
      }
    } catch (error) {
      console.error('创建多轮对话任务失败:', error);
      toast.error(t('tasks.createFailed', { defaultValue: '创建任务失败' }) + ': ' + error.message);
    }
  };

  return {
    // 状态
    processing,
    progress,

    // 方法
    handleBatchGenerateAnswers,
    handleAutoGenerateDatasets,
    handleAutoGenerateMultiTurnDatasets,
    handleAutoGenerateImageDatasets
  };
}