ComfyUI_frontend/src/composables/queue/useJobList.ts

import { computed, onUnmounted, ref, watch } from 'vue'
import { useI18n } from 'vue-i18n'

import { useQueueProgress } from '@/composables/queue/useQueueProgress'
import { st } from '@/i18n'
import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
import { useExecutionStore } from '@/stores/executionStore'
import { useQueueStore } from '@/stores/queueStore'
import type { TaskItemImpl } from '@/stores/queueStore'
import type { JobState } from '@/types/queue'
import {
  dateKey,
  formatClockTime,
  formatShortMonthDay,
  isToday,
  isYesterday
} from '@/utils/dateTimeUtil'
import { normalizeI18nKey } from '@/utils/formatUtil'
import { buildJobDisplay } from '@/utils/queueDisplay'
import { jobStateFromTask } from '@/utils/queueUtil'

/** Tabs for job list filtering */
export const jobTabs = ['All', 'Completed', 'Failed'] as const
export type JobTab = (typeof jobTabs)[number]

export const jobSortModes = ['mostRecent', 'totalGenerationTime'] as const
export type JobSortMode = (typeof jobSortModes)[number]

/**
 * UI item in the job list. Mirrors data previously prepared inline.
 */
export type JobListItem = {
  id: string
  title: string
  meta: string
  state: JobState
  iconName?: string
  iconImageUrl?: string
  showClear?: boolean
  taskRef?: any
  progressTotalPercent?: number
  progressCurrentPercent?: number
  runningNodeName?: string
  executionTimeMs?: number
  computeHours?: number
}

export type JobGroup = {
  key: string
  label: string
  items: JobListItem[]
}

const ADDED_HINT_DURATION_MS = 3000
const relativeTimeFormatterCache = new Map<string, Intl.RelativeTimeFormat>()
const taskIdToKey = (id: string | number | undefined) => {
  if (id === null || id === undefined) return null
  const key = String(id)
  return key.length ? key : null
}

/**
 * Returns localized Today/Yesterday (capitalized) or localized Mon DD.
 */
const dateLabelForTimestamp = (
  ts: number,
  locale: string,
  relativeFormatter: Intl.RelativeTimeFormat
) => {
  const formatRelativeDay = (value: number) => {
    const formatted = relativeFormatter.format(value, 'day')
    return formatted
      ? formatted[0].toLocaleUpperCase(locale) + formatted.slice(1)
      : formatted
  }
  if (isToday(ts)) {
    return formatRelativeDay(0)
  }
  if (isYesterday(ts)) {
    return formatRelativeDay(-1)
  }
  return formatShortMonthDay(ts, locale)
}

type TaskWithState = {
  task: TaskItemImpl
  state: JobState
}

const orderingTimestamp = (task: TaskItemImpl, state: JobState) => {
  if (state === 'completed' || state === 'failed') {
    return (
      task.executionEndTimestamp ??
      task.executionStartTimestamp ??
      task.createTime
    )
  }
  return task.createTime
}

const compareTasksByRecency = (a: TaskWithState, b: TaskWithState) => {
  const tsA = orderingTimestamp(a.task, a.state)
  const tsB = orderingTimestamp(b.task, b.state)

  if (tsA !== undefined && tsB !== undefined && tsA !== tsB) {
    return tsB - tsA
  }
  if (tsA !== undefined && tsB === undefined) return -1
  if (tsA === undefined && tsB !== undefined) return 1

  return b.task.queueIndex - a.task.queueIndex
}

/**
 * Builds the reactive job list, filters, and grouped view for the queue overlay.
 */
export function useJobList() {
  const { t, locale } = useI18n()
  const queueStore = useQueueStore()
  const executionStore = useExecutionStore()
  const workflowStore = useWorkflowStore()

  const recentlyAddedPendingIds = ref<Set<string>>(new Set())
  const addedHintTimeouts = new Map<string, ReturnType<typeof setTimeout>>()

  const clearAddedHintTimeout = (id: string) => {
    const timeoutId = addedHintTimeouts.get(id)
    if (timeoutId !== undefined) {
      clearTimeout(timeoutId)
      addedHintTimeouts.delete(id)
    }
  }

  const scheduleAddedHintExpiry = (id: string) => {
    clearAddedHintTimeout(id)
    const timeoutId = setTimeout(() => {
      addedHintTimeouts.delete(id)
      const updated = new Set(recentlyAddedPendingIds.value)
      if (updated.delete(id)) {
        recentlyAddedPendingIds.value = updated
      }
    }, ADDED_HINT_DURATION_MS)
    addedHintTimeouts.set(id, timeoutId)
  }

  watch(
    () =>
      queueStore.pendingTasks
        .map((task) => taskIdToKey(task.promptId))
        .filter((id): id is string => !!id),
    (pendingIds) => {
      const pendingSet = new Set(pendingIds)
      const next = new Set(recentlyAddedPendingIds.value)

      pendingIds.forEach((id) => {
        if (!next.has(id)) {
          next.add(id)
          scheduleAddedHintExpiry(id)
        }
      })

      for (const id of Array.from(next)) {
        if (!pendingSet.has(id)) {
          next.delete(id)
          clearAddedHintTimeout(id)
        }
      }

      recentlyAddedPendingIds.value = next
    },
    { immediate: true }
  )

  const shouldShowAddedHint = (task: TaskItemImpl, state: JobState) => {
    if (state !== 'pending') return false
    const id = taskIdToKey(task.promptId)
    if (!id) return false
    return recentlyAddedPendingIds.value.has(id)
  }

  onUnmounted(() => {
    addedHintTimeouts.forEach((timeoutId) => clearTimeout(timeoutId))
    addedHintTimeouts.clear()
    recentlyAddedPendingIds.value = new Set<string>()
  })

  const { totalPercent, currentNodePercent } = useQueueProgress()

  const relativeTimeFormatter = computed(() => {
    const localeValue = locale.value
    let formatter = relativeTimeFormatterCache.get(localeValue)
    if (!formatter) {
      formatter = new Intl.RelativeTimeFormat(localeValue, { numeric: 'auto' })
      relativeTimeFormatterCache.set(localeValue, formatter)
    }
    return formatter
  })
  const undatedLabel = computed(() => t('queue.jobList.undated') || 'Undated')

  const isJobInitializing = (promptId: string | number | undefined) =>
    executionStore.isPromptInitializing(promptId)

  const currentNodeName = computed(() => {
    const node = executionStore.executingNode
    if (!node) return t('g.emDash')
    const title = (node.title ?? '').toString().trim()
    if (title) return title
    const nodeType = (node.type ?? '').toString().trim() || t('g.untitled')
    const key = `nodeDefs.${normalizeI18nKey(nodeType)}.display_name`
    return st(key, nodeType)
  })

  const selectedJobTab = ref<JobTab>('All')
  const selectedWorkflowFilter = ref<'all' | 'current'>('all')
  const selectedSortMode = ref<JobSortMode>('mostRecent')

  const tasksWithJobState = computed<TaskWithState[]>(() => {
    const entries: TaskWithState[] = [
      ...queueStore.pendingTasks,
      ...queueStore.runningTasks,
      ...queueStore.historyTasks
    ].map((task) => ({
      task,
      state: jobStateFromTask(task, isJobInitializing(task?.promptId))
    }))

    return entries.sort(compareTasksByRecency)
  })

  const allTasksSorted = computed<TaskItemImpl[]>(() =>
    tasksWithJobState.value.map(({ task }) => task)
  )

  const hasFailedJobs = computed(() =>
    tasksWithJobState.value.some(({ state }) => state === 'failed')
  )

  watch(
    () => hasFailedJobs.value,
    (hasFailed) => {
      if (!hasFailed && selectedJobTab.value === 'Failed') {
        selectedJobTab.value = 'All'
      }
    }
  )

  const filteredTaskEntries = computed<TaskWithState[]>(() => {
    let entries = tasksWithJobState.value
    if (selectedJobTab.value === 'Completed') {
      entries = entries.filter(({ state }) => state === 'completed')
    } else if (selectedJobTab.value === 'Failed') {
      entries = entries.filter(({ state }) => state === 'failed')
    }

    if (selectedWorkflowFilter.value === 'current') {
      const activeId = workflowStore.activeWorkflow?.activeState?.id
      if (!activeId) return []
      entries = entries.filter(({ task }) => {
        const wid = task.workflow?.id
        return !!wid && wid === activeId
      })
    }
    return entries
  })

  const filteredTasks = computed<TaskItemImpl[]>(() =>
    filteredTaskEntries.value.map(({ task }) => task)
  )

  const jobItems = computed<JobListItem[]>(() => {
    return filteredTaskEntries.value.map(({ task, state }) => {
      const isActive =
        String(task.promptId ?? '') ===
        String(executionStore.activePromptId ?? '')
      const showAddedHint = shouldShowAddedHint(task, state)

      const display = buildJobDisplay(task, state, {
        t,
        locale: locale.value,
        formatClockTimeFn: formatClockTime,
        isActive,
        totalPercent: isActive ? totalPercent.value : undefined,
        currentNodePercent: isActive ? currentNodePercent.value : undefined,
        currentNodeName: isActive ? currentNodeName.value : undefined,
        showAddedHint
      })

      return {
        id: String(task.promptId),
        title: display.primary,
        meta: display.secondary,
        state,
        iconName: display.iconName,
        iconImageUrl: display.iconImageUrl,
        showClear: display.showClear,
        taskRef: task,
        progressTotalPercent:
          state === 'running' && isActive ? totalPercent.value : undefined,
        progressCurrentPercent:
          state === 'running' && isActive
            ? currentNodePercent.value
            : undefined,
        runningNodeName:
          state === 'running' && isActive ? currentNodeName.value : undefined,
        executionTimeMs: task.executionTime,
        computeHours:
          task.executionTime !== undefined
            ? task.executionTime / 3_600_000
            : undefined
      } as JobListItem
    })
  })

  const jobItemById = computed(() => {
    const m = new Map<string, JobListItem>()
    jobItems.value.forEach((ji) => m.set(ji.id, ji))
    return m
  })

  const groupedJobItems = computed<JobGroup[]>(() => {
    const groups: JobGroup[] = []
    const index = new Map<string, number>()
    const localeValue = locale.value
    for (const { task, state } of filteredTaskEntries.value) {
      const ts = orderingTimestamp(task, state)
      const key = ts === undefined ? 'undated' : dateKey(ts)
      let groupIdx = index.get(key)
      if (groupIdx === undefined) {
        const label =
          ts === undefined
            ? undatedLabel.value
            : dateLabelForTimestamp(
                ts,
                localeValue,
                relativeTimeFormatter.value
              )
        groups.push({ key, label, items: [] })
        groupIdx = groups.length - 1
        index.set(key, groupIdx)
      }
      const ji = jobItemById.value.get(String(task.promptId))
      if (ji) groups[groupIdx].items.push(ji)
    }

    if (selectedSortMode.value === 'totalGenerationTime') {
      const valueOrDefault = (value: JobListItem['executionTimeMs']) =>
        typeof value === 'number' && !Number.isNaN(value) ? value : -1
      const sortByExecutionTimeDesc = (a: JobListItem, b: JobListItem) =>
        valueOrDefault(b.executionTimeMs) - valueOrDefault(a.executionTimeMs)

      groups.forEach((group) => {
        group.items.sort(sortByExecutionTimeDesc)
      })
    }

    const undated = groups.filter((group) => group.key === 'undated')
    const dated = groups.filter((group) => group.key !== 'undated')

    return [...undated, ...dated]
  })

  return {
    // filters/state
    selectedJobTab,
    selectedWorkflowFilter,
    selectedSortMode,
    hasFailedJobs,
    // data sources
    allTasksSorted,
    filteredTasks,
    jobItems,
    groupedJobItems,
    currentNodeName
  }
}