import type {
	Extension,
	HostEvent,
	RuntimeSnapshot,
	RuntimeSummary,
	RuntimeView,
	UiState,
	Workspace,
	WorkspaceEvent,
	WorkspaceState,
} from "./types";

const idleProgress = {
	phase: "idle" as const,
	detail: "Ready for your next prompt",
	toolCount: 0,
};

function snapshotProgress(runtime: RuntimeSummary) {
	if (runtime.state === "streaming")
		return {
			phase: "working" as const,
			detail: "Thinking and preparing a response",
			toolCount: 0,
		};
	if (runtime.state === "recovering")
		return {
			phase: "recovering" as const,
			detail: "Restoring this Pi session",
			toolCount: 0,
		};
	if (runtime.state === "failed")
		return {
			phase: "error" as const,
			detail: runtime.error?.message ?? "This Pi session failed",
			toolCount: 0,
		};
	return idleProgress;
}

function messageText(message: RuntimeView["messages"][number]) {
	if (typeof message.content === "string") return message.content;
	return (message.content ?? [])
		.filter((part) => part.type === "text")
		.map((part) => part.text ?? "")
		.join("");
}

function reconcilePendingSubmissions(
	pendingSubmissions: RuntimeView["pendingSubmissions"],
	messages: RuntimeView["messages"],
) {
	const userMessages = messages
		.filter((message) => message.role === "user")
		.map(messageText);
	const consumed = new Set<number>();
	return pendingSubmissions.filter((pending) => {
		for (
			let index = pending.userMessageCount;
			index < userMessages.length;
			index++
		) {
			if (consumed.has(index) || userMessages[index] !== pending.text) continue;
			consumed.add(index);
			return false;
		}
		return true;
	});
}

export function runtimeView(
	summary: RuntimeSummary,
	previous?: RuntimeView,
): RuntimeView {
	return {
		summary,
		state: previous?.state ?? {},
		stats: previous?.stats ?? {},
		messages: previous?.messages ?? [],
		commands: previous?.commands ?? [],
		models: previous?.models ?? [],
		modelError: previous?.modelError,
		extensions: previous?.extensions ?? [],
		extensionStatuses: previous?.extensionStatuses ?? {},
		extensionWidgets: previous?.extensionWidgets ?? {},
		notifications: previous?.notifications ?? [],
		progress: previous?.progress ?? idleProgress,
		pendingSubmissions: previous?.pendingSubmissions ?? [],
		draft: previous?.draft ?? "",
		scroll: previous?.scroll ?? { atBottom: true },
		unread: previous?.unread ?? false,
		loading: previous?.loading ?? false,
		snapshotRetry: previous?.snapshotRetry ?? false,
		error: previous?.error,
		requestGeneration: previous?.requestGeneration ?? 0,
		loadedThroughSeq: previous?.loadedThroughSeq ?? 0,
	};
}

function resetRuntimeView(
	summary: RuntimeSummary,
	previous: RuntimeView,
): RuntimeView {
	return {
		...runtimeView(summary, previous),
		state: {},
		stats: {},
		messages: [],
		commands: [],
		models: [],
		modelError: undefined,
		extensions: [],
		extensionStatuses: {},
		extensionWidgets: {},
		notifications: [],
		pendingSubmissions: [],
		unread: false,
		progress: idleProgress,
		loading: false,
		snapshotRetry: false,
		error: undefined,
		requestGeneration: 0,
		loadedThroughSeq: 0,
	};
}

export const initialWorkspaceState: WorkspaceState = {
	bridge: { latestSeq: 0, status: "connecting" },
	directoriesByPath: {},
	directoryOrder: [],
	runtimesById: {},
	runtimeOrderByDirectory: {},
	windowActive: true,
	resourceWarning: false,
	defaultNewSessionWorktree: undefined,
	dismissedCollisionWarning: false,
	interfaceScale: 1,
};

function ordered(current: string[] | undefined, authoritative: string[]) {
	const valid = new Set(authoritative);
	return [
		...(current ?? []).filter((id) => valid.has(id)),
		...authoritative.filter((id) => !(current ?? []).includes(id)),
	];
}

export function reconcileWorkspace(
	state: WorkspaceState,
	workspace: Workspace,
	persisted?: UiState,
): WorkspaceState {
	const epochChanged = Boolean(
		state.bridge.instanceId &&
			state.bridge.instanceId !== workspace.bridgeInstanceId,
	);
	const directoriesByPath: WorkspaceState["directoriesByPath"] = {};
	const runtimesById: WorkspaceState["runtimesById"] = {};
	const runtimeOrderByDirectory: WorkspaceState["runtimeOrderByDirectory"] = {};
	for (const directory of workspace.directories) {
		directoriesByPath[directory.worktreePath] = directory;
		for (const runtime of directory.runtimes) {
			const previous = state.runtimesById[runtime.runtimeId];
			runtimesById[runtime.runtimeId] =
				epochChanged && previous
					? resetRuntimeView(runtime, previous)
					: runtimeView(runtime, previous);
		}
		const ids = directory.runtimes.map((runtime) => runtime.runtimeId);
		runtimeOrderByDirectory[directory.worktreePath] = ordered(
			persisted?.runtimeOrderByDirectory[directory.worktreePath] ??
				state.runtimeOrderByDirectory[directory.worktreePath],
			ids,
		);
	}
	const paths = workspace.directories.map((directory) => directory.worktreePath);
	const directoryOrder = ordered(
		persisted?.directoryOrder ?? state.directoryOrder,
		paths,
	);
	const preferred = persisted?.selectedRuntimeId ?? state.selectedRuntimeId;
	const selectedRuntimeId =
		preferred && runtimesById[preferred]
			? preferred
			: directoryOrder.flatMap((path) => runtimeOrderByDirectory[path] ?? [])[0];
	const selectedDirectoryPath = selectedRuntimeId
		? runtimesById[selectedRuntimeId]?.summary.worktreePath
		: directoryOrder[0];
	if (persisted) {
		for (const [id, view] of Object.entries(runtimesById)) {
			view.draft = persisted.drafts[id] ?? view.draft;
			view.scroll = persisted.scroll[id] ?? view.scroll;
			const seen = persisted.lastSeen[id];
			if (
				seen?.bridgeInstanceId === workspace.bridgeInstanceId &&
				seen.seq < workspace.latestSeq &&
				seen.lastActivity !== view.summary.lastActivity
			)
				view.unread = true;
		}
	}
	return {
		...state,
		bridge: {
			instanceId: workspace.bridgeInstanceId,
			latestSeq: workspace.latestSeq,
			status: "ready",
		},
		directoriesByPath,
		directoryOrder,
		runtimesById,
		runtimeOrderByDirectory,
		selectedRuntimeId,
		selectedDirectoryPath,
		dismissedCollisionWarning:
			persisted?.dismissedCollisionWarning ?? state.dismissedCollisionWarning,
		defaultNewSessionWorktree:
			persisted?.defaultNewSessionWorktree ?? state.defaultNewSessionWorktree,
		interfaceScale: persisted?.interfaceScale ?? state.interfaceScale,
	};
}

function meaningful(eventType?: string, event?: Record<string, unknown>) {
	if (["extension_ui_request", "agent_state"].includes(eventType ?? ""))
		return true;
	if (eventType === "transcript")
		return ["message_end", "turn_end", "agent_end"].includes(
			String(event?.type ?? ""),
		);
	return false;
}

function extensionPresentation(view: RuntimeView, extension: Extension) {
	if (["select", "confirm", "input", "editor"].includes(extension.method ?? ""))
		view.extensions = [
			...view.extensions.filter((entry) => entry.id !== extension.id),
			extension,
		];
	else if (extension.method === "notify")
		view.notifications = [...view.notifications, extension].slice(-4);
	else if (extension.method === "setStatus" && extension.statusKey) {
		const next = { ...view.extensionStatuses };
		if (extension.statusText) next[extension.statusKey] = extension.statusText;
		else delete next[extension.statusKey];
		view.extensionStatuses = next;
	} else if (extension.method === "setWidget" && extension.widgetKey) {
		const next = { ...view.extensionWidgets };
		if (extension.widgetLines)
			next[extension.widgetKey] = {
				lines: extension.widgetLines,
				placement: extension.widgetPlacement ?? "aboveEditor",
			};
		else delete next[extension.widgetKey];
		view.extensionWidgets = next;
	} else if (extension.method === "set_editor_text")
		view.draft = extension.text ?? "";
}

export function applyWorkspaceEvent(
	state: WorkspaceState,
	event: WorkspaceEvent,
): WorkspaceState {
	if (
		state.bridge.instanceId &&
		event.bridgeInstanceId !== state.bridge.instanceId
	)
		return {
			...state,
			bridge: {
				...state.bridge,
				status: "resyncing",
				error: "Bridge restarted",
			},
		};
	if (event.seq <= state.bridge.latestSeq) return state;
	if (event.seq !== state.bridge.latestSeq + 1)
		return {
			...state,
			bridge: {
				...state.bridge,
				status: "resyncing",
				error: "Workspace event gap",
			},
		};
	const id = event.data?.runtimeId;
	if (
		event.type === "runtime_renewed" &&
		id &&
		state.runtimesById[id] &&
		event.data?.runtime
	) {
		const renewed = resetRuntimeView(event.data.runtime, state.runtimesById[id]);
		return {
			...state,
			bridge: { ...state.bridge, latestSeq: event.seq, status: "resyncing" },
			runtimesById: {
				...state.runtimesById,
				[id]: {
					...renewed,
					draft: "",
					scroll: { atBottom: true },
				},
			},
		};
	}
	if (
		["runtime_opened", "runtime_recovered", "runtime_restore_failed"].includes(
			event.type ?? "",
		)
	)
		return {
			...state,
			bridge: { ...state.bridge, latestSeq: event.seq, status: "resyncing" },
		};
	if (!id || (!state.runtimesById[id] && !event.data?.runtime))
		return {
			...state,
			bridge: {
				...state.bridge,
				latestSeq: event.seq,
				status: "resyncing",
				error: "Unknown runtime",
			},
		};
	const runtimesById = { ...state.runtimesById };
	let view = runtimeView(
		event.data?.runtime ?? runtimesById[id].summary,
		runtimesById[id],
	);
	view = { ...view, summary: event.data?.runtime ?? view.summary };
	if (event.type === "runtime_closed") {
		delete runtimesById[id];
		const path = view.summary.worktreePath;
		const order = (state.runtimeOrderByDirectory[path] ?? []).filter(
			(runtimeId) => runtimeId !== id,
		);
		const runtimeOrderByDirectory = {
			...state.runtimeOrderByDirectory,
			[path]: order,
		};
		const selectedRuntimeId =
			state.selectedRuntimeId === id
				? (order[0] ?? Object.keys(runtimesById)[0])
				: state.selectedRuntimeId;
		return {
			...state,
			bridge: { ...state.bridge, latestSeq: event.seq },
			runtimesById,
			runtimeOrderByDirectory,
			selectedRuntimeId,
			selectedDirectoryPath: selectedRuntimeId
				? runtimesById[selectedRuntimeId]?.summary.worktreePath
				: path,
		};
	}
	const type = event.data?.eventType;
	const data = event.data?.eventData;
	const payload = data?.event ?? {};
	if (type === "agent_state" && data?.state) {
		view.summary = { ...view.summary, state: data.state };
		view.progress =
			data.state === "streaming"
				? {
						...view.progress,
						phase: "working",
						detail: "Thinking and preparing a response",
					}
				: data.state === "recovering"
					? {
							...view.progress,
							phase: "recovering",
							detail: "Recovering the Pi session",
						}
					: data.state === "error"
						? { ...view.progress, phase: "error", detail: "Pi needs attention" }
						: idleProgress;
	} else if (type === "tool") {
		const ending = String(payload.type ?? "").endsWith("_end");
		view.progress = {
			...view.progress,
			phase: "working",
			detail: ending ? "Continuing after tool result" : "Running a tool",
			tool: ending ? undefined : String(payload.toolName ?? "tool"),
			toolCount:
				view.progress.toolCount +
				(String(payload.type ?? "").endsWith("_start") ? 1 : 0),
		};
	} else if (type === "queue") {
		const queueCount =
			typeof payload.pendingMessageCount === "number"
				? payload.pendingMessageCount
				: (Array.isArray(payload.followUp) ? payload.followUp.length : 0) +
					(Array.isArray(payload.steering) ? payload.steering.length : 0);
		view.summary = { ...view.summary, queueCount };
		view.progress = {
			...view.progress,
			queueCount,
			detail: queueCount
				? `${queueCount} follow-up${queueCount === 1 ? "" : "s"} queued`
				: view.progress.detail,
		};
	} else if (type === "extension_ui_request")
		extensionPresentation(view, payload as Extension);
	if (id !== state.selectedRuntimeId && meaningful(type, payload))
		view.unread = true;
	view.loadedThroughSeq = Math.max(view.loadedThroughSeq, event.seq);
	runtimesById[id] = view;
	return {
		...state,
		bridge: { ...state.bridge, latestSeq: event.seq, status: "ready" },
		runtimesById,
	};
}

export type WorkspaceAction =
	| { type: "workspaceLoaded"; workspace: Workspace; persisted?: UiState }
	| { type: "hostEvent"; host: HostEvent }
	| { type: "snapshotRequested"; runtimeId: string; generation: number }
	| {
			type: "snapshotLoaded";
			runtimeId: string;
			generation: number;
			snapshot: RuntimeSnapshot;
	  }
	| {
			type: "snapshotFailed";
			runtimeId: string;
			generation: number;
			error: string;
	  }
	| { type: "selectRuntime"; runtimeId: string }
	| { type: "selectDirectory"; worktreePath: string }
	| { type: "draftChanged"; runtimeId: string; draft: string }
	| {
			type: "scrollChanged";
			runtimeId: string;
			atBottom: boolean;
			scrollTop?: number;
	  }
	| { type: "markRead"; runtimeId: string }
	| { type: "setWindowActive"; active: boolean }
	| { type: "setScale"; scale: number }
	| { type: "setResourceWarning"; value: boolean }
	| { type: "setDefaultNewSessionWorktree"; worktreePath?: string }
	| { type: "newSessionLaunch"; launch?: WorkspaceState["newSessionLaunch"] }
	| { type: "dismissCollisionWarning" }
	| {
			type: "submissionAdded";
			runtimeId: string;
			submission: RuntimeView["pendingSubmissions"][number];
	  }
	| {
			type: "submissionSent";
			runtimeId: string;
			id: number;
			delivery: "prompt" | "follow_up";
	  }
	| { type: "submissionRemoved"; runtimeId: string; id: number };

export function workspaceReducer(
	state: WorkspaceState,
	action: WorkspaceAction,
): WorkspaceState {
	if (action.type === "workspaceLoaded")
		return reconcileWorkspace(state, action.workspace, action.persisted);
	if (action.type === "hostEvent") {
		if (action.host.kind === "event")
			return applyWorkspaceEvent(state, action.host.event);
		if (action.host.kind === "resetRequired")
			return {
				...state,
				bridge: {
					...state.bridge,
					latestSeq: 0,
					status: "resyncing",
					error: action.host.reason,
				},
			};
		if (action.host.kind === "disconnected")
			return {
				...state,
				bridge: {
					...state.bridge,
					status: "error",
					error: action.host.message,
				},
			};
		return {
			...state,
			bridge: {
				instanceId: action.host.bridgeInstanceId,
				latestSeq: Math.max(state.bridge.latestSeq, action.host.latestSeq),
				status: "ready",
			},
		};
	}
	if (action.type === "selectDirectory")
		return {
			...state,
			selectedDirectoryPath: action.worktreePath,
			selectedRuntimeId: state.runtimeOrderByDirectory[action.worktreePath]?.[0],
		};
	if (action.type === "selectRuntime")
		return {
			...state,
			selectedRuntimeId: action.runtimeId,
			selectedDirectoryPath:
				state.runtimesById[action.runtimeId]?.summary.worktreePath,
		};
	if (action.type === "setWindowActive")
		return { ...state, windowActive: action.active };
	if (action.type === "setScale")
		return { ...state, interfaceScale: action.scale };
	if (action.type === "setResourceWarning")
		return { ...state, resourceWarning: action.value };
	if (action.type === "setDefaultNewSessionWorktree")
		return { ...state, defaultNewSessionWorktree: action.worktreePath };
	if (action.type === "newSessionLaunch")
		return { ...state, newSessionLaunch: action.launch };
	if (action.type === "dismissCollisionWarning")
		return { ...state, dismissedCollisionWarning: true };
	const runtimeId = "runtimeId" in action ? action.runtimeId : undefined;
	if (!runtimeId || !state.runtimesById[runtimeId]) return state;
	const view = { ...state.runtimesById[runtimeId] };
	if (action.type === "snapshotRequested") {
		view.loading = true;
		view.snapshotRetry = false;
		view.error = undefined;
		view.requestGeneration = action.generation;
	} else if (action.type === "snapshotFailed") {
		if (view.requestGeneration !== action.generation) return state;
		view.loading = false;
		view.snapshotRetry = false;
		view.error = action.error;
	} else if (action.type === "snapshotLoaded") {
		if (view.requestGeneration !== action.generation) return state;
		if (
			action.snapshot.bridgeInstanceId !== state.bridge.instanceId ||
			action.snapshot.latestSeq < view.loadedThroughSeq
		) {
			view.loading = false;
			view.snapshotRetry = true;
			return {
				...state,
				runtimesById: { ...state.runtimesById, [runtimeId]: view },
			};
		}
		view.loading = false;
		view.snapshotRetry = false;
		view.summary = action.snapshot.runtime;
		view.state = action.snapshot.state?.data ?? {};
		view.stats = action.snapshot.stats?.data ?? {};
		view.messages = action.snapshot.transcript?.data?.messages ?? [];
		view.pendingSubmissions = reconcilePendingSubmissions(
			view.pendingSubmissions,
			view.messages,
		);
		view.commands = action.snapshot.commands?.data?.commands ?? [];
		if (action.snapshot.models?.error) {
			view.modelError = `Could not load available Pi models: ${action.snapshot.models.error}`;
		} else {
			view.models = action.snapshot.models?.data?.models ?? [];
			view.modelError = undefined;
		}
		view.extensions = action.snapshot.extensions ?? [];
		view.loadedThroughSeq = action.snapshot.latestSeq;
		view.progress = snapshotProgress(action.snapshot.runtime);
	} else if (action.type === "draftChanged") view.draft = action.draft;
	else if (action.type === "scrollChanged")
		view.scroll = { atBottom: action.atBottom, scrollTop: action.scrollTop };
	else if (action.type === "markRead") {
		if (
			state.selectedRuntimeId === runtimeId &&
			state.windowActive &&
			view.scroll.atBottom
		)
			view.unread = false;
	} else if (action.type === "submissionAdded")
		view.pendingSubmissions = [...view.pendingSubmissions, action.submission];
	else if (action.type === "submissionSent")
		view.pendingSubmissions = view.pendingSubmissions.map((item) =>
			item.id === action.id
				? { ...item, phase: "sent", delivery: action.delivery }
				: item,
		);
	else if (action.type === "submissionRemoved")
		view.pendingSubmissions = view.pendingSubmissions.filter(
			(item) => item.id !== action.id,
		);
	return {
		...state,
		runtimesById: { ...state.runtimesById, [runtimeId]: view },
	};
}