useReportStudio.ts 4.67 KB
import { computed, ref, shallowRef } from 'vue'

import type { ReportStatusPayload, ReportTask } from '@/types'
import { fetchJson, postJson } from '@/utils/http'
import { bindPersistentState, loadStoredValue } from './usePersistentState'

interface ReportStatusResponse extends ReportStatusPayload {
  success: boolean
}

interface GenerateResponse {
  success: boolean
  task_id: string
  message: string
  task: ReportTask
}

const TEMPLATE_STORAGE_KEY = 'bettafish.reportTemplate.v1'
const REPORT_STATE_STORAGE_KEY = 'bettafish.reportState.v2'
const RESEARCH_STATE_STORAGE_KEY = 'bettafish.frontendState.v2'
const TERMINAL_STATUSES = new Set(['completed', 'error', 'cancelled'])

export function useReportStudio() {
  const storedTemplate = loadStoredValue(TEMPLATE_STORAGE_KEY, {
    content: '',
  })
  const storedState = loadStoredValue(REPORT_STATE_STORAGE_KEY, {
    selectedTaskId: '',
  })

  const loading = shallowRef(false)
  const generating = shallowRef(false)
  const selectedTaskId = shallowRef(String(storedState.selectedTaskId || ''))
  const customTemplate = shallowRef(String(storedTemplate.content || ''))
  const status = ref<ReportStatusPayload>({
    initialized: false,
    engines_ready: false,
    files_found: [],
    missing_files: [],
    current_task: null,
    tasks: [],
  })

  let progressTimer: number | null = null

  const selectedTask = computed(() => (
    status.value.tasks.find((item) => item.task_id === selectedTaskId.value)
    || status.value.current_task
    || null
  ))

  const previewUrl = computed(() => {
    const taskId = selectedTask.value?.task_id
    return taskId ? `/api/report/result/${taskId}` : ''
  })

  async function refreshStatus() {
    loading.value = true
    try {
      const payload = await fetchJson<ReportStatusResponse>('/api/report/status')
      status.value = {
        initialized: payload.initialized,
        engines_ready: payload.engines_ready,
        files_found: payload.files_found,
        missing_files: payload.missing_files,
        current_task: payload.current_task,
        tasks: payload.tasks,
      }

      if (!selectedTaskId.value && payload.tasks.length > 0) {
        selectedTaskId.value = payload.tasks[0].task_id
      }
    } finally {
      loading.value = false
    }
  }

  async function refreshTasks() {
    const payload = await fetchJson<ReportStatusResponse>('/api/report/tasks')
    status.value = {
      ...status.value,
      current_task: payload.current_task,
      tasks: payload.tasks,
    }
    if (!selectedTaskId.value && payload.tasks.length > 0) {
      selectedTaskId.value = payload.tasks[0].task_id
    }
  }

  function stopPolling() {
    if (progressTimer !== null) {
      window.clearInterval(progressTimer)
      progressTimer = null
    }
  }

  function startPolling(taskId: string) {
    stopPolling()
    progressTimer = window.setInterval(async () => {
      await refreshStatus()
      const task = status.value.tasks.find((item) => item.task_id === taskId)
      if (!task || TERMINAL_STATUSES.has(task.status)) {
        generating.value = false
        stopPolling()
      }
    }, 2000)
  }

  async function generate(query: string, researchTaskId = '') {
    generating.value = true
    try {
      const persistedResearchState = loadStoredValue(RESEARCH_STATE_STORAGE_KEY, {
        selectedResearchTaskId: '',
      })
      const resolvedResearchTaskId = String(
        researchTaskId || persistedResearchState.selectedResearchTaskId || '',
      ).trim()
      const payload = await postJson<GenerateResponse>('/api/report/generate', {
        query,
        custom_template: customTemplate.value,
        research_task_id: resolvedResearchTaskId,
      })
      selectedTaskId.value = payload.task_id
      await refreshStatus()
      startPolling(payload.task_id)
      return payload
    } catch (error) {
      generating.value = false
      throw error
    }
  }

  function selectTask(taskId: string) {
    selectedTaskId.value = taskId
  }

  function download(kind: 'html' | 'md' | 'pdf', taskId: string) {
    const target = {
      html: `/api/report/download/${taskId}`,
      md: `/api/report/export/md/${taskId}`,
      pdf: `/api/report/export/pdf/${taskId}`,
    }[kind]
    window.open(target, '_blank', 'noopener')
  }

  bindPersistentState(
    TEMPLATE_STORAGE_KEY,
    computed(() => ({
      content: customTemplate.value,
    })),
  )
  bindPersistentState(
    REPORT_STATE_STORAGE_KEY,
    computed(() => ({
      selectedTaskId: selectedTaskId.value,
    })),
  )

  return {
    loading,
    generating,
    status,
    selectedTask,
    selectedTaskId,
    customTemplate,
    previewUrl,
    refreshStatus,
    refreshTasks,
    generate,
    selectTask,
    download,
    stopPolling,
  }
}