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 { getMyDatasets, type AreaDataset } from '~/api/v3/dify/area-datasets';
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 [availableDatasets, setAvailableDatasets] = useState<AreaDataset[]>([]);
    const [loadingAvailableDatasets, setLoadingAvailableDatasets] = 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]);

    /**
     * 加载用户可访问的知识库列表（基于权限）
     */
    const loadAvailableDatasets = useCallback(async () => {
        setLoadingAvailableDatasets(true);
        try {
            console.log('[DatasetManager] 加载用户可访问的知识库列表...');
            const response = await getMyDatasets();
            console.log('[DatasetManager] 用户知识库列表响应:', response);

            if (response && response.code === 0 && response.data) {
                const dataList = Array.isArray(response.data.data) ? response.data.data : [];
                setAvailableDatasets(dataList);
                return dataList;
            } else {
                console.error('[DatasetManager] 获取用户知识库列表失败:', response);
                setAvailableDatasets([]);
                return [];
            }
        } catch (err: any) {
            console.error('[DatasetManager] 加载用户知识库列表失败:', err);
            setAvailableDatasets([]);
            return [];
        } finally {
            setLoadingAvailableDatasets(false);
        }
    }, []);

    /**
     * 根据 dataset_id 加载知识库详情
     */
    const loadDatasetById = useCallback(async (datasetId: string) => {
        setLoadingDataset(true);
        try {
            console.log('[DatasetManager] 加载知识库详情:', datasetId);
            const fullDataset = await fetchDataset(datasetId);
            console.log('[DatasetManager] 知识库详情响应:', fullDataset);

            setDataset(fullDataset);
            // 立即加载文档
            await loadDocuments(datasetId, 1);
        } catch (err: any) {
            console.error('[DatasetManager] 加载知识库详情失败:', err);
            setError(err.message || '加载知识库失败');
            message.error('加载知识库失败');
        } finally {
            setLoadingDataset(false);
        }
    }, [loadDocuments]);

    /**
     * 加载知识库（获取第一个知识库，再获取详情以包含 retrieval_model）
     */
    const loadDataset = useCallback(async () => {
        setLoadingDataset(true);
        try {
            console.log('[DatasetManager] 加载知识库...');

            // 先加载用户可访问的知识库列表
            const userDatasets = await loadAvailableDatasets();

            if (userDatasets.length > 0) {
                // 找到默认知识库或第一个知识库
                const defaultDataset = userDatasets.find(ds => ds.is_default) || userDatasets[0];
                const datasetId = defaultDataset.dataset_id;

                console.log('[DatasetManager] 使用知识库:', defaultDataset.dataset_name, datasetId);

                // 获取知识库详情
                const fullDataset = await fetchDataset(datasetId);
                console.log('[DatasetManager] 知识库详情响应:', fullDataset);

                setDataset(fullDataset);
                // 立即加载文档
                await loadDocuments(datasetId, 1);
            } else {
                // 回退到原有逻辑：直接从 Dify 获取
                console.log('[DatasetManager] 用户无绑定知识库，使用默认逻辑...');
                const response = await fetchDatasets(1, 1);
                console.log('[DatasetManager] Dify知识库列表响应:', response);

                if (response && response.data && response.data.length > 0) {
                    const firstDatasetId = response.data[0].id;
                    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, loadAvailableDatasets]);

    /**
     * 切换知识库
     */
    const handleDatasetChange = useCallback(async (datasetId: string) => {
        console.log('[DatasetManager] 切换知识库:', datasetId);
        // 重置状态
        setSelectedDocument(null);
        setActiveTab('documents');
        // 加载新知识库
        await loadDatasetById(datasetId);
    }, [loadDatasetById]);

    /**
     * 处理文档页码变化
     */
    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,

        // 知识库列表（基于权限）
        availableDatasets,
        loadingAvailableDatasets,

        // 方法
        loadDataset,
        loadDocuments,
        handlePageChange,
        handleDocumentDeleted,
        handleDocumentStatusChanged,
        handleRefresh,
        handleViewDocument,
        handleBackToDocuments,
        handleTabChange,
        handleDatasetUpdated,
        handleDatasetChange,
    };
}

export type UseDatasetManagerReturn = ReturnType<typeof useDatasetManager>;