leaudit-platform-frontend/app/hooks/dify-dataset-manager/document-detail.ts

import { message } from 'antd';
import { useCallback, useEffect, useRef, useState } from 'react';
import {
    fetchUploadFileInfo,
    downloadOriginalFile,
    updateDocumentByFile,
    fetchIndexingStatus,
} from '~/api/dify-dataset/api/documentApi';
import { fetchSegments } from '~/api/dify-dataset/api/segmentApi';
import type { Segment, IndexingStatus } from '~/api/dify-dataset/type';
import type { Document } from '~/api/dify-dataset/type/documentTypes';
import type { DocumentDetailSegmentationSettings } from '~/types/dify-dataset-manager/document-detail';
import { DEFAULT_DOCUMENT_DETAIL_SETTINGS } from '~/types/dify-dataset-manager/document-detail';

/**
 * 文档详情状态管理 Hook
 */
export function useDocumentDetail(datasetId: string, document: Document | null) {
    // 分段设置状态
    const [settings, setSettings] = useState<DocumentDetailSegmentationSettings>(DEFAULT_DOCUMENT_DETAIL_SETTINGS);

    // 预览状态
    const [previewSegments, setPreviewSegments] = useState<Segment[]>([]);
    const [previewLoading, setPreviewLoading] = useState(false);
    const [showPreview, setShowPreview] = useState(false);

    // 保存状态
    const [saving, setSaving] = useState(false);

    // 处理状态（嵌入进度）
    const [isProcessing, setIsProcessing] = useState(false);
    const [indexingStatus, setIndexingStatus] = useState<IndexingStatus | null>(null);
    const pollingTimerRef = useRef<ReturnType<typeof setInterval> | null>(null);

    /**
     * 停止轮询
     */
    const stopPolling = useCallback(() => {
        if (pollingTimerRef.current) {
            clearInterval(pollingTimerRef.current);
            pollingTimerRef.current = null;
        }
    }, []);

    /**
     * 轮询索引状态
     */
    const pollIndexingStatus = useCallback(async (batch: string) => {
        try {
            const response = await fetchIndexingStatus(datasetId, batch);
            const status = response.data?.[0];

            if (status) {
                setIndexingStatus(status.indexing_status);

                if (status.indexing_status === 'completed') {
                    // 停止轮询
                    stopPolling();
                    setIsProcessing(false);
                    message.success('文档处理完成');

                    // 刷新分段预览
                    setPreviewLoading(true);
                    try {
                        const segmentResponse = await fetchSegments(datasetId, document?.id || '', 1, 50);
                        setPreviewSegments(segmentResponse.data || []);
                        setShowPreview(true);
                    } catch (err) {
                        console.error('刷新分段失败:', err);
                    } finally {
                        setPreviewLoading(false);
                    }
                } else if (status.indexing_status === 'error') {
                    // 停止轮询
                    stopPolling();
                    setIsProcessing(false);
                    message.error(status.error || '处理失败');
                }
            }
        } catch (err) {
            console.error('获取索引状态失败:', err);
        }
    }, [datasetId, document?.id, stopPolling]);

    /**
     * 开始轮询
     */
    const startPolling = useCallback((batch: string) => {
        stopPolling();
        pollingTimerRef.current = setInterval(() => {
            pollIndexingStatus(batch);
        }, 2000);
        // 立即执行一次
        pollIndexingStatus(batch);
    }, [stopPolling, pollIndexingStatus]);

    // 当文档变化时，从文档已有的 process_rule 回显设置，无则使用默认值
    useEffect(() => {
        if (document) {
            const rule = (document as any).process_rule;
            if (rule?.mode === 'custom' && rule?.rules) {
                const seg = rule.rules.segmentation || {};
                const preRules = rule.rules.pre_processing_rules || [];
                setSettings({
                    separator: (seg.separator || '\\n\\n').replace(/\n/g, '\\n'),
                    maxTokens: seg.max_tokens || DEFAULT_DOCUMENT_DETAIL_SETTINGS.maxTokens,
                    chunkOverlap: DEFAULT_DOCUMENT_DETAIL_SETTINGS.chunkOverlap,
                    removeExtraSpaces: preRules.find((r: any) => r.id === 'remove_extra_spaces')?.enabled ?? DEFAULT_DOCUMENT_DETAIL_SETTINGS.removeExtraSpaces,
                    removeUrlsEmails: preRules.find((r: any) => r.id === 'remove_urls_emails')?.enabled ?? DEFAULT_DOCUMENT_DETAIL_SETTINGS.removeUrlsEmails,
                    indexingTechnique: DEFAULT_DOCUMENT_DETAIL_SETTINGS.indexingTechnique,
                });
            } else {
                setSettings(DEFAULT_DOCUMENT_DETAIL_SETTINGS);
            }
            setPreviewSegments([]);
            setShowPreview(false);
            setIsProcessing(false);
            setIndexingStatus(null);
            stopPolling();
        }
    }, [document?.id, stopPolling]);

    // 组件卸载时清理定时器
    useEffect(() => {
        return () => {
            stopPolling();
        };
    }, [stopPolling]);

    /**
     * 更新设置
     */
    const updateSettings = useCallback((key: keyof DocumentDetailSegmentationSettings, value: any) => {
        setSettings(prev => ({ ...prev, [key]: value }));
    }, []);

    /**
     * 重置设置
     */
    const handleReset = useCallback(() => {
        setSettings(DEFAULT_DOCUMENT_DETAIL_SETTINGS);
        setPreviewSegments([]);
        setShowPreview(false);
    }, []);

    /**
     * 预览分段
     */
    const handlePreview = useCallback(async () => {
        if (!document) return;

        setPreviewLoading(true);
        setShowPreview(true);
        try {
            // 获取当前文档的分段作为预览
            const response = await fetchSegments(datasetId, document.id, 1, 50);
            setPreviewSegments(response.data || []);
            if (response.data?.length === 0) {
                message.info('该文档暂无分段数据');
            }
        } catch (err: any) {
            console.error('预览分段失败:', err);
            message.error(err.message || '预览失败');
        } finally {
            setPreviewLoading(false);
        }
    }, [datasetId, document]);

    /**
     * 保存并处理
     * 流程：获取原始文件 → 下载 → 用新参数重新上传 → 轮询嵌入状态
     */
    const handleSaveAndProcess = useCallback(async () => {
        if (!document) return;

        setIsProcessing(true);
        setSaving(true);
        setIndexingStatus('waiting');
        try {
            // 1. 获取原始文件信息
            message.loading({ content: '正在获取原始文件信息...', key: 'save-process' });
            const uploadFileInfo = await fetchUploadFileInfo(datasetId, document.id);

            // 2. 下载原始文件（通过代理路由）
            message.loading({ content: '正在下载原始文件...', key: 'save-process' });
            const file = await downloadOriginalFile(uploadFileInfo);

            // 3. 用新参数重新上传
            message.loading({ content: '正在应用新配置并重新处理...', key: 'save-process' });
            const result = await updateDocumentByFile(datasetId, document.id, file, {
                indexing_technique: settings.indexingTechnique,
                process_rule: {
                    mode: 'custom',
                    rules: {
                        pre_processing_rules: [
                            { id: 'remove_extra_spaces', enabled: settings.removeExtraSpaces },
                            { id: 'remove_urls_emails', enabled: settings.removeUrlsEmails },
                        ],
                        segmentation: {
                            separator: settings.separator.replace(/\\n/g, '\n'),
                            max_tokens: settings.maxTokens,
                            chunk_overlap: settings.chunkOverlap,
                        },
                    },
                },
            });

            message.success({ content: '文档正在处理中...', key: 'save-process' });

            // 4. 开始轮询嵌入状态
            startPolling(result.batch);
        } catch (err: any) {
            console.error('保存设置失败:', err);
            message.error({ content: err.message || '保存失败', key: 'save-process' });
            setIsProcessing(false);
            setIndexingStatus(null);
        } finally {
            setSaving(false);
        }
    }, [datasetId, document, settings, startPolling]);

    return {
        // 状态
        settings,
        previewSegments,
        previewLoading,
        showPreview,
        saving,
        isProcessing,
        indexingStatus,

        // 方法
        updateSettings,
        handleReset,
        handlePreview,
        handleSaveAndProcess,
    };
}

export type UseDocumentDetailReturn = ReturnType<typeof useDocumentDetail>;