pi-ext

// ultra — interpolation resolver (Task 3).
//
// A closed, deliberately tiny mini-syntax: NOT a general expression evaluator
// and NO `eval`. Every token is parsed against a fixed grammar; anything
// outside it raises a `SpecValidationError`. See the design's
// "Interpolation mini-syntax" section for the authoritative semantics.
//
//   {args.NAME}                 — workflow input value
//   {item}                      — current fanout item (string verbatim; else compact JSON)
//   {PHASE.results}             — a prior phase's results array (nulls RETAINED positionally)
//   {PHASE.failures}            — a prior phase's typed failure array
//   {PHASE.results[].FIELD}              — project each result's FIELD (null/missing values SKIPPED)
//   {PHASE.results[].FIELD[]}             — flatten: concat each result's FIELD array (null elems / missing FIELDs SKIPPED)
//   {EXPR | where FIELD}                  — keep elements whose boolean FIELD is truthy (null elems / missing FIELD SKIPPED)
//   {EXPR | where FIELD == "VALUE"}       — keep elements whose string FIELD exactly equals VALUE
//   {EXPR | groupBy FIELD}                — group records by exact string FIELD value
//
// Pure module: no SDK imports.

/** Raised when a template/selector falls outside the documented v1 grammar. */
export class SpecValidationError extends Error {
	constructor(message: string) {
		super(message);
		this.name = "SpecValidationError";
	}
}

/**
 * Resolution context. This shape is the contract the phase engine (Task 6)
 * populates: `args` from the workflow inputs, `item` for the current fanout
 * element, and `results` mapping each prior phase id to its results array.
 * The three namespaces are segregated so a phase named `args` or `item`
 * cannot collide with the reserved roots.
 */
export interface InterpSource {
	phase: string;
	index: number;
}

export interface InterpContext {
	args?: Record<string, unknown>;
	item?: unknown;
	itemSources?: readonly InterpSource[];
	results?: Record<string, unknown[]>;
	failures?: Record<string, unknown[]>;
}

export interface SourcedValue {
	value: unknown;
	sources: InterpSource[];
}

interface ResolvedNode extends SourcedValue {
	elements?: ResolvedNode[];
}

const IDENT = "[A-Za-z_][A-Za-z0-9_]*";
const RE_ARGS = new RegExp(`^args\\.(${IDENT})$`);
const RE_RESULTS = new RegExp(`^(${IDENT})\\.results$`);
const RE_FAILURES = new RegExp(`^(${IDENT})\\.failures$`);
const RE_FLATTEN = new RegExp(
	`^(${IDENT})\\.results\\[\\]\\.(${IDENT})\\[\\]$`,
);
const RE_PROJECT = new RegExp(`^(${IDENT})\\.results\\[\\]\\.(${IDENT})$`);
const JSON_STRING = String.raw`"(?:\\.|[^"\\])*"`;
const RE_WHERE = new RegExp(
	`^where\\s+(${IDENT})(?:\\s*==\\s*(${JSON_STRING}))?$`,
);
const RE_GROUP_BY = new RegExp(`^groupBy\\s+(${IDENT})$`);
// A selector value (`over`/`when`/`return`/`report`) is exactly one braced token, no nesting.
const RE_SINGLE_TOKEN = /^\{([^{}]*)\}$/;
// Template token scanner: every `{...}` (no nesting in v1).
const RE_TOKEN = /\{([^{}]*)\}/g;

type Base =
	| { kind: "args"; name: string }
	| { kind: "item" }
	| { kind: "results"; phase: string }
	| { kind: "failures"; phase: string }
	| { kind: "flatten"; phase: string; field: string }
	| { kind: "project"; phase: string; field: string };

type Expr =
	| Base
	| { kind: "where"; base: Base; field: string; equals?: string }
	| { kind: "groupBy"; base: Base; field: string };

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function parseBase(raw: string): Base {
	const s = raw.trim();
	if (s === "item") return { kind: "item" };

	const args = s.match(RE_ARGS);
	if (args) return { kind: "args", name: args[1] };

	const results = s.match(RE_RESULTS);
	if (results) return { kind: "results", phase: results[1] };

	const failures = s.match(RE_FAILURES);
	if (failures) return { kind: "failures", phase: failures[1] };

	const flatten = s.match(RE_FLATTEN);
	if (flatten) {
		return { kind: "flatten", phase: flatten[1], field: flatten[2] };
	}

	const project = s.match(RE_PROJECT);
	if (project) {
		return { kind: "project", phase: project[1], field: project[2] };
	}

	throw new SpecValidationError(
		`unknown interpolation token "{${raw}}" (outside the v1 grammar; allowed: args.NAME, item, PHASE.results, PHASE.failures, PHASE.results[].FIELD, PHASE.results[].FIELD[], and selector operators where/groupBy)`,
	);
}

function selectorPipeIndex(input: string, start = 0): number {
	let quoted = false;
	let escaped = false;
	for (let index = start; index < input.length; index++) {
		const char = input[index];
		if (quoted) {
			if (escaped) escaped = false;
			else if (char === "\\") escaped = true;
			else if (char === '"') quoted = false;
		} else if (char === '"') quoted = true;
		else if (char === "|") return index;
	}
	return -1;
}

function parseExpr(inner: string): Expr {
	const trimmed = inner.trim();
	if (trimmed === "") {
		throw new SpecValidationError('empty interpolation token "{}"');
	}

	const pipeIdx = selectorPipeIndex(trimmed);
	if (pipeIdx === -1) return parseBase(trimmed);

	const left = trimmed.slice(0, pipeIdx).trim();
	const right = trimmed.slice(pipeIdx + 1).trim();
	if (selectorPipeIndex(trimmed, pipeIdx + 1) !== -1) {
		throw new SpecValidationError(
			`only a single selector operator is allowed in "{${trimmed}}"`,
		);
	}
	const base = parseBase(left);
	const groupBy = right.match(RE_GROUP_BY);
	if (groupBy) return { kind: "groupBy", base, field: groupBy[1] };

	const filter = right.match(RE_WHERE);
	if (filter) {
		let equals: string | undefined;
		if (filter[2] !== undefined) {
			try {
				equals = JSON.parse(filter[2]) as string;
			} catch {
				throw new SpecValidationError(
					`where equality requires a valid JSON string in "{${trimmed}}"`,
				);
			}
		}
		return {
			kind: "where",
			base,
			field: filter[1],
			...(equals !== undefined ? { equals } : {}),
		};
	}
	throw new SpecValidationError(
		`unsupported selector operator "${right}" in "{${trimmed}}" (use "| where FIELD", '| where FIELD == "VALUE"', or "| groupBy FIELD")`,
	);
}

function getResults(ctx: InterpContext, phase: string): unknown[] {
	const results = ctx.results?.[phase];
	if (results === undefined) {
		throw new SpecValidationError(
			`interpolation references unknown phase "${phase}"`,
		);
	}
	if (!Array.isArray(results)) {
		throw new SpecValidationError(`phase "${phase}" results is not an array`);
	}
	return results;
}

function getFailures(ctx: InterpContext, phase: string): unknown[] {
	const failures = ctx.failures?.[phase];
	if (failures === undefined) {
		throw new SpecValidationError(
			`interpolation references unknown phase "${phase}"`,
		);
	}
	if (!Array.isArray(failures)) {
		throw new SpecValidationError(`phase "${phase}" failures is not an array`);
	}
	return failures;
}

function uniqueSources(sources: readonly InterpSource[]): InterpSource[] {
	const seen = new Set<string>();
	const out: InterpSource[] = [];
	for (const source of sources) {
		const key = `${source.phase}#${source.index}`;
		if (seen.has(key)) continue;
		seen.add(key);
		out.push(source);
	}
	return out;
}

function node(
	value: unknown,
	sources: readonly InterpSource[] = [],
): ResolvedNode {
	const ownSources = uniqueSources(sources);
	return {
		value,
		sources: ownSources,
		...(Array.isArray(value)
			? { elements: value.map((element) => node(element, ownSources)) }
			: {}),
	};
}

function arrayNode(elements: ResolvedNode[]): ResolvedNode {
	return {
		value: elements.map((element) => element.value),
		sources: uniqueSources(elements.flatMap((element) => element.sources)),
		elements,
	};
}

function resultNode(ctx: InterpContext, phase: string): ResolvedNode {
	return arrayNode(
		getResults(ctx, phase).map((value, index) =>
			node(
				value,
				value === null || value === undefined ? [] : [{ phase, index }],
			),
		),
	);
}

function resolveBaseNode(base: Base, ctx: InterpContext): ResolvedNode {
	switch (base.kind) {
		case "item":
			return node(ctx.item, ctx.itemSources);
		case "args":
			return node(ctx.args?.[base.name]);
		case "results":
			// Retain nulls positionally — callers depend on length.
			return resultNode(ctx, base.phase);
		case "failures":
			// Failure records are emitted by Ultra, not authored by the failed agent.
			return node(getFailures(ctx, base.phase));
		case "flatten": {
			const out: ResolvedNode[] = [];
			for (const element of resultNode(ctx, base.phase).elements ?? []) {
				if (element.value === null || element.value === undefined) continue;
				if (!isRecord(element.value)) continue;
				const field = element.value[base.field];
				// Only arrays contribute; missing / null / non-array FIELDs skipped.
				if (!Array.isArray(field)) continue;
				out.push(...field.map((value) => node(value, element.sources)));
			}
			return arrayNode(out);
		}
		case "project": {
			const out: ResolvedNode[] = [];
			for (const element of resultNode(ctx, base.phase).elements ?? []) {
				if (element.value === null || element.value === undefined) continue;
				if (!isRecord(element.value)) continue;
				const field = element.value[base.field];
				// Missing / null FIELDs contribute no projected value.
				if (field !== undefined && field !== null)
					out.push(node(field, element.sources));
			}
			return arrayNode(out);
		}
	}
}

function resolveExprNode(expr: Expr, ctx: InterpContext): ResolvedNode {
	if (expr.kind === "where") {
		const source = resolveBaseNode(expr.base, ctx);
		if (!Array.isArray(source.value)) {
			throw new SpecValidationError(
				`"| where ${expr.field}" requires an array on its left-hand side`,
			);
		}
		const out: ResolvedNode[] = [];
		for (const element of source.elements ?? []) {
			if (element.value === null || element.value === undefined) continue;
			if (!isRecord(element.value)) continue;
			const value = element.value[expr.field];
			if (
				expr.equals !== undefined
					? typeof value === "string" && value === expr.equals
					: Boolean(value)
			)
				out.push(element);
		}
		return arrayNode(out);
	}
	if (expr.kind === "groupBy") {
		const source = resolveBaseNode(expr.base, ctx);
		if (!Array.isArray(source.value)) {
			throw new SpecValidationError(
				`"| groupBy ${expr.field}" requires an array on its left-hand side`,
			);
		}
		const groups: Array<{
			key: string;
			items: unknown[];
			sources: InterpSource[];
		}> = [];
		const byKey = new Map<string, (typeof groups)[number]>();
		for (const element of source.elements ?? []) {
			if (!isRecord(element.value)) continue;
			const value = element.value[expr.field];
			if (typeof value !== "string") continue;
			const group = byKey.get(value);
			if (group) {
				group.items.push(element.value);
				group.sources = uniqueSources([...group.sources, ...element.sources]);
				continue;
			}
			const created = {
				key: value,
				items: [element.value],
				sources: [...element.sources],
			};
			byKey.set(value, created);
			groups.push(created);
		}
		return arrayNode(
			groups.map((group) =>
				node({ key: group.key, items: group.items }, group.sources),
			),
		);
	}
	return resolveBaseNode(expr, ctx);
}

function resolveExpr(expr: Expr, ctx: InterpContext): unknown {
	return resolveExprNode(expr, ctx).value;
}

/** Render a resolved value into template text: strings verbatim, else compact JSON. */
function renderValue(value: unknown): string {
	if (typeof value === "string") return value;
	if (value === undefined) return "";
	return JSON.stringify(value);
}

/**
 * Substitute every `{...}` token in a prompt template. Each token is resolved
 * against the grammar and rendered (strings verbatim, objects/arrays/scalars as
 * compact JSON). An out-of-grammar token throws.
 */
export function resolveTemplate(
	template: string,
	ctx: InterpContext,
	render: (value: unknown) => string = renderValue,
): string {
	return template.replace(RE_TOKEN, (_match, inner: string) =>
		render(resolveExpr(parseExpr(inner), ctx)),
	);
}

/** Resolve prompt substitutions with immediate producer metadata while leaving
 * ordinary workflow arguments and literal fanout items unattributed. */
export function resolveTemplateWithSources(
	template: string,
	ctx: InterpContext,
	render: (content: string, sources: readonly InterpSource[]) => string,
): string {
	return template.replace(RE_TOKEN, (_match, inner: string) => {
		const resolved = resolveExprNode(parseExpr(inner), ctx);
		return render(renderValue(resolved.value), resolved.sources);
	});
}

/**
 * Resolve a single braced selector expression (`over` / `when` / `return` /
 * `report`) to its raw
 * value. The whole input must be exactly one `{...}` token.
 */
export function resolveSelector(selector: string, ctx: InterpContext): unknown {
	const match = selector.trim().match(RE_SINGLE_TOKEN);
	if (!match) {
		throw new SpecValidationError(
			`selector must be a single brace-wrapped token, got "${selector}"`,
		);
	}
	return resolveExpr(parseExpr(match[1]), ctx);
}

/**
 * Resolve a `fanout` `over`: an array literal is returned as-is; a string is an
 * interpolation expression that must resolve to an array.
 */
export function resolveOver(over: unknown, ctx: InterpContext): unknown[] {
	return resolveOverWithSources(over, ctx).map((item) => item.value);
}

/** Resolve fanout items together with immediate producer metadata. */
export function resolveOverWithSources(
	over: unknown,
	ctx: InterpContext,
): SourcedValue[] {
	if (Array.isArray(over)) return over.map((value) => ({ value, sources: [] }));
	if (typeof over === "string") {
		const match = over.trim().match(RE_SINGLE_TOKEN);
		if (!match) {
			throw new SpecValidationError(
				`selector must be a single brace-wrapped token, got "${over}"`,
			);
		}
		const resolved = resolveExprNode(parseExpr(match[1]), ctx);
		if (!Array.isArray(resolved.value)) {
			throw new SpecValidationError(
				`"over" interpolation "${over}" did not resolve to an array`,
			);
		}
		return (resolved.elements ?? []).map(({ value, sources }) => ({
			value,
			sources,
		}));
	}
	throw new SpecValidationError(
		`"over" must be an array literal or an interpolation string, got ${typeof over}`,
	);
}