apps/canvas/front/src/lib/state.ts

import { Category, defaultCategories } from "./categories";
import { CreateValidators, Validator } from "./config";
import { GitHubService, GitHubServiceImpl, GitHubRepository } from "./github";
import type { Edge, OnConnect, OnEdgesChange, OnNodesChange, Viewport as ReactFlowViewport } from "@xyflow/react";
import {
	addEdge,
	applyEdgeChanges,
	applyNodeChanges,
	Connection,
	EdgeChange,
	useNodes,
	XYPosition,
} from "@xyflow/react";
import type { DeepPartial } from "react-hook-form";
import { v4 as uuidv4 } from "uuid";
import { create } from "zustand";
import {
	AppNode,
	Env,
	NodeType,
	VolumeNode,
	GatewayTCPData,
	envSchema,
	AgentAccess,
	ViewportTransform,
	GraphSchema,
} from "config";

export function nodeLabel(n: AppNode): string {
	try {
		switch (n.type) {
			case "network":
				return n.data.domain;
			case "app":
				return n.data.label || "Service";
			case "github":
				return n.data.repository?.fullName || "Github";
			case "gateway-https": {
				if (n.data && n.data.subdomain) {
					return `${n.data.subdomain}`;
				} else {
					return "HTTPS Gateway";
				}
			}
			case "gateway-tcp": {
				if (n.data && n.data.subdomain) {
					return `${n.data.subdomain}`;
				} else {
					return "TCP Gateway";
				}
			}
			case "mongodb":
				return n.data.label || "MongoDB";
			case "postgresql":
				return n.data.label || "PostgreSQL";
			case "volume":
				return n.data.label || "Volume";
			case undefined:
				throw new Error(`nodeLabel: Node type is undefined. Node ID: ${n.id}, Data: ${JSON.stringify(n.data)}`);
		}
	} catch (e) {
		console.error("opaa", e);
	} finally {
		console.log("done");
	}
	return "Unknown Node";
}

export function nodeLabelFull(n: AppNode): string {
	if (n.type === "gateway-https") {
		return `https://${n.data.subdomain}.${n.data.network}`;
	} else {
		return nodeLabel(n);
	}
}

export function nodeIsConnectable(n: AppNode, handle: string): boolean {
	switch (n.type) {
		case "network":
			return true;
		case "app":
			if (handle === "ports") {
				return n.data !== undefined && n.data.ports !== undefined && n.data.ports.length > 0;
			} else if (handle === "repository") {
				if (!n.data || !n.data.repository || !n.data.repository.id) {
					return true;
				}
				return false;
			}
			return false;
		case "github":
			if (n.data.repository?.id !== undefined) {
				return true;
			}
			return false;
		case "gateway-https":
			if (handle === "subdomain") {
				return n.data.network === undefined;
			}
			return n.data === undefined || n.data.https === undefined;
		case "gateway-tcp":
			if (handle === "subdomain") {
				return n.data.network === undefined;
			}
			return true;
		case "mongodb":
			return true;
		case "postgresql":
			return true;
		case "volume":
			if (n.data === undefined || n.data.type === undefined) {
				return false;
			}
			if (n.data.type === "ReadWriteOnce" || n.data.type === "ReadWriteOncePod") {
				return n.data.attachedTo === undefined || n.data.attachedTo.length === 0;
			}
			return true;
		case undefined:
			throw new Error(
				`nodeIsConnectable: Node type is undefined. Node ID: ${n.id}, Handle: ${handle}, Data: ${JSON.stringify(n.data)}`,
			);
	}
}

export function nodeEnvVarNamePort(n: AppNode, portName: string): string {
	return `DODO_SERVICE_${n.data.label.toUpperCase()}_ADDRESS_${portName.toUpperCase()}`;
}

export function nodeEnvVarNames(n: AppNode): string[] {
	switch (n.type) {
		case "app":
			return [
				`DODO_SERVICE_${n.data.label.toUpperCase()}_ADDRESS`,
				...(n.data.ports || []).map((p) => nodeEnvVarNamePort(n, p.name)),
			];
		case "github":
			return [];
		case "gateway-https":
			return [];
		case "gateway-tcp":
			return [];
		case "mongodb":
			return [`DODO_MONGODB_${n.data.label.toUpperCase()}_URL`];
		case "postgresql":
			return [`DODO_POSTGRESQL_${n.data.label.toUpperCase()}_URL`];
		case "volume":
			return [`DODO_VOLUME_${n.data.label.toUpperCase()}_PATH`];
		case undefined:
			throw new Error("MUST NOT REACH");
		default:
			throw new Error("MUST NOT REACH");
	}
}

export type MessageType = "INFO" | "WARNING" | "FATAL";

export type Message = {
	id: string;
	type: MessageType;
	nodeId?: string;
	message: string;
	onHighlight?: (state: AppState) => void;
	onLooseHighlight?: (state: AppState) => void;
	onClick?: (state: AppState) => void;
};

const defaultEnv: Env = {
	deployKeyPublic: undefined,
	instanceId: undefined,
	networks: [],
	integrations: {
		github: false,
		gemini: false,
		anthropic: false,
	},
	services: [],
	user: {
		id: "",
		username: "",
	},
	access: [],
};

export type Project = {
	id: string;
	name: string;
};

export type IntegrationsConfig = {
	github: boolean;
};

type NodeUpdate<T extends NodeType> = DeepPartial<Extract<AppNode, { type: T }>>;
type NodeDataUpdate<T extends NodeType> = DeepPartial<Extract<AppNode, { type: T }>["data"]>;

let refreshEnvIntervalId: number | null = null;

export type AppState = {
	projectId: string | undefined;
	mode: "edit" | "deploy";
	buildMode: "overview" | "canvas";
	projects: Project[];
	nodes: AppNode[];
	edges: Edge[];
	zoom: ReactFlowViewport;
	categories: Category[];
	messages: Message[];
	env: Env;
	viewport: ViewportTransform;
	setViewport: (viewport: ViewportTransform) => void;
	githubService: GitHubService | null;
	githubRepositories: GitHubRepository[];
	githubRepositoriesLoading: boolean;
	githubRepositoriesError: string | null;
	stateEventSource: EventSource | null;
	setHighlightCategory: (name: string, active: boolean) => void;
	onNodesChange: OnNodesChange<AppNode>;
	onEdgesChange: OnEdgesChange;
	onConnect: OnConnect;
	addNode: (node: Omit<AppNode, "position">) => void;
	setNodes: (nodes: AppNode[]) => void;
	setEdges: (edges: Edge[]) => void;
	setProject: (projectId: string | undefined) => Promise<void>;
	setMode: (mode: "edit" | "deploy") => void;
	setBuildMode: (buildMode: "overview" | "canvas") => void;
	updateNode: <T extends NodeType>(id: string, node: NodeUpdate<T>) => void;
	updateNodeData: <T extends NodeType>(id: string, data: NodeDataUpdate<T>) => void;
	replaceEdge: (c: Connection, id?: string) => void;
	refreshEnv: () => Promise<void>;
	fetchGithubRepositories: () => Promise<void>;
};

const projectIdSelector = (state: AppState) => state.projectId;
const categoriesSelector = (state: AppState) => state.categories;
const messagesSelector = (state: AppState) => state.messages;
const envSelector = (state: AppState) => state.env;
const zoomSelector = (state: AppState) => state.zoom;
const githubRepositoriesSelector = (state: AppState) => state.githubRepositories;
const githubRepositoriesLoadingSelector = (state: AppState) => state.githubRepositoriesLoading;
const githubRepositoriesErrorSelector = (state: AppState) => state.githubRepositoriesError;
const buildModeSelector = (state: AppState) => state.buildMode;

export function useZoom(): ReactFlowViewport {
	return useStateStore(zoomSelector);
}

export function useProjectId(): string | undefined {
	return useStateStore(projectIdSelector);
}

export function useSetProject(): (projectId: string | undefined) => void {
	return useStateStore((state) => state.setProject);
}

export function useCategories(): Category[] {
	return useStateStore(categoriesSelector);
}

export function useMessages(): Message[] {
	return useStateStore(messagesSelector);
}

export function useNodeMessages(id: string): Message[] {
	return useMessages().filter((m) => m.nodeId === id);
}

export function useNodeLabel(id: string): string {
	return nodeLabel(useNodes<AppNode>().find((n) => n.id === id)!);
}

export function useNodePortName(id: string, portId: string): string {
	return (useNodes<AppNode>().find((n) => n.id === id)!.data.ports || []).find((p) => p.id === portId)!.name;
}

export function useEnv(): Env {
	return useStateStore(envSelector);
}

export function useAgents(): AgentAccess[] {
	return useStateStore(envSelector).access.filter(
		(acc): acc is AgentAccess => acc.type === "https" && acc.agentName != null,
	);
}

export function useLeadAgent(): AgentAccess | undefined {
	const agents = useAgents();
	return agents.find((a) => a.agentName === "dodo") || agents[0];
}

export function useGithubService(): boolean {
	return useStateStore(envSelector).integrations.github;
}

export function useGeminiService(): boolean {
	return useStateStore(envSelector).integrations.gemini;
}

export function useAnthropicService(): boolean {
	return useStateStore(envSelector).integrations.anthropic;
}

export function useGithubRepositories(): GitHubRepository[] {
	return useStateStore(githubRepositoriesSelector);
}

export function useGithubRepositoriesLoading(): boolean {
	return useStateStore(githubRepositoriesLoadingSelector);
}

export function useGithubRepositoriesError(): string | null {
	return useStateStore(githubRepositoriesErrorSelector);
}

export function useFetchGithubRepositories(): () => Promise<void> {
	return useStateStore((state) => state.fetchGithubRepositories);
}

export function useMode(): "edit" | "deploy" {
	return useStateStore((state) => state.mode);
}

export function useBuildMode(): "overview" | "canvas" {
	return useStateStore(buildModeSelector);
}

const v: Validator = CreateValidators();

function getRandomPosition({ width, height, transformX, transformY, transformZoom }: ViewportTransform): XYPosition {
	const zoomMultiplier = 1 / transformZoom;
	const realWidth = width * zoomMultiplier;
	const realHeight = height * zoomMultiplier;
	const paddingMultiplier = 0.8;
	const ret = {
		x: -transformX * zoomMultiplier + Math.random() * realWidth * paddingMultiplier,
		y: -transformY * zoomMultiplier + Math.random() * realHeight * paddingMultiplier,
	};
	return ret;
}

export const useStateStore = create<AppState>((setOg, get): AppState => {
	const set = (state: Partial<AppState>) => {
		setOg(state);
	};
	const setN = (nodes: AppNode[]) => {
		const env = get().env;
		console.log("---env", env);
		set({
			nodes,
			messages: v(nodes, env),
		});
	};

	const startRefreshEnvInterval = () => {
		if (refreshEnvIntervalId) {
			clearInterval(refreshEnvIntervalId);
		}
		if (get().projectId && typeof document !== "undefined" && document.visibilityState === "visible") {
			console.log("Starting refreshEnv interval for project:", get().projectId);
			refreshEnvIntervalId = setInterval(async () => {
				if (get().projectId && typeof document !== "undefined" && document.visibilityState === "visible") {
					console.log("Interval: Calling refreshEnv for project:", get().projectId);
					await get().refreshEnv();
				} else if (refreshEnvIntervalId) {
					console.log(
						"Interval: Conditions not met (project removed or tab hidden), stopping interval from inside.",
					);
					clearInterval(refreshEnvIntervalId);
					refreshEnvIntervalId = null;
				}
			}, 5000) as unknown as number;
		} else {
			console.log(
				"Not starting refreshEnv interval. Project ID:",
				get().projectId,
				"Visibility:",
				typeof document !== "undefined" ? document.visibilityState : "SSR",
			);
		}
	};

	const stopRefreshEnvInterval = () => {
		if (refreshEnvIntervalId) {
			console.log("Stopping refreshEnv interval for project:", get().projectId);
			clearInterval(refreshEnvIntervalId);
			refreshEnvIntervalId = null;
		}
	};

	if (typeof document !== "undefined") {
		document.addEventListener("visibilitychange", () => {
			if (document.visibilityState === "visible") {
				console.log("Tab became visible, attempting to start refreshEnv interval.");
				startRefreshEnvInterval();
			} else {
				console.log("Tab became hidden, stopping refreshEnv interval.");
				stopRefreshEnvInterval();
			}
		});
	}

	const injectNetworkNodes = () => {
		const newNetworks = get().env.networks.filter(
			(x) => !get().nodes.some((n) => n.type === "network" && n.data.domain === x.domain),
		);
		newNetworks.forEach((n) => {
			get().addNode({
				id: n.domain,
				type: "network",
				connectable: true,
				data: {
					domain: n.domain,
					label: n.domain,
					envVars: [],
					ports: [],
					state: "success", // TODO(gio): monitor network health
				},
			});
			console.log("added network", n.domain);
		});
	};

	function updateNodeData<T extends NodeType>(id: string, data: NodeDataUpdate<T>): void {
		setN(
			get().nodes.map((n) => {
				if (n.id === id) {
					return {
						...n,
						data: {
							...n.data,
							...data,
						},
					} as Extract<AppNode, { type: T }>;
				}
				return n;
			}),
		);
	}

	function updateNode<T extends NodeType>(id: string, node: NodeUpdate<T>): void {
		setN(
			get().nodes.map((n) => {
				if (n.id === id) {
					return {
						...n,
						...node,
					} as Extract<AppNode, { type: T }>;
				}
				return n;
			}),
		);
	}

	function onConnect(c: Connection) {
		const { nodes, edges } = get();
		set({
			edges: addEdge(c, edges),
		});
		const sn = nodes.filter((n) => n.id === c.source)[0]!;
		const tn = nodes.filter((n) => n.id === c.target)[0]!;
		if (tn.type === "network") {
			if (sn.type === "gateway-https") {
				updateNodeData<"gateway-https">(sn.id, {
					network: tn.data.domain,
				});
			} else if (sn.type === "gateway-tcp") {
				updateNodeData<"gateway-tcp">(sn.id, {
					network: tn.data.domain,
				});
			}
		}
		if (tn.type === "app") {
			if (c.sourceHandle === "env_var" && c.targetHandle === "env_var") {
				const sourceEnvVars = nodeEnvVarNames(sn);
				if (sourceEnvVars.length === 0) {
					throw new Error(
						`onConnect (env_var): Source node ${sn.id} (type: ${sn.type}) has no env vars to connect from.`,
					);
				}
				const id = uuidv4();
				if (sourceEnvVars.length === 1) {
					updateNode<"app">(c.target, {
						...tn,
						data: {
							...tn.data,
							envVars: [
								...(tn.data.envVars || []),
								{
									id: id,
									source: c.source,
									name: sourceEnvVars[0],
									isEditting: false,
								},
							],
						},
					});
				} else {
					updateNode<"app">(c.target, {
						...tn,
						data: {
							...tn.data,
							envVars: [
								...(tn.data.envVars || []),
								{
									id: id,
									source: c.source,
								},
							],
						},
					});
				}
			}
			if (c.sourceHandle === "ports" && c.targetHandle === "env_var") {
				const sourcePorts = sn.data.ports || [];
				const id = uuidv4();
				if (sourcePorts.length === 1) {
					updateNode<"app">(c.target, {
						...tn,
						data: {
							...tn.data,
							envVars: [
								...(tn.data.envVars || []),
								{
									id: id,
									source: c.source,
									name: nodeEnvVarNamePort(sn, sourcePorts[0].name),
									portId: sourcePorts[0].id,
									isEditting: false,
								},
							],
						},
					});
				}
			}
			if (c.targetHandle === "repository") {
				const sourceNode = nodes.find((n) => n.id === c.source);
				if (sourceNode && sourceNode.type === "github" && sourceNode.data.repository) {
					updateNodeData<"app">(tn.id, {
						repository: {
							id: sourceNode.data.repository.id,
							repoNodeId: c.source,
						},
					});
				}
			}
		}
		if (c.sourceHandle === "volume") {
			updateNodeData<"volume">(c.source, {
				attachedTo: ((sn as VolumeNode).data.attachedTo || []).concat(c.source),
			});
		}
		if (c.targetHandle === "https") {
			if ((sn.data.ports || []).length === 1) {
				updateNodeData<"gateway-https">(c.target, {
					https: {
						serviceId: c.source,
						portId: sn.data.ports![0].id,
					},
				});
			} else {
				updateNodeData<"gateway-https">(c.target, {
					https: {
						serviceId: c.source,
						portId: "", // TODO(gio)
					},
				});
			}
		}
		if (c.targetHandle === "tcp") {
			const td = tn.data as GatewayTCPData;
			if ((sn.data.ports || []).length === 1) {
				updateNodeData<"gateway-tcp">(c.target, {
					exposed: (td.exposed || []).concat({
						serviceId: c.source,
						portId: sn.data.ports![0].id,
					}),
				});
			} else {
				updateNodeData<"gateway-tcp">(c.target, {
					selected: {
						serviceId: c.source,
						portId: undefined,
					},
				});
			}
		}
		if (sn.type === "app") {
			if (c.sourceHandle === "ports") {
				updateNodeData<"app">(sn.id, {
					isChoosingPortToConnect: true,
				});
			}
		}
	}

	const fetchGithubRepositories = async () => {
		const { githubService, projectId } = get();
		if (!githubService || !projectId) {
			set({
				githubRepositories: [],
				githubRepositoriesError: "GitHub service or Project ID not available.",
				githubRepositoriesLoading: false,
			});
			return;
		}

		set({ githubRepositoriesLoading: true, githubRepositoriesError: null });
		try {
			const repos = await githubService.getRepositories();
			set({ githubRepositories: repos, githubRepositoriesLoading: false });
		} catch (error) {
			console.error("Failed to fetch GitHub repositories in store:", error);
			const errorMessage = error instanceof Error ? error.message : "Unknown error fetching repositories";
			set({ githubRepositories: [], githubRepositoriesError: errorMessage, githubRepositoriesLoading: false });
		}
	};

	const disconnectFromStateStream = () => {
		const { stateEventSource } = get();
		if (stateEventSource) {
			stateEventSource.close();
			set({ stateEventSource: null });
		}
	};

	const connectToStateStream = (projectId: string, mode: "deploy" | "edit") => {
		disconnectFromStateStream();

		const eventSource = new EventSource(
			`/api/project/${projectId}/state/stream/${mode === "edit" ? "draft" : "deploy"}`,
		);
		set({ stateEventSource: eventSource });

		eventSource.onmessage = (event) => {
			const { mode, viewport } = get();
			const inst = GraphSchema.parse(JSON.parse(event.data));
			let positionChanged = false;
			let nodes = inst.nodes;
			if (mode === "edit") {
				nodes = inst.nodes.map((n) => {
					if (n.position.x === 0 && n.position.y === 0) {
						positionChanged = true;
						return {
							...n,
							position: getRandomPosition(viewport),
						};
					} else {
						return n;
					}
				});
			}
			setN(nodes);
			set({ edges: inst.edges });
			injectNetworkNodes();
			if (positionChanged) {
				fetch(`/api/project/${projectId}/saved`, {
					method: "POST",
					headers: {
						"Content-Type": "application/json",
					},
					body: JSON.stringify({
						type: "graph",
						graph: {
							...inst,
							nodes,
						},
					}),
				});
			}
		};

		eventSource.onerror = (err) => {
			console.error("EventSource failed:", err);
			eventSource.close();
			set({ stateEventSource: null });
		};
	};

	return {
		projectId: undefined,
		mode: "edit",
		buildMode: "overview",
		projects: [],
		nodes: [],
		edges: [],
		categories: defaultCategories,
		messages: [],
		env: defaultEnv,
		viewport: {
			transformX: 0,
			transformY: 0,
			transformZoom: 1,
			width: 800,
			height: 600,
		},
		zoom: {
			x: 0,
			y: 0,
			zoom: 1,
		},
		githubService: null,
		githubRepositories: [],
		githubRepositoriesLoading: false,
		githubRepositoriesError: null,
		stateEventSource: null,
		setViewport: (viewport) => {
			const { viewport: vp } = get();
			if (
				viewport.transformX !== vp.transformX ||
				viewport.transformY !== vp.transformY ||
				viewport.transformZoom !== vp.transformZoom ||
				viewport.width !== vp.width ||
				viewport.height !== vp.height
			) {
				set({ viewport });
			}
		},
		setHighlightCategory: (name, active) => {
			set({
				categories: get().categories.map((c) => {
					if (c.title.toLowerCase() !== name.toLowerCase()) {
						return c;
					} else {
						return {
							...c,
							active,
						};
					}
				}),
			});
		},
		onNodesChange: (changes) => {
			const nodes = applyNodeChanges(changes, get().nodes);
			setN(nodes);
		},
		onEdgesChange: (changes) => {
			set({
				edges: applyEdgeChanges(changes, get().edges),
			});
		},
		addNode: (node) => {
			const { viewport, nodes } = get();
			setN(
				nodes.concat({
					...node,
					position: getRandomPosition(viewport),
				} as AppNode),
			);
		},
		setNodes: (nodes) => {
			setN(nodes);
		},
		setEdges: (edges) => {
			set({ edges });
		},
		replaceEdge: (c, id) => {
			let change: EdgeChange;
			if (id === undefined) {
				change = {
					type: "add",
					item: {
						id: uuidv4(),
						...c,
					},
				};
				onConnect(c);
			} else {
				change = {
					type: "replace",
					id,
					item: {
						id,
						...c,
					},
				};
			}
			set({
				edges: applyEdgeChanges([change], get().edges),
			});
		},
		updateNode,
		updateNodeData,
		onConnect,
		refreshEnv: async () => {
			const projectId = get().projectId;
			let env: Env = defaultEnv;
			try {
				if (projectId) {
					const response = await fetch(`/api/project/${projectId}/env`);
					if (response.ok) {
						const data = await response.json();
						const result = envSchema.safeParse(data);
						if (result.success) {
							env = result.data;
						} else {
							console.error("Invalid env data:", result.error);
						}
					}
				}
			} catch (error) {
				console.error("Failed to fetch integrations:", error);
			} finally {
				const oldEnv = get().env;
				const oldGithubIntegrationStatus = oldEnv.integrations.github;
				if (JSON.stringify(oldEnv) !== JSON.stringify(env)) {
					set({ env });
					injectNetworkNodes();
					let ghService = null;
					if (env.integrations.github) {
						ghService = new GitHubServiceImpl(projectId!);
					}
					if (get().githubService !== ghService || (ghService && !get().githubService)) {
						set({ githubService: ghService });
					}
					if (
						ghService &&
						(oldGithubIntegrationStatus !== env.integrations.github || !oldEnv.integrations.github)
					) {
						get().fetchGithubRepositories();
					}
					if (!env.integrations.github) {
						set({
							githubRepositories: [],
							githubRepositoriesError: null,
							githubRepositoriesLoading: false,
						});
					}
				}
			}
		},
		setMode: (mode) => {
			disconnectFromStateStream();
			set({ mode });
			const projectId = get().projectId;
			if (projectId) {
				connectToStateStream(projectId, mode);
			}
		},
		setBuildMode: (buildMode) => {
			set({ buildMode });
		},
		setProject: async (projectId) => {
			const currentProjectId = get().projectId;
			if (projectId === currentProjectId) {
				return;
			}
			stopRefreshEnvInterval();
			disconnectFromStateStream();
			set({
				projectId,
				githubRepositories: [],
				githubRepositoriesLoading: false,
				githubRepositoriesError: null,
			});
			if (projectId) {
				await get().refreshEnv();
				if (get().env.instanceId) {
					set({ mode: "deploy" });
				} else {
					set({ mode: "edit" });
				}
				const mode = get().mode;
				connectToStateStream(projectId, mode);
				startRefreshEnvInterval();
			} else {
				set({
					nodes: [],
					edges: [],
					env: defaultEnv,
					githubService: null,
					githubRepositories: [],
					githubRepositoriesLoading: false,
					githubRepositoriesError: null,
				});
			}
		},
		fetchGithubRepositories: fetchGithubRepositories,
	};
});