pi-gui/ui/src/workspace/useWorkspace.ts

import { invoke } from "@tauri-apps/api/core";
import { listen } from "@tauri-apps/api/event";
import { getCurrentWindow } from "@tauri-apps/api/window";
import { useCallback, useEffect, useReducer, useRef } from "react";
import { initialWorkspaceState, workspaceReducer } from "./reducer";
import {
	legacyPreferences,
	loadPresentation,
	presentationState,
	removeLegacyPreferences,
	savePresentation,
} from "./persistence";
import type {
	HostEvent,
	RuntimeSnapshot,
	UiState,
	Workspace,
	WorkspaceSummary,
} from "./types";

export function useWorkspace() {
	const [state, dispatch] = useReducer(workspaceReducer, initialWorkspaceState);
	const stateRef = useRef(state);
	stateRef.current = state;
	const loadGeneration = useRef<Record<string, number>>({});
	const timers = useRef<Record<string, number>>({});
	const initialized = useRef(false);
	const legacyMigrationPending = useRef(false);
	const disposed = useRef(true);
	const refreshGeneration = useRef(0);

	const loadSnapshot = useCallback(async (runtimeId: string) => {
		if (disposed.current) return;
		const generation = (loadGeneration.current[runtimeId] ?? 0) + 1;
		loadGeneration.current[runtimeId] = generation;
		dispatch({ type: "snapshotRequested", runtimeId, generation });
		try {
			const snapshot = await invoke<RuntimeSnapshot>(
				"get_session_runtime_snapshot",
				{ runtimeId },
			);
			if (!disposed.current)
				dispatch({ type: "snapshotLoaded", runtimeId, generation, snapshot });
		} catch (error) {
			if (!disposed.current)
				dispatch({
					type: "snapshotFailed",
					runtimeId,
					generation,
					error: String(error),
				});
		}
	}, []);

	const scheduleSnapshot = useCallback(
		(runtimeId: string) => {
			if (timers.current[runtimeId]) return;
			timers.current[runtimeId] = window.setTimeout(() => {
				delete timers.current[runtimeId];
				void loadSnapshot(runtimeId);
			}, 120);
		},
		[loadSnapshot],
	);

	const refresh = useCallback(
		async (initialUiState?: UiState, reloadAllSnapshots = false) => {
			const generation = ++refreshGeneration.current;
			const cancelled = () =>
				disposed.current || generation !== refreshGeneration.current;
			const workspace = await invoke<Workspace>("get_workspace");
			if (cancelled()) return;
			let nextPersisted = initialUiState;
			if (nextPersisted) {
				const legacy = legacyPreferences(localStorage);
				const candidate = nextPersisted.selectedRuntimeId
					? undefined
					: workspace.directories.find(
							(directory) => directory.worktreePath === legacy.path,
						)?.runtimes[0];
				nextPersisted = {
					...nextPersisted,
					...(candidate ? { selectedRuntimeId: candidate.runtimeId } : {}),
					...(nextPersisted.interfaceScale === undefined &&
					legacy.scale !== undefined
						? { interfaceScale: legacy.scale }
						: {}),
				};
				legacyMigrationPending.current = Boolean(
					(!initialUiState?.selectedRuntimeId && candidate) ||
						(initialUiState?.interfaceScale === undefined &&
							legacy.scale !== undefined),
				);
			}
			dispatch({
				type: "workspaceLoaded",
				workspace,
				persisted: nextPersisted,
			});
			const summary = await invoke<WorkspaceSummary>("get_workspace_summary");
			if (cancelled()) return;
			dispatch({ type: "setResourceWarning", value: summary.resourceWarning });
			const selected =
				nextPersisted?.selectedRuntimeId &&
				workspace.directories.some((directory) =>
					directory.runtimes.some(
						(runtime) => runtime.runtimeId === nextPersisted?.selectedRuntimeId,
					),
				)
					? nextPersisted.selectedRuntimeId
					: workspace.directories[0]?.runtimes[0]?.runtimeId;
			const runtimeIds = workspace.directories.flatMap((directory) =>
				directory.runtimes.map((runtime) => runtime.runtimeId),
			);
			for (const runtimeId of reloadAllSnapshots
				? runtimeIds
				: selected
					? [selected]
					: []) {
				if (cancelled()) return;
				void loadSnapshot(runtimeId);
			}
			if (cancelled()) return;
			await invoke("subscribe_workspace", {
				cursor: {
					bridgeInstanceId: workspace.bridgeInstanceId,
					seq: workspace.latestSeq,
				},
			});
			if (cancelled()) await invoke("unsubscribe_workspace").catch(() => {});
		},
		[loadSnapshot],
	);

	useEffect(() => {
		disposed.current = false;
		let stop: (() => void) | undefined;
		let stopNewRuntime: (() => void) | undefined;
		void listen<{
			phase: "starting" | "ready" | "error";
			runtimeId?: string;
			detail: string;
		}>("workspace-new-session", (message) => {
			const launch = message.payload;
			if (launch.phase === "starting") {
				dispatch({
					type: "newSessionLaunch",
					launch: { phase: "starting", detail: launch.detail },
				});
			} else if (launch.phase === "ready" && launch.runtimeId) {
				dispatch({ type: "newSessionLaunch", launch: undefined });
				void refresh().then(() => {
					if (!disposed.current)
						dispatch({ type: "selectRuntime", runtimeId: launch.runtimeId! });
				});
			} else if (launch.phase === "error") {
				dispatch({
					type: "newSessionLaunch",
					launch: { phase: "error", detail: launch.detail },
				});
			}
		}).then((unlisten) => {
			if (disposed.current) unlisten();
			else stopNewRuntime = unlisten;
		});
		void listen<HostEvent>("workspace-bridge", (message) => {
			if (disposed.current) return;
			const host = message.payload;
			dispatch({ type: "hostEvent", host });
			if (host.kind === "resetRequired") void refresh(undefined, true);
			else if (host.kind === "event") {
				const runtimeId = host.event.data?.runtimeId;
				const eventType = host.event.type;
				const piType = String(host.event.data?.eventData?.event?.type ?? "");
				if (
					[
						"runtime_opened",
						"runtime_closed",
						"runtime_recovered",
						"runtime_restore_failed",
					].includes(eventType ?? "")
				) {
					void refresh();
				} else if (
					runtimeId &&
					(host.event.data?.eventType === "transcript" ||
						piType === "agent_settled" ||
						piType === "compaction_end")
				) {
					scheduleSnapshot(runtimeId);
				}
			}
		}).then((unlisten) => {
			if (disposed.current) unlisten();
			else stop = unlisten;
		});
		void (async () => {
			try {
				const loaded = await loadPresentation();
				if (disposed.current) return;
				await refresh(loaded);
				if (disposed.current) return;
				initialized.current = true;
			} catch (error) {
				if (!disposed.current)
					dispatch({
						type: "hostEvent",
						host: {
							kind: "disconnected",
							message: String(error),
							retryInMs: 0,
						},
					});
			}
		})();
		const appWindow = getCurrentWindow();
		const visibility = () =>
			dispatch({
				type: "setWindowActive",
				active: document.visibilityState === "visible" && document.hasFocus(),
			});
		window.addEventListener("focus", visibility);
		window.addEventListener("blur", visibility);
		document.addEventListener("visibilitychange", visibility);
		return () => {
			disposed.current = true;
			refreshGeneration.current += 1;
			stop?.();
			stopNewRuntime?.();
			void invoke("unsubscribe_workspace").catch(() => {});
			window.removeEventListener("focus", visibility);
			window.removeEventListener("blur", visibility);
			document.removeEventListener("visibilitychange", visibility);
			Object.values(timers.current).forEach(clearTimeout);
			void appWindow;
		};
	}, [refresh, scheduleSnapshot]);

	useEffect(() => {
		if (!initialized.current) return;
		const timer = window.setTimeout(() => {
			void savePresentation(presentationState(state))
				.then(() => {
					if (legacyMigrationPending.current) {
						removeLegacyPreferences(localStorage);
						legacyMigrationPending.current = false;
					}
				})
				.catch(() => {});
		}, 300);
		return () => clearTimeout(timer);
	}, [
		state.directoryOrder,
		state.runtimeOrderByDirectory,
		state.selectedRuntimeId,
		state.defaultNewSessionWorktree,
		state.interfaceScale,
		state.bridge.latestSeq,
		state.runtimesById,
	]);

	useEffect(() => {
		const id = state.selectedRuntimeId;
		if (!id) return;
		if (
			!state.runtimesById[id]?.loading &&
			!state.runtimesById[id]?.messages.length
		)
			void loadSnapshot(id);
	}, [state.selectedRuntimeId, loadSnapshot]);

	return { state, dispatch, refresh, loadSnapshot, scheduleSnapshot };
}