crosssubtitle-ai/src/stores/tasks.ts

import { defineStore } from 'pinia'
import { invoke } from '@tauri-apps/api/core'
import { listen, type UnlistenFn } from '@tauri-apps/api/event'
import type {
  ErrorEvent,
  LogEvent,
  ProgressEvent,
  ResetSegmentsEvent,
  SegmentEvent,
  StartTaskPayload,
  SubtitleSegment,
  SubtitleTask,
} from '../lib/types'

type ExportFormat = 'srt' | 'vtt' | 'ass'

function sortSegments(segments: SubtitleSegment[]) {
  segments.sort((left, right) => {
    if (left.start !== right.start) return left.start - right.start
    if (left.end !== right.end) return left.end - right.end
    return left.id.localeCompare(right.id)
  })
}

export const useTaskStore = defineStore('tasks', {
  state: () => ({
    tasks: [] as SubtitleTask[],
    logsByTaskId: {} as Record<string, string[]>,
    selectedTaskId: '' as string,
    ready: false,
    unlisteners: [] as UnlistenFn[],
    pendingPayloads: {} as Record<string, StartTaskPayload>,
  }),
  getters: {
    selectedTask(state) {
      return state.tasks.find((task) => task.id === state.selectedTaskId) ?? null
    },
    selectedTaskLogs(state) {
      return state.logsByTaskId[state.selectedTaskId] ?? []
    },
  },
  actions: {
    async initialize() {
      if (this.ready) return

      this.tasks = await invoke<SubtitleTask[]>('list_tasks')

      const progressUnlisten = await listen<ProgressEvent>('task:progress', ({ payload }) => {
        const task = this.tasks.find((item) => item.id === payload.taskId)
        if (!task) return
        task.status = payload.status
        task.progress = payload.progress
      })

      const segmentUnlisten = await listen<SegmentEvent>('task:segment', ({ payload }) => {
        const task = this.tasks.find((item) => item.id === payload.taskId)
        if (!task) return

        const existing = task.segments.find((segment) => segment.id === payload.segment.id)
        if (existing) {
          Object.assign(existing, payload.segment)
        } else {
          task.segments.push(payload.segment)
        }
        sortSegments(task.segments)
      })

      const resetSegmentsUnlisten = await listen<ResetSegmentsEvent>('task:segments_reset', ({ payload }) => {
        const task = this.tasks.find((item) => item.id === payload.taskId)
        if (!task) return
        task.segments = []
      })

      const errorUnlisten = await listen<ErrorEvent>('task:error', ({ payload }) => {
        const task = this.tasks.find((item) => item.id === payload.taskId)
        if (!task) return
        task.status = 'failed'
        task.error = payload.message
      })

      const logUnlisten = await listen<LogEvent>('task:log', ({ payload }) => {
        if (!this.logsByTaskId[payload.taskId]) {
          this.logsByTaskId[payload.taskId] = []
        }
        this.logsByTaskId[payload.taskId].push(payload.message)
        if (this.logsByTaskId[payload.taskId].length > 300) {
          this.logsByTaskId[payload.taskId] = this.logsByTaskId[payload.taskId].slice(-300)
        }
      })

      const doneUnlisten = await listen<SubtitleTask>('task:done', ({ payload }) => {
        sortSegments(payload.segments)
        const index = this.tasks.findIndex((item) => item.id === payload.id)
        if (index >= 0) {
          this.tasks[index] = payload
        } else {
          this.tasks.unshift(payload)
        }
      })

      this.unlisteners = [
        progressUnlisten,
        segmentUnlisten,
        resetSegmentsUnlisten,
        errorUnlisten,
        logUnlisten,
        doneUnlisten,
      ]
      this.ready = true
    },

    async startTask(payload: StartTaskPayload) {
      const task = await invoke<SubtitleTask>('start_subtitle_task', { payload })
      this.tasks.unshift(task)
      this.pendingPayloads[task.id] = payload
      this.logsByTaskId[task.id] = []
      this.selectedTaskId = task.id
    },

    async retryTask(taskId: string) {
      const payload = this.pendingPayloads[taskId]
      if (!payload) return

      const index = this.tasks.findIndex((t) => t.id === taskId)
      if (index >= 0) {
        this.tasks.splice(index, 1)
      }
      delete this.logsByTaskId[taskId]

      await this.startTask(payload)
    },

    selectTask(taskId: string) {
      this.selectedTaskId = taskId
    },

    async updateSegment(segment: SubtitleSegment) {
      const updated = await invoke<SubtitleTask>('update_segment_text', { segment })
      const index = this.tasks.findIndex((item) => item.id === updated.id)
      if (index >= 0) this.tasks[index] = updated
    },

    async deleteTask(taskId: string) {
      await invoke('delete_task', { taskId })
      const index = this.tasks.findIndex((t) => t.id === taskId)
      if (index >= 0) {
        this.tasks.splice(index, 1)
      }
      delete this.logsByTaskId[taskId]
      if (this.selectedTaskId === taskId) {
        this.selectedTaskId = ''
      }
    },

    async exportTask(taskId: string, format: ExportFormat) {
      return invoke<string>('export_subtitles', { taskId, format })
    },
  },
})