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 { app } from "electron/main";
import type { DesktopTaskEventMessage, DesktopTaskInfo, JsonValue } from "@geo/shared-types";

export type MonitorSchedulerRouting = "rabbitmq-primary" | "db-recovery";
type MonitorTaskState = "queued" | "leasing" | "active";

export interface MonitorSchedulerTaskRecord {
  taskId: string;
  jobId: string;
  accountId: string;
  clientId: string;
  platform: string;
  routing: MonitorSchedulerRouting;
  state: MonitorTaskState;
  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;
}

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 = 1;
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;
  }

  next.version = typeof input.version === "number" ? input.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",
        state: task.state === "active" || task.state === "leasing" ? task.state : "queued",
        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);
  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,
      state: existing?.state === "active" ? "active" : "queued",
      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,
      state: "active",
      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 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.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 platformCooldown = draft.platformCooldowns[candidate.platform] ?? 0;
    if (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 (questionCooldown > now) {
        continue;
      }
    }

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

  writePersistedState(draft);
  return null;
}

export function getMonitorSchedulerSnapshot(): MonitorSchedulerSnapshot {
  const state = readPersistedState();
  const tasks = Object.values(state.tasks).sort((left, right) => 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>,
  keys: string[],
): string | 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 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 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;
}