/**
 * Phase 3: 批处理模式 - 状态管理Hook
 */

import { useState, useCallback } from 'react';
import { message } from 'antd';
import * as batchApi from '../api/batchApi';
import type { BatchTemplate, BatchTask, BatchResult } from '../api/batchApi';

export interface BatchTaskState {
  // 模板选择
  templates: BatchTemplate[];
  selectedTemplate: BatchTemplate | null;
  customPrompt: string;
  
  // 文献选择
  selectedDocuments: string[];
  
  // 任务执行
  currentTask: BatchTask | null;
  taskResults: BatchResult[];
  isExecuting: boolean;
  
  // UI状态
  step: 'define' | 'executing' | 'results'; // 当前步骤
}

export function useBatchTask() {
  const [state, setState] = useState<BatchTaskState>({
    templates: [],
    selectedTemplate: null,
    customPrompt: '',
    selectedDocuments: [],
    currentTask: null,
    taskResults: [],
    isExecuting: false,
    step: 'define',
  });

  // 加载模板列表
  const loadTemplates = useCallback(async () => {
    try {
      const templates = await batchApi.getTemplates();
      setState(prev => ({ ...prev, templates }));
      
      // 默认选择第一个模板
      if (templates.length > 0) {
        setState(prev => ({ ...prev, selectedTemplate: templates[0] }));
      }
    } catch (error: any) {
      message.error(`加载模板失败: ${error.message}`);
    }
  }, []);

  // 选择模板
  const selectTemplate = useCallback((template: BatchTemplate | null) => {
    setState(prev => ({ ...prev, selectedTemplate: template }));
  }, []);

  // 设置自定义提示词
  const setCustomPrompt = useCallback((prompt: string) => {
    setState(prev => ({ ...prev, customPrompt: prompt }));
  }, []);

  // 选择文献
  const setSelectedDocuments = useCallback((documentIds: string[]) => {
    setState(prev => ({ ...prev, selectedDocuments: documentIds }));
  }, []);

  // 执行批处理任务
  const executeBatch = useCallback(async (params: {
    kbId: string;
    modelType: string;
    taskName?: string;
  }) => {
    const { kbId, modelType, taskName } = params;

    // 验证
    if (state.selectedDocuments.length < 3) {
      message.error('请至少选择3篇文献');
      return;
    }

    if (state.selectedDocuments.length > 50) {
      message.error('最多只能选择50篇文献');
      return;
    }

    const isPreset = !!state.selectedTemplate;
    if (!isPreset && !state.customPrompt.trim()) {
      message.error('请输入自定义提示词');
      return;
    }

    setState(prev => ({ ...prev, isExecuting: true, step: 'executing' }));

    try {
      // 执行批处理
      const result = await batchApi.executeBatch({
        kb_id: kbId,
        document_ids: state.selectedDocuments,
        template_type: isPreset ? 'preset' : 'custom',
        template_id: state.selectedTemplate?.id,
        custom_prompt: state.customPrompt || undefined,
        model_type: modelType,
        task_name: taskName,
      });

      message.success('批处理任务已开始');

      // 轮询任务状态
      const task = await batchApi.pollTaskUntilComplete(
        result.task_id,
        (task) => {
          setState(prev => ({ ...prev, currentTask: task }));
        }
      );

      // 获取结果
      const { task: finalTask, results } = await batchApi.getTaskResults(result.task_id);

      setState(prev => ({
        ...prev,
        currentTask: finalTask,
        taskResults: results,
        isExecuting: false,
        step: 'results',
      }));

      if (finalTask.failed_count > 0) {
        message.warning(`任务完成，${finalTask.failed_count}个文档处理失败`);
      } else {
        message.success('批处理任务完成！');
      }
    } catch (error: any) {
      message.error(`批处理任务失败: ${error.message}`);
      setState(prev => ({ ...prev, isExecuting: false, step: 'define' }));
    }
  }, [state.selectedDocuments, state.selectedTemplate, state.customPrompt]);

  // 重试失败的文档
  const retryFailed = useCallback(async () => {
    if (!state.currentTask) return;

    try {
      await batchApi.retryFailed(state.currentTask.id);
      message.success('已开始重试失败的文档');

      // 重新轮询状态
      const task = await batchApi.pollTaskUntilComplete(
        state.currentTask.id,
        (task) => {
          setState(prev => ({ ...prev, currentTask: task }));
        }
      );

      // 获取新结果
      const { task: finalTask, results } = await batchApi.getTaskResults(state.currentTask.id);

      setState(prev => ({
        ...prev,
        currentTask: finalTask,
        taskResults: results,
      }));

      message.success('重试完成');
    } catch (error: any) {
      message.error(`重试失败: ${error.message}`);
    }
  }, [state.currentTask]);

  // 重置状态
  const reset = useCallback(() => {
    setState({
      templates: state.templates, // 保留模板列表
      selectedTemplate: state.templates[0] || null,
      customPrompt: '',
      selectedDocuments: [],
      currentTask: null,
      taskResults: [],
      isExecuting: false,
      step: 'define',
    });
  }, [state.templates]);

  return {
    state,
    loadTemplates,
    selectTemplate,
    setCustomPrompt,
    setSelectedDocuments,
    executeBatch,
    retryFailed,
    reset,
  };
}

