maestro/ui/src/components/chat/movementHistory.ts

import type { LocalTaskComment, MovementHistoryEvent } from '../../api';

export type MovementRailItem =
  | { id: string; type: 'movement'; name: string; status: 'completed' | 'active'; ts: string; targetId?: string }
  | { id: string; type: 'return'; name: string; ts: string; targetId?: string }
  | { id: string; type: 'llm_retry'; name: string; ts: string; targetId: string; detail: string }
  | { id: string; type: 'job_retry'; name: string; ts: string; targetId: string; detail: string }
  | { id: string; type: 'user_input'; name: string; ts: string; targetId: string; detail: string };

interface JobRetryData {
  type: 'job_retry';
  disposition: 'retry' | 'requeued_unhealthy';
  attempt: number;
  nextAttempt: number;
  maxAttempts: number;
}

export const parseJobRetry = (body: string): JobRetryData | null => {
  try {
    const data = JSON.parse(body) as Partial<JobRetryData>;
    if (data.type !== 'job_retry'
      || (data.disposition !== 'retry' && data.disposition !== 'requeued_unhealthy')
      || typeof data.attempt !== 'number' || !Number.isFinite(data.attempt)
      || typeof data.nextAttempt !== 'number' || !Number.isFinite(data.nextAttempt)
      || typeof data.maxAttempts !== 'number' || !Number.isFinite(data.maxAttempts)) return null;
    return data as JobRetryData;
  } catch {
    return null;
  }
};

const userComment = (comment: LocalTaskComment): boolean =>
  ['request', 'comment', 'interjection'].includes(comment.kind)
  && comment.author !== 'agent'
  && comment.author !== 'system';

export const buildMovementRailItems = (
  events: MovementHistoryEvent[],
  comments: LocalTaskComment[],
  movementNames: string[],
): MovementRailItem[] => {
  const order = new Map(movementNames.map((name, index) => [name, index]));
  const items: MovementRailItem[] = [];
  let previousMovement: string | null = null;
  const completionTargets = new Map<string, string[]>();
  for (const comment of comments) {
    try {
      const data = JSON.parse(comment.body) as { movement?: unknown; durationMs?: unknown };
      if (comment.kind !== 'progress' || typeof data.movement !== 'string' || typeof data.durationMs !== 'number') continue;
      const targets = completionTargets.get(data.movement) ?? [];
      targets.push(`movement-completion-${comment.id}`);
      completionTargets.set(data.movement, targets);
    } catch { /* not a movement completion */ }
  }
  const occurrence = new Map<string, number>();
  const openMovements = new Map<string, number[]>();

  const sortedEvents = [...events].sort((a, b) =>
    a.ts.localeCompare(b.ts) || (a.runId === b.runId ? a.seq - b.seq : a.line - b.line));
  for (const event of sortedEvents) {
    const name = event.movement ?? 'unknown';
    if (event.kind === 'movement_start') {
      const index = occurrence.get(name) ?? 0;
      occurrence.set(name, index + 1);
      const movementTarget = completionTargets.get(name)?.[index];
      const previousIndex = previousMovement ? order.get(previousMovement) : undefined;
      const nextIndex = order.get(name);
      if (previousIndex !== undefined && nextIndex !== undefined && nextIndex < previousIndex) {
        items.push({ id: `return-${event.eventId}`, type: 'return', name, ts: event.ts, targetId: movementTarget });
      }
      const itemIndex = items.length;
      items.push({ id: event.eventId, type: 'movement', name, status: 'active', ts: event.ts, targetId: movementTarget });
      const runMovement = `${event.runId}:${name}`;
      const open = openMovements.get(runMovement) ?? [];
      open.push(itemIndex);
      openMovements.set(runMovement, open);
      previousMovement = name;
    } else if (event.kind === 'movement_complete') {
      const index = Math.max(0, (occurrence.get(name) ?? 1) - 1);
      const open = openMovements.get(`${event.runId}:${name}`);
      const itemIndex = open?.shift();
      if (itemIndex !== undefined) {
        const started = items[itemIndex];
        if (started?.type === 'movement') {
          items[itemIndex] = { ...started, status: 'completed', targetId: completionTargets.get(name)?.[index] ?? started.targetId };
        }
      } else {
        items.push({ id: event.eventId, type: 'movement', name, status: 'completed', ts: event.ts, targetId: completionTargets.get(name)?.[index] });
      }
    } else {
      const attempt = event.payload.attempt ?? '?';
      const max = event.payload.maxAttempts ?? '?';
      const error = event.payload.errorClass ?? (event.payload.httpStatus ? `HTTP ${event.payload.httpStatus}` : 'LLM error');
      items.push({ id: event.eventId, type: 'llm_retry', name, ts: event.ts, targetId: `trace-event-${event.eventId}`, detail: `${attempt}/${max} · ${error}` });
    }
  }

  for (const comment of comments) {
    const retry = comment.kind === 'progress' && comment.author === 'system' ? parseJobRetry(comment.body) : null;
    if (retry) {
      const detail = retry.disposition === 'requeued_unhealthy'
        ? 'Worker connection retry'
        : `${retry.nextAttempt}/${retry.maxAttempts}`;
      items.push({ id: `comment-${comment.id}`, type: 'job_retry', name: 'Retry', ts: comment.createdAt, targetId: `comment-${comment.id}`, detail });
    } else if (userComment(comment)) {
      items.push({
        id: `comment-${comment.id}`,
        type: 'user_input',
        name: comment.kind === 'request' ? 'Request' : 'Comment',
        ts: comment.createdAt,
        targetId: `comment-${comment.id}`,
        detail: comment.body.replace(/\s+/g, ' ').trim().slice(0, 80),
      });
    }
  }
  return items.sort((a, b) => a.ts.localeCompare(b.ts));
};

export const assignRetryEvents = (
  events: MovementHistoryEvent[],
  completions: Array<{ id: number; movement: string }>,
): { byCompletionId: Map<number, MovementHistoryEvent[]>; unassigned: MovementHistoryEvent[] } => {
  const targets = new Map<string, number[]>();
  for (const completion of completions) {
    const ids = targets.get(completion.movement) ?? [];
    ids.push(completion.id);
    targets.set(completion.movement, ids);
  }
  const byCompletionId = new Map<number, MovementHistoryEvent[]>();
  const active = new Map<string, MovementHistoryEvent[]>();
  const unassigned: MovementHistoryEvent[] = [];
  const sorted = [...events].sort((a, b) => a.ts.localeCompare(b.ts) || a.line - b.line);
  for (const event of sorted) {
    const key = `${event.runId}:${event.movement ?? 'unknown'}`;
    if (event.kind === 'movement_start') {
      active.set(key, []);
    } else if (event.kind === 'llm_call_retry') {
      const retries = active.get(key);
      if (retries) retries.push(event);
      else unassigned.push(event);
    } else if (event.kind === 'movement_complete') {
      const completionId = targets.get(event.movement ?? 'unknown')?.shift();
      const retries = active.get(key) ?? [];
      if (completionId !== undefined) byCompletionId.set(completionId, retries);
      else unassigned.push(...retries);
      active.delete(key);
    }
  }
  for (const retries of active.values()) unassigned.push(...retries);
  return { byCompletionId, unassigned };
};