platform/src/servers/api/tasks/pipeline-job-manager.ts

import type { ServerWebSocket } from 'bun';
import { randomUUID } from 'crypto';
import { EventEmitter } from 'events';
import {
  createPipelineJob,
  getPipelineJob,
  updatePipelineJob,
  getPipelineJobsForUser,
  getOldestPendingJob,
  countPendingJobs,
  deletePipelineJob,
  deleteTerminalJobsForUser,
  markInterruptedJobs,
  getUserById,
} from 'officerdb';
import { unlink } from 'node:fs/promises';
import { toShellUsername } from '../../data-path';
import { executePipeline } from './pipeline-executor';
import type { OutMessage } from './pipeline-executor';
import { executeScript, jobLogPath, type ScriptEvent } from './execute-script';

// Everything a job can stream — pipeline structural events plus script stdout/stderr/exit.
type JobEvent = OutMessage | ScriptEvent;
type JobMode = 'pipeline' | 'script' | 'agentic';

type WSData = {
  userId: number;
  email: string;
  username: string;
};

type LiveJob = {
  jobId: string;
  userId: number;
  abortSignal: { aborted: boolean };
  emitter: EventEmitter;
  eventBuffer: JobEvent[];
  viewers: Set<ServerWebSocket<WSData>>;
  progressDirty: boolean;
  lastProgress: unknown;
  lastCost: unknown;
};

const EVENT_BUFFER_SIZE = 500;
const PROGRESS_FLUSH_MS = 3000;

const liveJobs = new Map<string, LiveJob>();

function sendToViewer(ws: ServerWebSocket<WSData>, jobId: string, event: JobEvent) {
  if (ws.readyState === 1) {
    ws.send(JSON.stringify({ jobId, ...event }));
  }
}

function broadcast(job: LiveJob, event: JobEvent) {
  // Buffer for replay — skip high-frequency streams (deltas / script stdout+stderr). Those are
  // live-only; a late viewer backfills script output from the persisted log file instead.
  if (event.type !== 'assistant:delta' && event.type !== 'stdout' && event.type !== 'stderr') {
    job.eventBuffer.push(event);
    if (job.eventBuffer.length > EVENT_BUFFER_SIZE) {
      job.eventBuffer.shift();
    }
  }

  // Forward to all viewers
  for (const ws of job.viewers) {
    sendToViewer(ws, job.jobId, event);
  }
}

type StartJobParams = {
  userId: number;
  email: string;
  username: string;
  mode?: JobMode; // defaults to 'pipeline' for back-compat with the existing pipeline caller
  taskDirName: string;
  taskName: string;
  inputs: Record<string, string>;
  cwd?: string;
  config: unknown;
  model?: string;
  startAt?: number;
};

export function runningCount(): number {
  return liveJobs.size;
}

// Create a job. action 'start' runs it now; 'queue' runs it only if nothing is running, else it stays
// 'pending' and gets promoted when the running job finishes. (Single user → one global queue.)
export async function enqueueJob(
  params: StartJobParams,
  action: 'start' | 'queue',
): Promise<{ jobId: string; status: 'running' | 'pending' }> {
  const jobId = randomUUID();
  const mode: JobMode = params.mode ?? 'pipeline';
  const run = action === 'start' || runningCount() === 0;
  await createPipelineJob({
    id: jobId,
    userId: params.userId,
    taskDirName: params.taskDirName,
    taskName: params.taskName,
    mode,
    status: run ? 'running' : 'pending',
    inputs: params.inputs,
    cwd: params.cwd ?? null,
    config: params.config ?? {}, // column is NOT NULL; many script tasks have no config block
    startedAt: run ? new Date() : null,
  });
  if (run) launch(jobId, mode, params);
  return { jobId, status: run ? 'running' : 'pending' };
}

// Back-compat: the pipeline WS caller expects startJob(params) -> jobId, running immediately.
export async function startJob(params: StartJobParams): Promise<string> {
  return (await enqueueJob(params, 'start')).jobId;
}

// Build live state + run the executor for a job whose DB row is already status 'running'.
function launch(jobId: string, mode: JobMode, params: StartJobParams) {
  // Create live state
  const job: LiveJob = {
    jobId,
    userId: params.userId,
    abortSignal: { aborted: false },
    emitter: new EventEmitter(),
    eventBuffer: [],
    viewers: new Set(),
    progressDirty: false,
    lastProgress: null,
    lastCost: null,
  };
  liveJobs.set(jobId, job);

  // Emit callback — broadcasts to viewers and tracks progress
  const emit = (event: JobEvent) => {
    broadcast(job, event);

    // Track progress for DB persistence (skip streaming messages — only structural events)
    if (event.type === 'pipeline:init' || event.type.startsWith('step:') || event.type.startsWith('iteration:')) {
      job.progressDirty = true;
      job.lastProgress = extractProgress(event, job.lastProgress);
    }
    if (event.type === 'step:complete' || event.type === 'iteration:complete') {
      const cost = 'cost' in event ? event.cost : undefined;
      if (cost) {
        const prev = (job.lastCost as { inputTokens: number; outputTokens: number; totalUSD: number } | null) ?? {
          inputTokens: 0,
          outputTokens: 0,
          totalUSD: 0,
        };
        job.lastCost = {
          inputTokens: prev.inputTokens + cost.inputTokens,
          outputTokens: prev.outputTokens + cost.outputTokens,
          totalUSD: prev.totalUSD + cost.totalUSD,
        };
        job.progressDirty = true;
      }
    }
  };

  // Periodic progress flush to DB
  const flushInterval = setInterval(async () => {
    if (job.progressDirty) {
      job.progressDirty = false;
      try {
        await updatePipelineJob(jobId, {
          progress: job.lastProgress as Record<string, unknown>,
          totalCost: job.lastCost as Record<string, unknown>,
        });
      } catch (err) {
        console.error(`[pipeline-jobs] failed to flush progress for ${jobId}:`, err);
      }
    }
  }, PROGRESS_FLUSH_MS);

  // Run the job in the background — dispatch by mode.
  const runner: Promise<void | { exitCode: number }> =
    mode === 'script'
      ? executeScript({
          jobId,
          email: params.email,
          taskDirName: params.taskDirName,
          inputs: params.inputs,
          cwd: params.cwd,
          abortSignal: job.abortSignal,
          emit,
        })
      : executePipeline({
          userId: params.userId,
          email: params.email,
          username: params.username,
          taskDirName: params.taskDirName,
          inputs: params.inputs,
          cwd: params.cwd,
          model: params.model,
          startAt: params.startAt,
          abortSignal: job.abortSignal,
          emit,
        });

  runner
    .then(async (result) => {
      clearInterval(flushInterval);
      // Script jobs resolve with an exit code — a non-zero exit is a failure. Pipelines resolve void.
      const exitCode = result && typeof result === 'object' && 'exitCode' in result ? result.exitCode : null;
      const failed = exitCode !== null && exitCode !== 0;
      await updatePipelineJob(jobId, {
        status: failed ? 'failed' : 'completed',
        exitCode,
        progress: job.lastProgress as Record<string, unknown>,
        totalCost: job.lastCost as Record<string, unknown>,
        error: failed ? `Script exited with code ${exitCode}` : undefined,
        completedAt: new Date(),
      }).catch((err) => console.error(`[pipeline-jobs] failed to finalize ${jobId}:`, err));
      liveJobs.delete(jobId);
      void promoteNext();
    })
    .catch(async (err) => {
      clearInterval(flushInterval);
      const message = err instanceof Error ? err.message : String(err);
      const isStopped = job.abortSignal.aborted;
      broadcast(job, isStopped ? { type: 'stopped' } : { type: 'error', message });
      await updatePipelineJob(jobId, {
        status: isStopped ? 'stopped' : 'failed',
        progress: job.lastProgress as Record<string, unknown>,
        totalCost: job.lastCost as Record<string, unknown>,
        error: isStopped ? undefined : message,
        completedAt: new Date(),
      }).catch((e) => console.error(`[pipeline-jobs] failed to finalize ${jobId}:`, e));
      liveJobs.delete(jobId);
      void promoteNext();
    });
}

// When a job finishes (and nothing else is running), promote the oldest queued job. Also called on
// startup to resume a backlog. Rebuilds the user context from userId since there's no request here.
async function promoteNext(): Promise<void> {
  if (runningCount() > 0) return;
  const next = await getOldestPendingJob();
  if (!next) return;
  const user = await getUserById(next.userId);
  if (!user) {
    await updatePipelineJob(next.id, { status: 'failed', error: 'user not found', completedAt: new Date() }).catch(
      () => {},
    );
    return promoteNext();
  }
  await updatePipelineJob(next.id, { status: 'running', startedAt: new Date() });
  const nextMode = (next.mode as JobMode) ?? 'pipeline';
  launch(next.id, nextMode, {
    userId: next.userId,
    email: user.email,
    username: toShellUsername(user.username ?? '', user.email),
    mode: nextMode,
    taskDirName: next.taskDirName,
    taskName: next.taskName,
    inputs: next.inputs as Record<string, string>,
    cwd: next.cwd ?? undefined,
    config: next.config,
  });
}

export function attachViewer(jobId: string, ws: ServerWebSocket<WSData>) {
  const job = liveJobs.get(jobId);
  if (!job) return false;

  job.viewers.add(ws);

  // Replay buffered events so the viewer catches up
  for (const event of job.eventBuffer) {
    sendToViewer(ws, jobId, event);
  }

  return true;
}

export function detachViewer(jobId: string, ws: ServerWebSocket<WSData>) {
  const job = liveJobs.get(jobId);
  if (job) job.viewers.delete(ws);
}

export function detachAllForViewer(ws: ServerWebSocket<WSData>) {
  for (const job of liveJobs.values()) {
    job.viewers.delete(ws);
  }
}

export function stopJob(jobId: string): boolean {
  const job = liveJobs.get(jobId);
  if (!job) return false;
  job.abortSignal.aborted = true;
  return true;
}

// Stop a running job (cooperative abort) OR cancel a queued one (mark stopped so it won't promote).
export async function requestStop(jobId: string): Promise<'stopped' | 'cancelled' | 'noop'> {
  if (stopJob(jobId)) return 'stopped';
  const job = await getPipelineJob(jobId);
  if (job && job.status === 'pending') {
    await updatePipelineJob(jobId, { status: 'stopped', completedAt: new Date() });
    return 'cancelled';
  }
  return 'noop';
}

export function isJobLive(jobId: string): boolean {
  return liveJobs.has(jobId);
}

export function getLiveJobForUser(userId: number): string | null {
  for (const [jobId, job] of liveJobs) {
    if (job.userId === userId) return jobId;
  }
  return null;
}

export async function getJobsForUser(userId: number) {
  const jobs = await getPipelineJobsForUser(userId);
  return jobs.map((j) => ({
    ...j,
    isLive: liveJobs.has(j.id),
  }));
}

// Delete a job (row + log). Refuses a running job — it must be stopped first.
export async function deleteJob(jobId: string): Promise<'deleted' | 'running' | 'noop'> {
  if (liveJobs.has(jobId)) return 'running';
  const job = await getPipelineJob(jobId);
  if (!job) return 'noop';
  await deletePipelineJob(jobId);
  await unlink(jobLogPath(jobId)).catch(() => {});
  return 'deleted';
}

// Clear all finished jobs (rows + logs). Running/queued jobs are untouched.
export async function clearHistory(userId: number): Promise<number> {
  const ids = await deleteTerminalJobsForUser(userId);
  await Promise.all(ids.map((id) => unlink(jobLogPath(id)).catch(() => {})));
  return ids.length;
}

// Lightweight header-badge summary: how many of the user's jobs are running / queued, and which one
// is running (for the "running" badge's link).
export async function getCounts(
  userId: number,
): Promise<{ running: number; runningJobId: string | null; queued: number }> {
  let running = 0;
  let runningJobId: string | null = null;
  for (const [id, job] of liveJobs) {
    if (job.userId === userId) {
      running++;
      if (!runningJobId) runningJobId = id;
    }
  }
  const queued = await countPendingJobs(userId);
  return { running, runningJobId, queued };
}

export async function getJob(jobId: string) {
  const job = await getPipelineJob(jobId);
  if (!job) return null;
  return { ...job, isLive: liveJobs.has(job.id) };
}

export async function cleanupOnStartup() {
  const count = await markInterruptedJobs();
  if (count > 0) {
    console.log(`[pipeline-jobs] marked ${count} interrupted jobs on startup`);
  }
  // Resume any queued backlog that survived the restart.
  await promoteNext();
}

/** Extract minimal progress info from events for DB storage. Script events fall through to `prev`. */
function extractProgress(event: JobEvent, prev: unknown): unknown {
  const p = (prev as Record<string, unknown>) ?? {};

  switch (event.type) {
    case 'pipeline:init':
      return { ...p, steps: event.steps, currentStepIndex: -1 };

    case 'step:start':
      return { ...p, currentStepIndex: event.stepIndex, currentIteration: event.iteration ?? null };

    case 'step:complete':
      return { ...p, currentStepIndex: event.stepIndex };

    case 'step:parallel':
      return {
        ...p,
        currentStepIndex: event.stepIndex,
        parallel: {
          taskName: event.taskName,
          concurrency: event.concurrency,
          iterations: event.iterations.map((l) => ({ label: l, status: 'pending' })),
        },
      };

    case 'iteration:start':
      if (p.parallel && Array.isArray((p.parallel as Record<string, unknown>).iterations)) {
        const parallel = p.parallel as { iterations: Array<{ label: string; status: string }> };
        return {
          ...p,
          parallel: {
            ...parallel,
            iterations: parallel.iterations.map((it) => (it.label === event.label ? { ...it, status: 'running' } : it)),
          },
        };
      }
      return p;

    case 'iteration:complete':
    case 'iteration:error':
      if (p.parallel && Array.isArray((p.parallel as Record<string, unknown>).iterations)) {
        const parallel = p.parallel as { iterations: Array<{ label: string; status: string }> };
        const status = event.type === 'iteration:complete' ? 'complete' : 'error';
        return {
          ...p,
          parallel: {
            ...parallel,
            iterations: parallel.iterations.map((it) => (it.label === event.label ? { ...it, status } : it)),
          },
        };
      }
      return p;

    default:
      return p;
  }
}