dify/web/app/components/datasets/create/step-two/hooks/use-preview-state.ts

import type { NotionPage } from '@/models/common'
import type {
  CrawlResultItem,
  CustomFile,
  DocumentItem,
  FullDocumentDetail,
} from '@/models/datasets'
import { useCallback, useState } from 'react'
import { DataSourceType } from '@/models/datasets'

type UsePreviewStateOptions = {
  dataSourceType: DataSourceType
  files: CustomFile[]
  notionPages: NotionPage[]
  websitePages: CrawlResultItem[]
  documentDetail?: FullDocumentDetail
  datasetId?: string
}

export const usePreviewState = (options: UsePreviewStateOptions) => {
  const { dataSourceType, files, notionPages, websitePages, documentDetail, datasetId } = options

  // File preview state
  const [previewFile, setPreviewFile] = useState<DocumentItem>(
    datasetId && documentDetail ? documentDetail.file : files[0],
  )

  // Notion page preview state
  const [previewNotionPage, setPreviewNotionPage] = useState<NotionPage>(
    datasetId && documentDetail ? documentDetail.notion_page : notionPages[0],
  )

  // Website page preview state
  const [previewWebsitePage, setPreviewWebsitePage] = useState<CrawlResultItem>(
    datasetId && documentDetail ? documentDetail.website_page : websitePages[0],
  )

  // Get preview items for document picker based on data source type
  const getPreviewPickerItems = useCallback(() => {
    if (dataSourceType === DataSourceType.FILE) {
      return files as Array<Required<CustomFile>>
    }
    if (dataSourceType === DataSourceType.NOTION) {
      return notionPages.map((page) => ({
        id: page.page_id,
        name: page.page_name,
        extension: 'md',
      }))
    }
    if (dataSourceType === DataSourceType.WEB) {
      return websitePages.map((page) => ({
        id: page.source_url,
        name: page.title,
        extension: 'md',
      }))
    }
    return []
  }, [dataSourceType, files, notionPages, websitePages])

  // Get current preview value for picker
  const getPreviewPickerValue = useCallback(() => {
    if (dataSourceType === DataSourceType.FILE) {
      return previewFile as Required<CustomFile>
    }
    if (dataSourceType === DataSourceType.NOTION) {
      return {
        id: previewNotionPage?.page_id || '',
        name: previewNotionPage?.page_name || '',
        extension: 'md',
      }
    }
    if (dataSourceType === DataSourceType.WEB) {
      return {
        id: previewWebsitePage?.source_url || '',
        name: previewWebsitePage?.title || '',
        extension: 'md',
      }
    }
    return { id: '', name: '', extension: '' }
  }, [dataSourceType, previewFile, previewNotionPage, previewWebsitePage])

  // Handle preview change
  const handlePreviewChange = useCallback(
    (selected: { id: string; name: string }) => {
      if (dataSourceType === DataSourceType.FILE) {
        setPreviewFile(selected as DocumentItem)
      } else if (dataSourceType === DataSourceType.NOTION) {
        const selectedPage = notionPages.find((page) => page.page_id === selected.id)
        if (selectedPage) setPreviewNotionPage(selectedPage)
      } else if (dataSourceType === DataSourceType.WEB) {
        const selectedPage = websitePages.find((page) => page.source_url === selected.id)
        if (selectedPage) setPreviewWebsitePage(selectedPage)
      }
    },
    [dataSourceType, notionPages, websitePages],
  )

  return {
    // File preview
    previewFile,
    setPreviewFile,

    // Notion preview
    previewNotionPage,
    setPreviewNotionPage,

    // Website preview
    previewWebsitePage,
    setPreviewWebsitePage,

    // Picker helpers
    getPreviewPickerItems,
    getPreviewPickerValue,
    handlePreviewChange,
  }
}