storkit/frontend/src/components/LozengeFlyContext.tsx

/**
 * LozengeFlyContext – FLIP-style animation system for agent lozenges.
 *
 * When an agent is assigned to a story, a fixed-positioned clone of the
 * agent lozenge "flies" from the roster badge in AgentPanel to the slot
 * in StagePanel (or vice-versa when the agent is removed).  The overlay
 * travels above all other UI elements (z-index 9999) so it is never
 * clipped by the layout.
 */
import * as React from "react";
import * as ReactDOM from "react-dom";
import type { PipelineState } from "../api/client";

const {
	createContext,
	useCallback,
	useContext,
	useEffect,
	useLayoutEffect,
	useMemo,
	useRef,
	useState,
} = React;

// ─── Public context shape ─────────────────────────────────────────────────────

export interface LozengeFlyContextValue {
	/** Register/unregister a roster badge DOM element by agent name. */
	registerRosterEl: (agentName: string, el: HTMLElement | null) => void;
	/**
	 * Save the latest DOMRect for a story's lozenge slot.
	 * Called on every render of AgentLozenge via useLayoutEffect.
	 */
	saveSlotRect: (storyId: string, rect: DOMRect) => void;
	/**
	 * Set of storyIds whose slot lozenges should be hidden because a
	 * fly-in animation is in progress.
	 */
	pendingFlyIns: ReadonlySet<string>;
	/**
	 * Set of agent names whose roster badge should be hidden.
	 * An agent is hidden while it is assigned to a work item OR while its
	 * fly-out animation (work item → roster) is still in flight.
	 */
	hiddenRosterAgents: ReadonlySet<string>;
}

const noop = () => {};
const emptySet: ReadonlySet<string> = new Set();

export const LozengeFlyContext = createContext<LozengeFlyContextValue>({
	registerRosterEl: noop,
	saveSlotRect: noop,
	pendingFlyIns: emptySet,
	hiddenRosterAgents: emptySet,
});

// ─── Internal flying-lozenge state ───────────────────────────────────────────

interface FlyingLozenge {
	id: string;
	label: string;
	isActive: boolean;
	startX: number;
	startY: number;
	endX: number;
	endY: number;
	/** false = positioned at start, true = CSS transition to end */
	flying: boolean;
}

interface PendingFlyIn {
	storyId: string;
	agentName: string;
	label: string;
	isActive: boolean;
}

interface PendingFlyOut {
	storyId: string;
	agentName: string;
	label: string;
}

// ─── Provider ─────────────────────────────────────────────────────────────────

interface LozengeFlyProviderProps {
	children: React.ReactNode;
	pipeline: PipelineState;
}

export function LozengeFlyProvider({
	children,
	pipeline,
}: LozengeFlyProviderProps) {
	const rosterElsRef = useRef<Map<string, HTMLElement>>(new Map());
	const savedSlotRectsRef = useRef<Map<string, DOMRect>>(new Map());
	const prevPipelineRef = useRef<PipelineState | null>(null);

	// Actions detected in useLayoutEffect, consumed in useEffect
	const pendingFlyInActionsRef = useRef<PendingFlyIn[]>([]);
	const pendingFlyOutActionsRef = useRef<PendingFlyOut[]>([]);

	const [pendingFlyIns, setPendingFlyIns] = useState<ReadonlySet<string>>(
		new Set(),
	);
	const [flyingLozenges, setFlyingLozenges] = useState<FlyingLozenge[]>([]);

	// Agents currently assigned to a work item (derived from pipeline state).
	const assignedAgentNames = useMemo(() => {
		const names = new Set<string>();
		for (const item of [
			...pipeline.upcoming,
			...pipeline.current,
			...pipeline.qa,
			...pipeline.merge,
		]) {
			if (item.agent) names.add(item.agent.agent_name);
		}
		return names;
	}, [pipeline]);

	// Agents whose fly-out (work item → roster) animation is still in flight.
	// Kept hidden until the clone lands so no duplicate badge flashes.
	const [flyingOutAgents, setFlyingOutAgents] = useState<ReadonlySet<string>>(
		new Set(),
	);

	// Union: hide badge whenever the agent is assigned OR still flying back.
	const hiddenRosterAgents = useMemo(() => {
		if (flyingOutAgents.size === 0) return assignedAgentNames;
		const combined = new Set(assignedAgentNames);
		for (const name of flyingOutAgents) combined.add(name);
		return combined;
	}, [assignedAgentNames, flyingOutAgents]);

	const registerRosterEl = useCallback(
		(agentName: string, el: HTMLElement | null) => {
			if (el) {
				rosterElsRef.current.set(agentName, el);
			} else {
				rosterElsRef.current.delete(agentName);
			}
		},
		[],
	);

	const saveSlotRect = useCallback((storyId: string, rect: DOMRect) => {
		savedSlotRectsRef.current.set(storyId, rect);
	}, []);

	// ── Detect pipeline changes (runs before paint) ───────────────────────────
	// Sets pendingFlyIns so slot lozenges hide before the browser paints,
	// preventing a one-frame "flash" of the visible lozenge before fly-in.
	useLayoutEffect(() => {
		if (prevPipelineRef.current === null) {
			prevPipelineRef.current = pipeline;
			return;
		}

		const prev = prevPipelineRef.current;
		const allPrev = [
			...prev.upcoming,
			...prev.current,
			...prev.qa,
			...prev.merge,
		];
		const allCurr = [
			...pipeline.upcoming,
			...pipeline.current,
			...pipeline.qa,
			...pipeline.merge,
		];

		const newFlyInStoryIds = new Set<string>();

		for (const curr of allCurr) {
			const prevItem = allPrev.find((p) => p.story_id === curr.story_id);
			const agentChanged =
				curr.agent &&
				(!prevItem?.agent ||
					prevItem.agent.agent_name !== curr.agent.agent_name);
			if (agentChanged && curr.agent) {
				const label = curr.agent.model
					? `${curr.agent.agent_name} ${curr.agent.model}`
					: curr.agent.agent_name;
				pendingFlyInActionsRef.current.push({
					storyId: curr.story_id,
					agentName: curr.agent.agent_name,
					label,
					isActive: curr.agent.status === "running",
				});
				newFlyInStoryIds.add(curr.story_id);
			}
		}

		for (const prevItem of allPrev) {
			if (!prevItem.agent) continue;
			const currItem = allCurr.find((c) => c.story_id === prevItem.story_id);
			const agentRemoved =
				!currItem?.agent ||
				currItem.agent.agent_name !== prevItem.agent.agent_name;
			if (agentRemoved) {
				const label = prevItem.agent.model
					? `${prevItem.agent.agent_name} ${prevItem.agent.model}`
					: prevItem.agent.agent_name;
				pendingFlyOutActionsRef.current.push({
					storyId: prevItem.story_id,
					agentName: prevItem.agent.agent_name,
					label,
				});
			}
		}

		prevPipelineRef.current = pipeline;

		// Only hide slots for stories that have a matching roster element
		if (newFlyInStoryIds.size > 0) {
			const hideable = new Set<string>();
			for (const storyId of newFlyInStoryIds) {
				const action = pendingFlyInActionsRef.current.find(
					(a) => a.storyId === storyId,
				);
				if (action && rosterElsRef.current.has(action.agentName)) {
					hideable.add(storyId);
				}
			}
			if (hideable.size > 0) {
				setPendingFlyIns((prev) => {
					const next = new Set(prev);
					for (const id of hideable) next.add(id);
					return next;
				});
			}
		}
	}, [pipeline]);

	// ── Execute animations (runs after paint, DOM positions are stable) ───────
	useEffect(() => {
		const flyIns = [...pendingFlyInActionsRef.current];
		pendingFlyInActionsRef.current = [];
		const flyOuts = [...pendingFlyOutActionsRef.current];
		pendingFlyOutActionsRef.current = [];

		for (const action of flyIns) {
			const rosterEl = rosterElsRef.current.get(action.agentName);
			const slotRect = savedSlotRectsRef.current.get(action.storyId);

			if (!rosterEl || !slotRect) {
				// No roster element: immediately reveal the slot lozenge
				setPendingFlyIns((prev) => {
					const next = new Set(prev);
					next.delete(action.storyId);
					return next;
				});
				continue;
			}

			const rosterRect = rosterEl.getBoundingClientRect();
			const id = `fly-in-${action.agentName}-${action.storyId}-${Date.now()}`;

			setFlyingLozenges((prev) => [
				...prev,
				{
					id,
					label: action.label,
					isActive: action.isActive,
					startX: rosterRect.left,
					startY: rosterRect.top,
					endX: slotRect.left,
					endY: slotRect.top,
					flying: false,
				},
			]);

			// FLIP "Play" step: after two frames the transition begins
			requestAnimationFrame(() => {
				requestAnimationFrame(() => {
					setFlyingLozenges((prev) =>
						prev.map((l) => (l.id === id ? { ...l, flying: true } : l)),
					);
				});
			});

			// After the transition completes, remove clone and reveal slot lozenge
			setTimeout(() => {
				setFlyingLozenges((prev) => prev.filter((l) => l.id !== id));
				setPendingFlyIns((prev) => {
					const next = new Set(prev);
					next.delete(action.storyId);
					return next;
				});
			}, 500);
		}

		for (const action of flyOuts) {
			const rosterEl = rosterElsRef.current.get(action.agentName);
			const slotRect = savedSlotRectsRef.current.get(action.storyId);
			if (!slotRect) continue;

			// Keep the roster badge hidden while the clone is flying back.
			setFlyingOutAgents((prev) => {
				const next = new Set(prev);
				next.add(action.agentName);
				return next;
			});

			const rosterRect = rosterEl?.getBoundingClientRect();
			const id = `fly-out-${action.agentName}-${action.storyId}-${Date.now()}`;

			setFlyingLozenges((prev) => [
				...prev,
				{
					id,
					label: action.label,
					isActive: false,
					startX: slotRect.left,
					startY: slotRect.top,
					endX: rosterRect?.left ?? slotRect.left,
					endY: rosterRect?.top ?? Math.max(0, slotRect.top - 80),
					flying: false,
				},
			]);

			requestAnimationFrame(() => {
				requestAnimationFrame(() => {
					setFlyingLozenges((prev) =>
						prev.map((l) => (l.id === id ? { ...l, flying: true } : l)),
					);
				});
			});

			setTimeout(() => {
				setFlyingLozenges((prev) => prev.filter((l) => l.id !== id));
				// Reveal the roster badge now that the clone has landed.
				setFlyingOutAgents((prev) => {
					const next = new Set(prev);
					next.delete(action.agentName);
					return next;
				});
			}, 500);
		}
	}, [pipeline]);

	const contextValue = useMemo(
		() => ({
			registerRosterEl,
			saveSlotRect,
			pendingFlyIns,
			hiddenRosterAgents,
		}),
		[registerRosterEl, saveSlotRect, pendingFlyIns, hiddenRosterAgents],
	);

	return (
		<LozengeFlyContext.Provider value={contextValue}>
			{children}
			{ReactDOM.createPortal(
				<FloatingLozengeSurface lozenges={flyingLozenges} />,
				document.body,
			)}
		</LozengeFlyContext.Provider>
	);
}

// ─── Portal surface ───────────────────────────────────────────────────────────

function FloatingLozengeSurface({ lozenges }: { lozenges: FlyingLozenge[] }) {
	return (
		<>
			{lozenges.map((l) => (
				<FlyingLozengeClone key={l.id} lozenge={l} />
			))}
		</>
	);
}

function FlyingLozengeClone({ lozenge }: { lozenge: FlyingLozenge }) {
	const color = lozenge.isActive ? "#3fb950" : "#e3b341";
	const x = lozenge.flying ? lozenge.endX : lozenge.startX;
	const y = lozenge.flying ? lozenge.endY : lozenge.startY;

	return (
		<div
			data-testid={`flying-lozenge-${lozenge.id}`}
			style={{
				position: "fixed",
				left: `${x}px`,
				top: `${y}px`,
				zIndex: 9999,
				pointerEvents: "none",
				transition: lozenge.flying
					? "left 0.4s cubic-bezier(0.4, 0, 0.2, 1), top 0.4s cubic-bezier(0.4, 0, 0.2, 1)"
					: "none",
				display: "inline-flex",
				alignItems: "center",
				gap: "5px",
				padding: "2px 8px",
				borderRadius: "999px",
				fontSize: "0.72em",
				fontWeight: 600,
				background: `${color}18`,
				color,
				border: `1px solid ${color}44`,
				whiteSpace: "nowrap",
			}}
		>
			{lozenge.isActive && (
				<span
					style={{
						width: "5px",
						height: "5px",
						borderRadius: "50%",
						background: color,
						animation: "pulse 1.5s infinite",
						flexShrink: 0,
					}}
				/>
			)}
			{lozenge.label}
		</div>
	);
}

// ─── Hook ─────────────────────────────────────────────────────────────────────

export function useLozengeFly(): LozengeFlyContextValue {
	return useContext(LozengeFlyContext);
}