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

import { useState, useEffect, useCallback } from 'react';
import { message } from 'antd';
import type { Dataset } from '~/api/dify-dataset/type/datasetTypes';
import type { Document } from '~/api/dify-dataset/type/documentTypes';
import { fetchDatasets, fetchDataset } from '~/api/dify-dataset/api/datasetApi';
import { fetchDocuments } from '~/api/dify-dataset/api/documentApi';
import type { MenuTab } from '~/types/dify-dataset-manager/layout';
import { DEFAULT_DOCUMENT_PAGE_SIZE } from '~/types/dify-dataset-manager/index';

/**
 * 知识库管理器状态管理 Hook
 */
export function useDatasetManager() {
    // 知识库状态
    const [dataset, setDataset] = useState<Dataset | null>(null);
    const [loadingDataset, setLoadingDataset] = useState(true);

    // 文档状态
    const [documents, setDocuments] = useState<Document[]>([]);
    const [loadingDocuments, setLoadingDocuments] = useState(false);
    const [documentTotal, setDocumentTotal] = useState(0);
    const [documentPage, setDocumentPage] = useState(1);
    const [documentPageSize] = useState(DEFAULT_DOCUMENT_PAGE_SIZE);

    // 初始化状态
    const [inited, setInited] = useState(false);
    const [error, setError] = useState<string | null>(null);

    // 菜单状态
    const [activeTab, setActiveTab] = useState<MenuTab>('documents');

    // 选中的文档（用于查看文档详情）
    const [selectedDocument, setSelectedDocument] = useState<Document | null>(null);

    /**
     * 加载文档列表
     */
    const loadDocuments = useCallback(async (datasetId: string, page: number = 1) => {
        if (!datasetId) return;

        setLoadingDocuments(true);
        try {
            console.log('[DatasetManager] 加载文档列表:', { datasetId, page });
            const response = await fetchDocuments(datasetId, page, documentPageSize);
            console.log('[DatasetManager] 文档列表响应:', response);

            if (response && response.data) {
                setDocuments(response.data);
                setDocumentTotal(response.total);
                setDocumentPage(page);
            }
        } catch (err: any) {
            console.error('[DatasetManager] 加载文档列表失败:', err);
            message.error('加载文档列表失败');
        } finally {
            setLoadingDocuments(false);
        }
    }, [documentPageSize]);

    /**
     * 加载知识库（获取第一个知识库，再获取详情以包含 retrieval_model）
     */
    const loadDataset = useCallback(async () => {
        setLoadingDataset(true);
        try {
            console.log('[DatasetManager] 加载知识库...');
            // 先获取列表，找到第一个知识库的 ID
            const response = await fetchDatasets(1, 1);
            console.log('[DatasetManager] 知识库列表响应:', response);

            if (response && response.data && response.data.length > 0) {
                const firstDatasetId = response.data[0].id;

                // 再获取详情，包含完整的 retrieval_model 等字段
                const fullDataset = await fetchDataset(firstDatasetId);
                console.log('[DatasetManager] 知识库详情响应:', fullDataset);

                setDataset(fullDataset);
                // 立即加载文档
                await loadDocuments(firstDatasetId, 1);
            } else {
                setError('未找到知识库，请先在Dify中创建知识库');
            }
        } catch (err: any) {
            console.error('[DatasetManager] 加载知识库失败:', err);
            setError(err.message || '加载知识库失败');
            message.error('加载知识库失败');
        } finally {
            setLoadingDataset(false);
            setInited(true);
        }
    }, [loadDocuments]);

    /**
     * 处理文档页码变化
     */
    const handlePageChange = useCallback((page: number) => {
        if (dataset) {
            loadDocuments(dataset.id, page);
        }
    }, [dataset, loadDocuments]);

    /**
     * 处理文档删除
     */
    const handleDocumentDeleted = useCallback((documentId: string) => {
        setDocuments((prev) => prev.filter((doc) => doc.id !== documentId));
        setDocumentTotal((prev) => prev - 1);

        // 更新知识库的文档数量
        setDataset((prev) => {
            if (prev) {
                return {
                    ...prev,
                    document_count: prev.document_count - 1
                };
            }
            return prev;
        });
    }, []);

    /**
     * 处理文档状态变化
     */
    const handleDocumentStatusChanged = useCallback((documentId: string, enabled: boolean) => {
        setDocuments((prev) =>
            prev.map((doc) =>
                doc.id === documentId ? { ...doc, enabled } : doc
            )
        );
    }, []);

    /**
     * 刷新文档列表
     */
    const handleRefresh = useCallback(() => {
        if (dataset) {
            loadDocuments(dataset.id, documentPage);
        }
    }, [dataset, documentPage, loadDocuments]);

    /**
     * 查看文档详情（分段管理）
     */
    const handleViewDocument = useCallback((doc: Document) => {
        console.log('[DatasetManager] 查看文档详情:', doc);
        setSelectedDocument(doc);
    }, []);

    /**
     * 返回文档列表
     */
    const handleBackToDocuments = useCallback(() => {
        setSelectedDocument(null);
    }, []);

    /**
     * 处理菜单切换
     */
    const handleTabChange = useCallback((tab: MenuTab) => {
        setActiveTab(tab);
        // 切换菜单时清除选中的文档
        if (tab !== 'documents') {
            setSelectedDocument(null);
        }
    }, []);

    /**
     * 处理知识库更新
     */
    const handleDatasetUpdated = useCallback((updatedDataset: Dataset) => {
        setDataset(updatedDataset);
    }, []);

    // 初始化
    useEffect(() => {
        loadDataset();
    }, [loadDataset]);

    return {
        // 状态
        dataset,
        loadingDataset,
        documents,
        loadingDocuments,
        documentTotal,
        documentPage,
        documentPageSize,
        inited,
        error,
        activeTab,
        selectedDocument,
        
        // 方法
        loadDataset,
        loadDocuments,
        handlePageChange,
        handleDocumentDeleted,
        handleDocumentStatusChanged,
        handleRefresh,
        handleViewDocument,
        handleBackToDocuments,
        handleTabChange,
        handleDatasetUpdated,
    };
}

export type UseDatasetManagerReturn = ReturnType<typeof useDatasetManager>;