geo/apps/desktop-client/src/main/monitor-scheduler.ts

import { createHash } from 'node:crypto'
import { mkdirSync, readFileSync, writeFileSync } from 'node:fs'
import { dirname, join } from 'node:path'

import type { DesktopTaskEventMessage, DesktopTaskInfo, JsonValue } from '@geo/shared-types'
import { app } from 'electron/main'

export type MonitorSchedulerRouting = 'rabbitmq-primary' | 'db-recovery'
type MonitorTaskState = 'queued' | 'leasing' | 'active'
export type MonitorTaskSource = 'desktop_task' | 'monitoring_lease'

export interface MonitorSchedulerTaskRecord {
  taskId: string
  jobId: string
  accountId: string
  clientId: string
  platform: string
  routing: MonitorSchedulerRouting
  source: MonitorTaskSource
  state: MonitorTaskState
  priority: number
  lane: string | null
  laneWeight: number
  title: string | null
  businessDate: string | null
  questionKey: string | null
  questionText: string | null
  updatedAt: number
  enqueuedAt: number
  availableAt: number
  lastSeenAt: number
  lastLeaseAttemptAt: number | null
  lastStartedAt: number | null
  leaseMisses: number
}

interface PersistedMonitorSchedulerState {
  version: number
  tasks: Record<string, MonitorSchedulerTaskRecord>
  questionCooldowns: Record<string, number>
  platformCooldowns: Record<string, number>
  lastHydratedAt: number
}

export interface MonitorSchedulerSnapshot {
  hydratedAt: number | null
  queueDepth: number
  activeCount: number
  leasingCount: number
  queuedCount: number
  questionCooldownCount: number
  platformCooldownCount: number
  tasks: MonitorSchedulerTaskRecord[]
}

export interface MonitorSchedulerSelection {
  taskId: string
  routing: MonitorSchedulerRouting
  source: MonitorTaskSource
}

interface MonitorSchedulerSelectionOptions {
  now?: number
  maxConcurrency: number
  activePlatforms: ReadonlySet<string>
  activeQuestionKeys: ReadonlySet<string>
  activeCount: number
}

interface MonitorTaskMetadata {
  title: string | null
  businessDate: string | null
  questionKey: string | null
  questionText: string | null
}

const schedulerStateVersion = 5
const schedulerQuestionCooldownMs = 45_000
const schedulerPlatformCooldownMs = 5_000
const schedulerRestartRecoveryDelayMs = 90_000
const schedulerUnknownQuestionParallelismFloor = 1
const schedulerMaxLeaseMisses = 4

let persistedStateCache: PersistedMonitorSchedulerState | null = null

function persistedStatePath(): string {
  return join(app.getPath('userData'), 'desktop-monitor-scheduler.json')
}

function defaultPersistedState(): PersistedMonitorSchedulerState {
  return {
    version: schedulerStateVersion,
    tasks: {},
    questionCooldowns: {},
    platformCooldowns: {},
    lastHydratedAt: 0,
  }
}

function clonePersistedState(
  state: PersistedMonitorSchedulerState,
): PersistedMonitorSchedulerState {
  return {
    version: state.version,
    tasks: Object.fromEntries(
      Object.entries(state.tasks).map(([taskId, task]) => [taskId, { ...task }]),
    ),
    questionCooldowns: { ...state.questionCooldowns },
    platformCooldowns: { ...state.platformCooldowns },
    lastHydratedAt: state.lastHydratedAt,
  }
}

function readPersistedState(): PersistedMonitorSchedulerState {
  if (persistedStateCache) {
    return persistedStateCache
  }

  try {
    const parsed = JSON.parse(
      readFileSync(persistedStatePath(), 'utf8'),
    ) as PersistedMonitorSchedulerState
    persistedStateCache = normalizePersistedState(parsed)
  } catch {
    persistedStateCache = defaultPersistedState()
  }

  return persistedStateCache
}

function writePersistedState(next: PersistedMonitorSchedulerState): void {
  persistedStateCache = next
  const target = persistedStatePath()
  mkdirSync(dirname(target), { recursive: true })
  writeFileSync(target, JSON.stringify(next, null, 2), 'utf8')
}

function mutatePersistedState(
  mutator: (draft: PersistedMonitorSchedulerState) => void,
): PersistedMonitorSchedulerState {
  const draft = clonePersistedState(readPersistedState())
  mutator(draft)
  prunePersistedState(draft, Date.now())
  writePersistedState(draft)
  return draft
}

function normalizePersistedState(
  input: PersistedMonitorSchedulerState | null | undefined,
): PersistedMonitorSchedulerState {
  const next = defaultPersistedState()
  if (!input || typeof input !== 'object') {
    return next
  }

  // The monitor scheduler is only a local optimization layer. When its on-disk
  // shape changes, start from a clean slate and let the server-side queue/resume
  // APIs repopulate state. Keeping old records risks reviving already-finished
  // monitor desktop_tasks and blocking legacy fallback pulls.
  if (input.version !== schedulerStateVersion) {
    return next
  }

  next.version = schedulerStateVersion
  next.lastHydratedAt = typeof input.lastHydratedAt === 'number' ? input.lastHydratedAt : 0

  if (input.tasks && typeof input.tasks === 'object') {
    for (const [taskId, value] of Object.entries(input.tasks)) {
      if (!value || typeof value !== 'object') {
        continue
      }

      const task = value as Partial<MonitorSchedulerTaskRecord>
      if (
        typeof task.taskId !== 'string' ||
        task.taskId !== taskId ||
        typeof task.platform !== 'string'
      ) {
        continue
      }

      next.tasks[taskId] = {
        taskId,
        jobId: typeof task.jobId === 'string' ? task.jobId : '',
        accountId: typeof task.accountId === 'string' ? task.accountId : '',
        clientId: typeof task.clientId === 'string' ? task.clientId : '',
        platform: task.platform,
        routing: task.routing === 'db-recovery' ? 'db-recovery' : 'rabbitmq-primary',
        source: task.source === 'monitoring_lease' ? 'monitoring_lease' : 'desktop_task',
        state: task.state === 'active' || task.state === 'leasing' ? task.state : 'queued',
        priority: typeof task.priority === 'number' ? task.priority : 100,
        lane: normalizeOptionalString(task.lane),
        laneWeight:
          typeof task.laneWeight === 'number' ? task.laneWeight : laneWeightFromLane(task.lane),
        title: typeof task.title === 'string' ? task.title : null,
        businessDate: normalizeBusinessDate(task.businessDate),
        questionKey: normalizeOptionalString(task.questionKey),
        questionText: normalizeOptionalString(task.questionText),
        updatedAt: typeof task.updatedAt === 'number' ? task.updatedAt : Date.now(),
        enqueuedAt: typeof task.enqueuedAt === 'number' ? task.enqueuedAt : Date.now(),
        availableAt: typeof task.availableAt === 'number' ? task.availableAt : Date.now(),
        lastSeenAt: typeof task.lastSeenAt === 'number' ? task.lastSeenAt : Date.now(),
        lastLeaseAttemptAt:
          typeof task.lastLeaseAttemptAt === 'number' ? task.lastLeaseAttemptAt : null,
        lastStartedAt: typeof task.lastStartedAt === 'number' ? task.lastStartedAt : null,
        leaseMisses: typeof task.leaseMisses === 'number' ? Math.max(0, task.leaseMisses) : 0,
      }
    }
  }

  if (input.questionCooldowns && typeof input.questionCooldowns === 'object') {
    for (const [key, value] of Object.entries(input.questionCooldowns)) {
      if (typeof value === 'number' && value > 0) {
        next.questionCooldowns[key] = value
      }
    }
  }

  if (input.platformCooldowns && typeof input.platformCooldowns === 'object') {
    for (const [key, value] of Object.entries(input.platformCooldowns)) {
      if (typeof value === 'number' && value > 0) {
        next.platformCooldowns[key] = value
      }
    }
  }

  prunePersistedState(next, Date.now())
  return next
}

function prunePersistedState(state: PersistedMonitorSchedulerState, now: number): void {
  const currentBusinessDate = currentMonitoringBusinessDate(now)

  for (const [taskId, task] of Object.entries(state.tasks)) {
    if (task.businessDate && task.businessDate < currentBusinessDate) {
      delete state.tasks[taskId]
      continue
    }

    if (task.state === 'active' || task.state === 'leasing') {
      task.state = 'queued'
      task.availableAt = Math.max(task.availableAt, now + schedulerRestartRecoveryDelayMs)
    }

    if (task.leaseMisses >= schedulerMaxLeaseMisses) {
      delete state.tasks[taskId]
    }
  }

  for (const [questionKey, until] of Object.entries(state.questionCooldowns)) {
    if (until <= now) {
      delete state.questionCooldowns[questionKey]
    }
  }

  for (const [platform, until] of Object.entries(state.platformCooldowns)) {
    if (until <= now) {
      delete state.platformCooldowns[platform]
    }
  }
}

export function initMonitorScheduler(): void {
  const now = Date.now()
  const next = clonePersistedState(readPersistedState())
  prunePersistedState(next, now)
  for (const [taskId, task] of Object.entries(next.tasks)) {
    if (task.source === 'monitoring_lease') {
      delete next.tasks[taskId]
    }
  }
  next.lastHydratedAt = now
  writePersistedState(next)
}

export function enqueueMonitorTaskFromEvent(
  event: DesktopTaskEventMessage,
  routing: MonitorSchedulerRouting,
): void {
  const metadata = metadataFromEvent(event)
  const now = Date.now()

  mutatePersistedState((draft) => {
    const existing = draft.tasks[event.task_id]
    draft.tasks[event.task_id] = {
      taskId: event.task_id,
      jobId: event.job_id,
      accountId: event.target_account_id,
      clientId: event.target_client_id,
      platform: event.platform,
      routing,
      source: 'desktop_task',
      state: existing?.state === 'active' ? 'active' : 'queued',
      priority: normalizePriority(event.priority, existing?.priority),
      lane: normalizeLane(event.lane, existing?.lane),
      laneWeight: laneWeightFromLane(event.lane ?? existing?.lane ?? null),
      title: metadata.title ?? existing?.title ?? null,
      businessDate: metadata.businessDate ?? existing?.businessDate ?? null,
      questionKey: metadata.questionKey ?? existing?.questionKey ?? null,
      questionText: metadata.questionText ?? existing?.questionText ?? null,
      updatedAt: parseTimestamp(event.updated_at) ?? now,
      enqueuedAt: existing?.enqueuedAt ?? now,
      availableAt: existing?.availableAt ?? now,
      lastSeenAt: now,
      lastLeaseAttemptAt: existing?.lastLeaseAttemptAt ?? null,
      lastStartedAt: existing?.lastStartedAt ?? null,
      leaseMisses: existing?.leaseMisses ?? 0,
    }
  })
}

export function noteMonitorTaskLeased(
  task: DesktopTaskInfo,
  routing: MonitorSchedulerRouting,
): void {
  const now = Date.now()
  const metadata = metadataFromPayload(task.payload ?? null)

  mutatePersistedState((draft) => {
    const existing = draft.tasks[task.id]
    draft.tasks[task.id] = {
      taskId: task.id,
      jobId: task.job_id,
      accountId: task.target_account_id,
      clientId: task.target_client_id,
      platform: task.platform,
      routing,
      source: 'desktop_task',
      state: 'active',
      priority: normalizePriority(
        payloadNumber(task.payload ?? null, 'priority'),
        existing?.priority,
      ),
      lane: normalizeLane(firstString(task.payload ?? null, ['lane']), existing?.lane),
      laneWeight: laneWeightFromLane(
        firstString(task.payload ?? null, ['lane']) ?? existing?.lane ?? null,
      ),
      title: metadata.title ?? existing?.title ?? null,
      businessDate: metadata.businessDate ?? existing?.businessDate ?? null,
      questionKey: metadata.questionKey ?? existing?.questionKey ?? null,
      questionText: metadata.questionText ?? existing?.questionText ?? null,
      updatedAt: parseTimestamp(task.updated_at) ?? now,
      enqueuedAt: existing?.enqueuedAt ?? now,
      availableAt: now,
      lastSeenAt: now,
      lastLeaseAttemptAt: now,
      lastStartedAt: now,
      leaseMisses: existing?.leaseMisses ?? 0,
    }
  })
}

export function enqueueMonitorLeaseTask(input: {
  taskId: string
  jobId?: string
  accountId?: string
  clientId: string
  platform: string
  routing: MonitorSchedulerRouting
  priority?: number
  lane?: string | null
  title?: string | null
  businessDate?: string | null
  questionKey?: string | null
  questionText?: string | null
}): void {
  const now = Date.now()

  mutatePersistedState((draft) => {
    const existing = draft.tasks[input.taskId]
    draft.tasks[input.taskId] = {
      taskId: input.taskId,
      jobId: input.jobId ?? input.taskId,
      accountId: input.accountId ?? existing?.accountId ?? '',
      clientId: input.clientId,
      platform: input.platform,
      routing: input.routing,
      source: 'monitoring_lease',
      state: existing?.state === 'active' ? 'active' : 'queued',
      priority: normalizePriority(input.priority, existing?.priority),
      lane: normalizeLane(input.lane, existing?.lane),
      laneWeight: laneWeightFromLane(input.lane ?? existing?.lane ?? null),
      title: normalizeOptionalString(input.title) ?? existing?.title ?? null,
      businessDate: normalizeBusinessDate(input.businessDate) ?? existing?.businessDate ?? null,
      questionKey: normalizeOptionalString(input.questionKey) ?? existing?.questionKey ?? null,
      questionText: normalizeOptionalString(input.questionText) ?? existing?.questionText ?? null,
      updatedAt: now,
      enqueuedAt: existing?.enqueuedAt ?? now,
      availableAt: existing?.availableAt ?? now,
      lastSeenAt: now,
      lastLeaseAttemptAt: existing?.lastLeaseAttemptAt ?? null,
      lastStartedAt: existing?.lastStartedAt ?? null,
      leaseMisses: existing?.leaseMisses ?? 0,
    }
  })
}

export function noteMonitorTaskActivated(taskId: string): void {
  const now = Date.now()

  mutatePersistedState((draft) => {
    const existing = draft.tasks[taskId]
    if (!existing) {
      return
    }
    existing.state = 'active'
    existing.lastLeaseAttemptAt = now
    existing.lastStartedAt = now
    existing.lastSeenAt = now
    existing.availableAt = now
  })
}

export function noteMonitorTaskLeaseMiss(taskId: string): void {
  const now = Date.now()

  mutatePersistedState((draft) => {
    const existing = draft.tasks[taskId]
    if (!existing) {
      return
    }

    const leaseMisses = existing.leaseMisses + 1
    if (leaseMisses >= schedulerMaxLeaseMisses) {
      delete draft.tasks[taskId]
      return
    }

    existing.state = 'queued'
    existing.leaseMisses = leaseMisses
    existing.lastLeaseAttemptAt = now
    existing.availableAt = now + leaseMissBackoffMs(leaseMisses)
    existing.lastSeenAt = now
  })
}

export function noteMonitorTaskCompleted(taskId: string): void {
  const now = Date.now()

  mutatePersistedState((draft) => {
    const existing = draft.tasks[taskId]
    if (!existing) {
      return
    }

    if (existing.questionKey) {
      draft.questionCooldowns[existing.questionKey] = now + schedulerQuestionCooldownMs
    }
    draft.platformCooldowns[existing.platform] = now + schedulerPlatformCooldownMs
    delete draft.tasks[taskId]
  })
}

export function noteMonitorTaskCanceled(taskId: string): void {
  mutatePersistedState((draft) => {
    delete draft.tasks[taskId]
  })
}

export function selectNextMonitorTask(
  options: MonitorSchedulerSelectionOptions,
): MonitorSchedulerSelection | null {
  const now = options.now ?? Date.now()
  const draft = clonePersistedState(readPersistedState())
  prunePersistedState(draft, now)

  if (options.activeCount >= options.maxConcurrency) {
    writePersistedState(draft)
    return null
  }

  const candidates = Object.values(draft.tasks)
    .filter((task) => task.state === 'queued' && task.availableAt <= now)
    .sort((left, right) => {
      if (left.laneWeight !== right.laneWeight) {
        return right.laneWeight - left.laneWeight
      }
      if (left.priority !== right.priority) {
        return right.priority - left.priority
      }
      if (left.availableAt !== right.availableAt) {
        return left.availableAt - right.availableAt
      }
      if (left.enqueuedAt !== right.enqueuedAt) {
        return left.enqueuedAt - right.enqueuedAt
      }
      return left.updatedAt - right.updatedAt
    })

  for (const candidate of candidates) {
    if (options.activePlatforms.has(candidate.platform)) {
      continue
    }

    const bypassCooldowns = candidate.lane === 'high'

    const platformCooldown = draft.platformCooldowns[candidate.platform] ?? 0
    if (!bypassCooldowns && platformCooldown > now) {
      continue
    }

    if (!candidate.questionKey && options.activeCount >= schedulerUnknownQuestionParallelismFloor) {
      continue
    }

    if (candidate.questionKey) {
      if (options.activeQuestionKeys.has(candidate.questionKey)) {
        continue
      }
      const questionCooldown = draft.questionCooldowns[candidate.questionKey] ?? 0
      if (!bypassCooldowns && questionCooldown > now) {
        continue
      }
    }

    candidate.state = 'leasing'
    candidate.lastLeaseAttemptAt = now
    candidate.lastSeenAt = now
    writePersistedState(draft)
    return {
      taskId: candidate.taskId,
      routing: candidate.routing,
      source: candidate.source,
    }
  }

  writePersistedState(draft)
  return null
}

export function getMonitorSchedulerSnapshot(): MonitorSchedulerSnapshot {
  const state = readPersistedState()
  const tasks = Object.values(state.tasks).sort((left, right) => {
    if (left.laneWeight !== right.laneWeight) {
      return right.laneWeight - left.laneWeight
    }
    if (left.priority !== right.priority) {
      return right.priority - left.priority
    }
    return left.enqueuedAt - right.enqueuedAt
  })
  return {
    hydratedAt: state.lastHydratedAt || null,
    queueDepth: tasks.length,
    activeCount: tasks.filter((task) => task.state === 'active').length,
    leasingCount: tasks.filter((task) => task.state === 'leasing').length,
    queuedCount: tasks.filter((task) => task.state === 'queued').length,
    questionCooldownCount: Object.keys(state.questionCooldowns).length,
    platformCooldownCount: Object.keys(state.platformCooldowns).length,
    tasks,
  }
}

function leaseMissBackoffMs(leaseMisses: number): number {
  if (leaseMisses <= 1) {
    return 20_000
  }
  if (leaseMisses === 2) {
    return 60_000
  }
  return 3 * 60_000
}

function metadataFromEvent(event: DesktopTaskEventMessage): MonitorTaskMetadata {
  return {
    title: normalizeOptionalString(event.title),
    businessDate: normalizeBusinessDate(event.business_date),
    questionKey: normalizeOptionalString(event.scheduler_group_key),
    questionText: normalizeOptionalString(event.question_text),
  }
}

function metadataFromPayload(payload: Record<string, JsonValue> | null): MonitorTaskMetadata {
  if (!payload) {
    return {
      title: null,
      businessDate: null,
      questionKey: null,
      questionText: null,
    }
  }

  const title = firstString(payload, ['title', 'question_title', 'questionTitle'])
  const businessDate = normalizeBusinessDate(
    firstString(payload, ['business_date', 'businessDate', 'metric_date', 'date']),
  )
  const questionText = firstString(payload, [
    'question_text',
    'questionText',
    'query',
    'question',
    'prompt',
    'content',
  ])
  const questionHash = firstString(payload, ['question_hash', 'questionHash'])
  const questionID = firstString(payload, ['question_id', 'questionId'])

  return {
    title,
    businessDate,
    questionKey:
      normalizeOptionalString(firstString(payload, ['scheduler_group_key', 'schedulerGroupKey'])) ??
      normalizeOptionalString(questionHash) ??
      (questionID ? `qid:${questionID}` : questionText ? `qtxt:${digestText(questionText)}` : null),
    questionText,
  }
}

function firstString(payload: Record<string, JsonValue> | null, keys: string[]): string | null {
  if (!payload) {
    return null
  }
  for (const key of keys) {
    const value = payload[key]
    if (typeof value === 'string') {
      const normalized = value.trim()
      if (normalized) {
        return normalized
      }
      continue
    }
    if (typeof value === 'number' && Number.isFinite(value)) {
      return String(value)
    }
  }
  return null
}

function payloadNumber(payload: Record<string, JsonValue> | null, key: string): number | null {
  if (!payload) {
    return null
  }
  const value = payload[key]
  if (typeof value === 'number' && Number.isFinite(value)) {
    return value
  }
  if (typeof value === 'string' && value.trim()) {
    const parsed = Number.parseInt(value.trim(), 10)
    if (Number.isFinite(parsed)) {
      return parsed
    }
  }
  return null
}

function currentMonitoringBusinessDate(now: number): string {
  const parts = new Intl.DateTimeFormat('en-US', {
    timeZone: 'Asia/Shanghai',
    year: 'numeric',
    month: '2-digit',
    day: '2-digit',
  }).formatToParts(now)

  const year = parts.find((part) => part.type === 'year')?.value ?? ''
  const month = parts.find((part) => part.type === 'month')?.value ?? ''
  const day = parts.find((part) => part.type === 'day')?.value ?? ''
  return `${year}-${month}-${day}`
}

function normalizeBusinessDate(value: unknown): string | null {
  if (typeof value !== 'string') {
    return null
  }
  const normalized = value.trim()
  return /^\d{4}-\d{2}-\d{2}$/.test(normalized) ? normalized : null
}

function normalizeOptionalString(value: unknown): string | null {
  if (typeof value !== 'string') {
    return null
  }
  const normalized = value.trim()
  return normalized || null
}

function normalizePriority(value: unknown, fallback = 100): number {
  if (typeof value === 'number' && Number.isFinite(value)) {
    return value
  }
  return fallback
}

function normalizeLane(value: unknown, fallback: string | null = null): string | null {
  return normalizeOptionalString(value) ?? fallback
}

function laneWeightFromLane(lane: unknown): number {
  switch (normalizeOptionalString(lane)) {
    case 'high':
      return 70
    case 'normal_boosted':
      return 50
    case 'normal':
      return 40
    case 'retry':
      return 20
    default:
      return 0
  }
}

function digestText(value: string): string {
  return createHash('sha1').update(value).digest('hex')
}

function parseTimestamp(value: string | null | undefined): number | null {
  if (!value) {
    return null
  }
  const timestamp = Date.parse(value)
  return Number.isNaN(timestamp) ? null : timestamp
}