platform/src/servers/api/tasks/pipeline-executor.ts

import type { ServerWebSocket } from 'bun';
import { randomUUID } from 'crypto';
import { readdirSync, existsSync, mkdirSync, writeFileSync, chmodSync, rmSync } from 'node:fs';
import { join } from 'node:path';
import { tmpdir } from 'node:os';
import { getUserSettings } from 'officerdb';
import { getTaskByDirName } from './task-files';
import { getHomeDir, getOwnerHomeDir } from '../../data-path';
import { resolveBaseCwd } from '../chat/websocket';
import { sendClaudeCodeStreaming } from '../../channels/send-claude-code';
import type { TurnMessage, MessageCost } from '../chat/types';
import * as jobManager from './pipeline-job-manager';

const DEFAULT_MODEL = 'claude-code';

async function resolveModel(userId: number): Promise<string> {
  try {
    const settings = await getUserSettings(userId);
    const tasks = settings?.tasks as Record<string, unknown> | undefined;
    return (tasks?.defaultModel as string) || DEFAULT_MODEL;
  } catch {
    return DEFAULT_MODEL;
  }
}

// ── Types ──

type PipelineStep = {
  task: string;
  inputs?: Record<string, string>;
  foreach?: 'subdirectory';
  skip_if?: string;
  concurrency?: string | boolean;
};

type PipelineConfig = {
  steps: PipelineStep[];
};

// Messages sent to client
export type OutMessage =
  | { type: 'pipeline:init'; steps: Array<{ task: string; foreach?: string; concurrency?: number }> }
  | {
      type: 'step:start';
      stepIndex: number;
      taskName: string;
      iteration?: { current: number; total: number; label: string };
    }
  | { type: 'step:complete'; stepIndex: number; cost?: MessageCost }
  | { type: 'step:skip'; stepIndex: number; label: string; reason: string }
  | { type: 'step:parallel'; stepIndex: number; taskName: string; iterations: string[]; concurrency: number }
  | { type: 'step:waiting'; stepIndex: number; iterationLabel?: string; elapsed: number }
  | { type: 'iteration:start'; stepIndex: number; label: string }
  | { type: 'iteration:complete'; stepIndex: number; label: string; cost?: MessageCost }
  | { type: 'iteration:error'; stepIndex: number; label: string; error: string }
  | { type: 'assistant:delta'; text: string; stepIndex: number; iterationLabel?: string }
  | { type: 'assistant:text'; text: string; stepIndex: number; iterationLabel?: string }
  | {
      type: 'tool:start';
      toolCallId: string;
      toolName: string;
      toolInput: Record<string, unknown>;
      stepIndex: number;
      iterationLabel?: string;
    }
  | {
      type: 'tool:result';
      toolCallId: string;
      output: string;
      isError: boolean;
      stepIndex: number;
      iterationLabel?: string;
    }
  | { type: 'pipeline:complete'; totalCost: MessageCost }
  | { type: 'error'; message: string }
  | { type: 'stopped' };

type EmitEvent = (msg: OutMessage) => void;

type AbortSignal = { aborted: boolean };

// ── Step execution ──

type RunStepParams = {
  userId: number;
  email: string;
  username: string;
  taskDirName: string;
  prompt: string;
  cwd: string;
  model: string;
  abortSignal: AbortSignal;
  emit: EmitEvent;
  stepIndex: number;
  iterationLabel?: string;
};

async function refreshProxyToken(): Promise<void> {
  const port = process.env.ANTHROPIC_PROXY_PORT ?? '5051';
  try {
    await fetch(`http://127.0.0.1:${port}/refresh`, { method: 'POST' });
  } catch {
    // Best effort — proxy may not be running (e.g. using API key directly)
  }
}

// How long (ms) to wait with zero activity before timing out an agentic step.
// Only applies to non-Claude-Code models (local LLMs etc).
const ACTIVITY_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
const WAITING_INTERVAL_MS = 10 * 1000; // emit "waiting" every 10s

async function runAgenticStep({
  userId,
  email,
  username,
  taskDirName,
  prompt,
  cwd,
  model,
  abortSignal,
  emit,
  stepIndex,
  iterationLabel,
}: RunStepParams): Promise<MessageCost> {
  const sessionId = randomUUID();
  const isClaudeCode = model.startsWith('claude-code');

  // Ensure fresh OAuth token before spawning Claude Code
  if (isClaudeCode) await refreshProxyToken();

  console.log(
    `[pipeline] starting step ${stepIndex} for session ${sessionId} (model=${model})${iterationLabel ? ` [${iterationLabel}]` : ''}`,
  );

  return new Promise<MessageCost>(async (resolve, reject) => {
    if (abortSignal.aborted) return reject(new Error('Pipeline aborted'));

    let cleanup: (() => void) | null = null;
    let settled = false;
    let lastActivity = Date.now();
    const stepStart = Date.now();

    const settle = (fn: () => void) => {
      if (settled) return;
      settled = true;
      clearInterval(abortPoll);
      clearInterval(waitingPoll);
      fn();
    };

    // The agent sidecar now hands over finished turn messages rather than raw parser events, so this is a
    // re-label onto the pipeline's own event stream. The `seq` is ignored: pipeline steps have their own
    // durable record (the job's events), not the chat cursor.
    const onMessage = (msg: TurnMessage) => {
      if (abortSignal.aborted) return;
      lastActivity = Date.now();

      switch (msg.type) {
        case 'assistant:delta':
          emit({ type: 'assistant:delta', text: msg.text, stepIndex, iterationLabel });
          break;
        case 'assistant:text':
          emit({ type: 'assistant:text', text: msg.text, stepIndex, iterationLabel });
          break;
        case 'tool:start':
          emit({
            type: 'tool:start',
            toolCallId: msg.toolCallId,
            toolName: msg.toolName,
            toolInput: msg.toolInput,
            stepIndex,
            iterationLabel,
          });
          break;
        case 'tool:result':
          emit({
            type: 'tool:result',
            toolCallId: msg.toolCallId,
            output: msg.output,
            isError: msg.isError,
            stepIndex,
            iterationLabel,
          });
          break;
        case 'result': {
          const cost = msg.cost;
          settle(() => {
            cleanup?.();
            resolve(cost);
          });
          break;
        }
        case 'error': {
          const message = msg.message;
          settle(() => {
            cleanup?.();
            reject(new Error(message));
          });
          break;
        }
        case 'stopped':
          settle(() => {
            cleanup?.();
            reject(new Error('Step was stopped'));
          });
          break;
      }
    };

    // Poll for abort signal and activity timeout
    const abortPoll = setInterval(() => {
      if (abortSignal.aborted) {
        settle(() => {
          cleanup?.();
          reject(new Error('Pipeline was stopped'));
        });
        return;
      }
      // Activity timeout (skip for Claude Code which has its own mechanisms)
      if (!isClaudeCode && Date.now() - lastActivity > ACTIVITY_TIMEOUT_MS) {
        const elapsed = Math.round((Date.now() - stepStart) / 1000);
        console.error(`[pipeline] step ${stepIndex} timed out after ${elapsed}s of inactivity (session=${sessionId})`);
        settle(() => {
          cleanup?.();
          reject(new Error(`Step timed out — no response from model for ${Math.round(ACTIVITY_TIMEOUT_MS / 1000)}s`));
        });
      }
    }, 500);

    // Emit periodic "waiting" events so the frontend knows the step is alive
    const waitingPoll = setInterval(() => {
      if (settled) return;
      const elapsed = Math.round((Date.now() - stepStart) / 1000);
      emit({ type: 'step:waiting', stepIndex, iterationLabel, elapsed });
    }, WAITING_INTERVAL_MS);

    try {
      const handle = await sendClaudeCodeStreaming({
        userId,
        email,
        username,
        prompt,
        sessionKey: sessionId,
        cwd,
        model,
        durable: false,
        onMessage,
      });
      cleanup = handle.kill;
    } catch (err) {
      settle(() => {
        cleanup?.();
        reject(err);
      });
    }
  });
}

// ── Helpers ──

function resolveInputTemplate(template: string, variables: Record<string, string>): string {
  return template.replace(/\$\{(\w+)\}/g, (_, key) => variables[key] ?? '');
}

function buildStepPrompt(taskBody: string, inputs: Record<string, string>, targetDir?: string): string {
  const inputLines = Object.entries(inputs)
    .filter(([, v]) => v.trim())
    .map(([key, value]) => `- **${key}**: ${value}`)
    .join('\n');

  const contextLines: string[] = [];
  if (targetDir) contextLines.push(`- **Target directory**: ${targetDir}`);
  const contextSection = contextLines.length > 0 ? `\n\n## Context\n\n${contextLines.join('\n')}` : '';

  return `${taskBody.trim()}\n\n## Inputs\n\n${inputLines}${contextSection}`;
}

// ── Script step execution ──

type RunScriptStepParams = {
  email: string;
  task: { name: string; implementation: string; language: string; args?: string[] | null };
  inputs: Record<string, string>;
  cwd: string;
  abortSignal: AbortSignal;
  emit: EmitEvent;
  stepIndex: number;
};

function getRunner(language: string): string[] {
  switch (language) {
    case 'bash':
      return ['bash'];
    case 'python':
      return ['python3'];
    case 'typescript':
      return ['bun', 'run'];
    case 'javascript':
      return ['node'];
    default:
      return ['bash'];
  }
}

function getFileName(language: string): string {
  switch (language) {
    case 'bash':
      return 'run.sh';
    case 'python':
      return 'run.py';
    case 'typescript':
      return 'index.ts';
    case 'javascript':
      return 'index.js';
    default:
      return 'run.sh';
  }
}

async function runScriptStep({
  email,
  task,
  inputs,
  cwd,
  abortSignal,
  emit,
  stepIndex,
}: RunScriptStepParams): Promise<void> {
  const language = task.language ?? 'bash';

  // Write script to temp file
  const dir = join(tmpdir(), `officer-task-${Date.now()}-${Math.random().toString(36).slice(2)}`);
  mkdirSync(dir, { recursive: true });
  const fileName = getFileName(language);
  const scriptPath = join(dir, fileName);
  writeFileSync(scriptPath, task.implementation);
  chmodSync(scriptPath, 0o755);

  const cleanup = () => {
    try {
      rmSync(dir, { recursive: true, force: true });
    } catch {
      /* best effort */
    }
  };

  // Build env vars from inputs
  const inputEnv: Record<string, string> = {};
  for (const [key, value] of Object.entries(inputs)) {
    inputEnv[`INPUT_${key.toUpperCase()}`] = value;
  }

  // Build positional args
  const positionalArgs = task.args?.map((name) => inputs[name] ?? '') ?? [];

  const runner = getRunner(language);
  const cmd = [...runner, scriptPath, ...positionalArgs];

  const spawnEnv = { ...(process.env as Record<string, string>), ...inputEnv };

  console.log(`[pipeline] running script step ${stepIndex}: ${task.name} (cwd=${cwd})`);

  try {
    const proc = Bun.spawn(cmd, {
      cwd,
      env: spawnEnv,
      stdout: 'pipe',
      stderr: 'pipe',
    });

    const decoder = new TextDecoder();

    const readStream = async (reader: ReadableStreamDefaultReader<Uint8Array>) => {
      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
          const text = decoder.decode(value);
          emit({ type: 'assistant:delta', text, stepIndex });
        }
      } catch {
        // stream closed
      }
    };

    // Check abort periodically
    const abortCheck = setInterval(() => {
      if (abortSignal.aborted) {
        try {
          proc.kill();
        } catch {
          /* already dead */
        }
      }
    }, 500);

    const [, , exitCode] = await Promise.all([
      readStream(proc.stdout.getReader()),
      readStream(proc.stderr.getReader()),
      proc.exited,
    ]);

    clearInterval(abortCheck);
    cleanup();

    // Flush the streamed text as a complete message
    emit({ type: 'assistant:text', text: '', stepIndex });

    if (exitCode !== 0 && !abortSignal.aborted) {
      throw new Error(`Script exited with code ${exitCode}`);
    }
  } catch (err) {
    cleanup();
    throw err;
  }
}

// ── Parallel foreach ──

type ForeachParams = {
  userId: number;
  email: string;
  username: string;
  stepIdx: number;
  step: PipelineStep;
  stepTask: { name: string; body: string };
  subdirs: string[];
  baseCwd: string;
  inputs: Record<string, string>;
  cwd?: string;
  abortSignal: AbortSignal;
  totalCost: MessageCost;
  emit: EmitEvent;
  concurrency: number;
  model: string;
};

async function runForeach({
  userId,
  email,
  username,
  stepIdx,
  step,
  stepTask,
  subdirs,
  baseCwd,
  inputs,
  cwd,
  abortSignal,
  totalCost,
  emit,
  concurrency,
  model,
}: ForeachParams) {
  // Determine skip vs run
  const toSkip: string[] = [];
  const toRun: string[] = [];
  for (const subdir of subdirs) {
    const subdirPath = join(baseCwd, subdir);
    if (step.skip_if && existsSync(join(subdirPath, step.skip_if))) {
      toSkip.push(subdir);
    } else {
      toRun.push(subdir);
    }
  }

  // Send skips
  for (const label of toSkip) {
    emit({ type: 'step:skip', stepIndex: stepIdx, label, reason: `${step.skip_if} already exists` });
  }

  // Tell UI about parallel step
  emit({
    type: 'step:parallel',
    stepIndex: stepIdx,
    taskName: stepTask.name,
    iterations: toRun,
    concurrency,
  });

  // Pool-based execution
  const executing = new Set<Promise<void>>();

  for (const subdir of toRun) {
    if (abortSignal.aborted) break;

    const run = async () => {
      emit({ type: 'iteration:start', stepIndex: stepIdx, label: subdir });

      const iterVars = { ...inputs, folder_name: subdir };
      const iterInputs: Record<string, string> = {};
      if (step.inputs) {
        for (const [key, template] of Object.entries(step.inputs)) {
          iterInputs[key] = resolveInputTemplate(template, iterVars);
        }
      }

      const cwdRelative = cwd ? `${cwd}/${subdir}` : subdir;
      const resolvedCwd = resolveBaseCwd(getOwnerHomeDir(email), cwdRelative);
      const targetDir = resolvedCwd;
      const prompt = buildStepPrompt(stepTask.body!, iterInputs, targetDir);

      try {
        const cost = await runAgenticStep({
          userId,
          email,
          username,
          taskDirName: step.task,
          prompt,
          cwd: resolvedCwd,
          model,
          abortSignal,
          emit,
          stepIndex: stepIdx,
          iterationLabel: subdir,
        });

        totalCost.inputTokens += cost.inputTokens;
        totalCost.outputTokens += cost.outputTokens;
        totalCost.totalUSD += cost.totalUSD;

        emit({ type: 'iteration:complete', stepIndex: stepIdx, label: subdir, cost });
      } catch (err) {
        if (!abortSignal.aborted) {
          emit({
            type: 'iteration:error',
            stepIndex: stepIdx,
            label: subdir,
            error: err instanceof Error ? err.message : String(err),
          });
        }
      }
    };

    const p = run().then(() => {
      executing.delete(p);
    });
    executing.add(p);

    if (executing.size >= concurrency) {
      await Promise.race(executing);
    }
  }

  await Promise.all(executing);
}

// ── Pipeline orchestrator ──

export type ExecutePipelineParams = {
  userId: number;
  email: string;
  username: string;
  taskDirName: string;
  inputs: Record<string, string>;
  cwd?: string;
  model?: string;
  startAt?: number;
  abortSignal: AbortSignal;
  emit: EmitEvent;
};

export async function executePipeline({
  userId,
  email,
  username,
  taskDirName,
  inputs,
  cwd,
  model: modelOverride,
  startAt,
  abortSignal,
  emit,
}: ExecutePipelineParams): Promise<void> {
  const pipelineTask = await getTaskByDirName(taskDirName);
  if (!pipelineTask) {
    emit({ type: 'error', message: `Task not found: ${taskDirName}` });
    return;
  }
  if (pipelineTask.mode !== 'pipeline') {
    emit({ type: 'error', message: 'Task is not a pipeline-mode task' });
    return;
  }

  const config = pipelineTask.config as PipelineConfig | null;
  if (!config?.steps?.length) {
    emit({ type: 'error', message: 'Pipeline has no steps defined' });
    return;
  }

  const baseCwd = resolveBaseCwd(getOwnerHomeDir(email), cwd);
  let model = modelOverride || (await resolveModel(userId));
  // Claude-only: coerce any legacy non-Claude task-model preference to the Claude default.
  if (!model.startsWith('claude-code')) model = DEFAULT_MODEL;
  const totalCost: MessageCost = { inputTokens: 0, outputTokens: 0, totalUSD: 0 };

  // Resolve concurrency from user input (default 1)
  const runtimeConcurrency = Math.max(1, parseInt(inputs._concurrency ?? '1', 10));

  emit({
    type: 'pipeline:init',
    steps: config.steps.map((s) => ({
      task: s.task,
      foreach: s.foreach,
      concurrency: s.concurrency ? runtimeConcurrency : undefined,
    })),
  });

  for (let stepIdx = 0; stepIdx < config.steps.length; stepIdx++) {
    if (abortSignal.aborted) break;

    // Skip steps before startAt
    if (startAt && stepIdx < startAt) {
      emit({ type: 'step:skip', stepIndex: stepIdx, label: config.steps[stepIdx]!.task, reason: 'skipped (start at)' });
      continue;
    }

    const step = config.steps[stepIdx]!;

    const stepTask = await getTaskByDirName(step.task);
    if (!stepTask) {
      emit({ type: 'error', message: `Step task not found: ${step.task}` });
      return;
    }

    const isScript = stepTask.mode === 'script';
    if (!isScript && !stepTask.body) {
      emit({ type: 'error', message: `Step task "${step.task}" has no body` });
      return;
    }

    const resolvedInputs: Record<string, string> = {};
    if (step.inputs) {
      for (const [key, template] of Object.entries(step.inputs)) {
        resolvedInputs[key] = resolveInputTemplate(template, inputs);
      }
    }

    // Script-mode steps run directly (no agent), only single execution supported
    if (isScript) {
      if (!stepTask.implementation) {
        emit({ type: 'error', message: `Script task "${step.task}" has no implementation` });
        return;
      }

      emit({ type: 'step:start', stepIndex: stepIdx, taskName: stepTask.name });

      try {
        await runScriptStep({
          email,
          task: {
            name: stepTask.name,
            implementation: stepTask.implementation,
            language: stepTask.language ?? 'bash',
            args: stepTask.args as string[] | null,
          },
          inputs: resolvedInputs,
          cwd: baseCwd,
          abortSignal,
          emit,
          stepIndex: stepIdx,
        });
        emit({ type: 'step:complete', stepIndex: stepIdx });
      } catch (err) {
        const message = err instanceof Error ? err.message : String(err);
        emit({ type: 'error', message: `Script step "${step.task}" failed: ${message}` });
        return;
      }

      continue;
    }

    if (step.foreach === 'subdirectory') {
      let subdirs: string[];
      try {
        subdirs = readdirSync(baseCwd, { withFileTypes: true })
          .filter((d) => d.isDirectory())
          .map((d) => d.name)
          .filter((name) => !name.startsWith('.'))
          .sort();
      } catch {
        emit({ type: 'error', message: `Cannot read directory: ${baseCwd}` });
        return;
      }

      const concurrency = step.concurrency ? runtimeConcurrency : 1;

      await runForeach({
        userId,
        email,
        username,
        stepIdx,
        step,
        stepTask: { name: stepTask.name, body: stepTask.body! },
        subdirs,
        baseCwd,
        inputs,
        cwd,
        abortSignal,
        totalCost,
        emit,
        concurrency,
        model,
      });
    } else {
      // Single execution step
      emit({ type: 'step:start', stepIndex: stepIdx, taskName: stepTask.name });

      const targetDir = baseCwd;
      const prompt = buildStepPrompt(stepTask.body!, resolvedInputs, targetDir);

      const cost = await runAgenticStep({
        userId,
        email,
        username,
        taskDirName: step.task,
        prompt,
        cwd: baseCwd,
        model,
        abortSignal,
        emit,
        stepIndex: stepIdx,
      });

      totalCost.inputTokens += cost.inputTokens;
      totalCost.outputTokens += cost.outputTokens;
      totalCost.totalUSD += cost.totalUSD;

      emit({ type: 'step:complete', stepIndex: stepIdx, cost });
    }
  }

  if (abortSignal.aborted) {
    throw new Error('Pipeline was stopped');
  }

  emit({ type: 'pipeline:complete', totalCost });
}

// ── WebSocket handler (thin layer) ──

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

type ClientMessage =
  | { type: 'run'; taskDirName: string; inputs: Record<string, string>; cwd?: string; model?: string; startAt?: number }
  | { type: 'stop'; jobId: string }
  | { type: 'attach'; jobId: string }
  | { type: 'list' };

function send(ws: ServerWebSocket<WSData>, msg: Record<string, unknown>) {
  if (ws.readyState === 1) ws.send(JSON.stringify(msg));
}

export function open(_ws: ServerWebSocket<WSData>) {}

export async function message(ws: ServerWebSocket<WSData>, raw: string | Buffer) {
  const data = typeof raw === 'string' ? raw : raw.toString();

  try {
    const msg = JSON.parse(data) as ClientMessage;

    switch (msg.type) {
      case 'run': {
        const { userId, email, username } = ws.data;

        // Resolve task name for the DB record
        const task = await getTaskByDirName(msg.taskDirName);
        if (!task) {
          send(ws, { type: 'error', message: `Task not found: ${msg.taskDirName}` });
          return;
        }

        const jobId = await jobManager.startJob({
          userId,
          email,
          username,
          taskDirName: msg.taskDirName,
          taskName: task.name,
          inputs: msg.inputs,
          cwd: msg.cwd,
          config: task.config,
          model: msg.model,
          startAt: msg.startAt,
        });

        send(ws, { type: 'job:created', jobId });
        jobManager.attachViewer(jobId, ws);
        break;
      }

      case 'stop': {
        jobManager.stopJob(msg.jobId);
        break;
      }

      case 'attach': {
        const attached = jobManager.attachViewer(msg.jobId, ws);
        if (!attached) {
          // Job not live — send the DB state
          const job = await jobManager.getJob(msg.jobId);
          if (job) {
            send(ws, {
              type: 'job:state',
              jobId: msg.jobId,
              status: job.status,
              progress: job.progress,
              cost: job.totalCost,
            });
          } else {
            send(ws, { type: 'error', message: `Job not found: ${msg.jobId}` });
          }
        }
        break;
      }

      case 'list': {
        const jobs = await jobManager.getJobsForUser(ws.data.userId);
        send(ws, {
          type: 'job:list',
          jobs: jobs.map((j) => ({
            id: j.id,
            taskDirName: j.taskDirName,
            taskName: j.taskName,
            status: j.status,
            isLive: j.isLive,
            totalCost: j.totalCost,
            createdAt: j.createdAt,
            completedAt: j.completedAt,
          })),
        });
        break;
      }
    }
  } catch {
    send(ws, { type: 'error', message: 'Failed to parse message' });
  }
}

export function close(ws: ServerWebSocket<WSData>) {
  // Detach from all jobs but don't stop them
  jobManager.detachAllForViewer(ws);
}

export const pipelineWebsocket = {
  open,
  message,
  close,
  drain() {},
};