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

import { v4 as uuidv4 } from "uuid";
import { create } from 'zustand';
import { addEdge, applyNodeChanges, applyEdgeChanges, Connection, EdgeChange, useNodes } from '@xyflow/react';
import {
  type Edge,
  type Node,
  type OnNodesChange,
  type OnEdgesChange,
  type OnConnect,
} from '@xyflow/react';
import { DeepPartial } from "react-hook-form";
import { Category, defaultCategories } from "./categories";
import { CreateValidators, Validator } from "./config";
import { z } from "zod";

export type InitData = {
  label: string;
  envVars: BoundEnvVar[];
  ports: Port[];
};

export type NodeData = InitData & {
  activeField?: string | undefined;
};

export type PortConnectedTo = {
  serviceId: string;
  portId: string;
}

export type GatewayHttpsData = NodeData & {
  network?: string;
  subdomain?: string;
  https?: PortConnectedTo;
};

export type GatewayHttpsNode = Node<GatewayHttpsData> & {
  type: "gateway-https";
};

export type GatewayTCPData = NodeData & {
  network?: string;
  subdomain?: string;
  exposed: PortConnectedTo[];
  selected?: {
    serviceId?: string;
    portId?: string;
  };
};

export type GatewayTCPNode = Node<GatewayTCPData> & {
  type: "gateway-tcp";
};

export type Port = {
  id: string;
  name: string;
  value: number;
};

export const ServiceTypes = ["node-23.1.0", "nextjs:deno-2.0.0"] as const;
export type ServiceType = typeof ServiceTypes[number];

export type ServiceData = NodeData & {
  type: ServiceType;
  repository: {
    id: string;
    branch: string;
    rootDir: string;
  };
  env: string[];
  volume: string[];
  isChoosingPortToConnect: boolean;
};

export type ServiceNode = Node<ServiceData> & {
  type: "app";
};

export type VolumeType = "ReadWriteOnce" | "ReadOnlyMany" | "ReadWriteMany" | "ReadWriteOncePod";

export type VolumeData = NodeData & {
  type: VolumeType;
  size: string;
  attachedTo: string[];
};

export type VolumeNode = Node<VolumeData> & {
  type: "volume";
};

export type PostgreSQLData = NodeData & {
  volumeId: string;
};

export type PostgreSQLNode = Node<PostgreSQLData> & {
  type: "postgresql";
};

export type MongoDBData = NodeData & {
  volumeId: string;
};

export type MongoDBNode = Node<MongoDBData> & {
  type: "mongodb";
};

export type GithubData = NodeData & {
  address: string;
};

export type GithubNode = Node<GithubData> & {
  type: "github";
};

export type NANode = Node<NodeData> & {
  type: undefined;
};

export type AppNode = GatewayHttpsNode | GatewayTCPNode | ServiceNode | VolumeNode | PostgreSQLNode | MongoDBNode | GithubNode | NANode;

export function nodeLabel(n: AppNode): string {
  switch (n.type) {
    case "app": return n.data.label || "Service";
    case "github": return n.data.address || "Github";
    case "gateway-https": {
      if (n.data && n.data.network && n.data.subdomain) {
        return `https://${n.data.subdomain}.${n.data.network}`;
      } else {
        return "HTTPS Gateway";
      }
    }
    case "gateway-tcp": {
      if (n.data && n.data.network && n.data.subdomain) {
        return `${n.data.subdomain}.${n.data.network}`;
      } 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("MUST NOT REACH!");
  }
}

export function nodeIsConnectable(n: AppNode, handle: string): boolean {
  switch (n.type) {
    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 !== undefined && n.data.address) {
        return true;
      }
      return false;
    case "gateway-https":
      return n.data === undefined || n.data.https === undefined;
    case "gateway-tcp":
      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("MUST NOT REACH!");
  }
}

export type BoundEnvVar = {
  id: string;
  source: string | null;
} | {
  id: string;
  source: string | null;
  name: string;
  isEditting: boolean;
} | {
  id: string;
  source: string | null;
  name: string;
  alias: string;
  isEditting: boolean;
} | {
  id: string;
  source: string | null;
  portId: string;
  name: string;
  alias: string;
  isEditting: boolean;
};

export type EnvVar = {
  name: string;
  value: string;
};

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");
  }
}

export type NodeType = Exclude<Pick<AppNode, "type">["type"], undefined>;

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;
};

export const envSchema = z.object({
  deployKey: z.string(),
  networks: z.array(z.object({
    name: z.string(),
    domain: z.string(),
  })),
});

export type Env = z.infer<typeof envSchema>;

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

export type AppState = {
  projectId: string | undefined;
  projects: Project[];
  nodes: AppNode[];
  edges: Edge[];
  categories: Category[];
  messages: Message[];
  env?: Env;
  setHighlightCategory: (name: string, active: boolean) => void;
  onNodesChange: OnNodesChange<AppNode>;
  onEdgesChange: OnEdgesChange;
  onConnect: OnConnect;
  setNodes: (nodes: AppNode[]) => void;
  setEdges: (edges: Edge[]) => void;
  setProject: (projectId: string) => void;
  setProjects: (projects: Project[]) => void;
  updateNode: <T extends NodeType>(id: string, data: DeepPartial<(AppNode & (Pick<AppNode, "type"> | { type: T }))>) => void;
  updateNodeData: <T extends NodeType>(id: string, data: DeepPartial<(AppNode & (Pick<AppNode, "type"> | { type: T }))["data"]>) => void;
  replaceEdge: (c: Connection, id?: string) => void;
  refreshEnv: () => Promise<Env | undefined>;
};

const projectIdSelector = (state: AppState) => state.projectId;
const categoriesSelector = (state: AppState) => state.categories;
const messagesSelector = (state: AppState) => state.messages;
const envSelector = (state: AppState) => state.env;

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

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;
}

let envRefresh: Promise<Env | undefined> | null = null;

export function useEnv(): Env {
  return {
    "deployKey": "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIPK58vMu0MwIzdZT+mqpIBkhl48p9+/YwDCZv7MgTesF",
    "networks": [{
      "name": "Public",
      "domain": "v1.dodo.cloud",
    }, {
      "name": "Private",
      "domain": "p.v1.dodo.cloud",
    }],
  };
  const store = useStateStore();
  const env = envSelector(store);
  console.log(env);
  if (env != null) {
    return env;
  }
  if (envRefresh == null) {
    envRefresh = store.refreshEnv();
    envRefresh.finally(() => envRefresh = null);
  }
  return {
    deployKey: "",
    networks: [],
  };
}

const v: Validator = CreateValidators();

export const useStateStore = create<AppState>((set, get): AppState => {
  set({ env: {
    "deployKey": "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIPK58vMu0MwIzdZT+mqpIBkhl48p9+/YwDCZv7MgTesF",
    "networks": [{
      "name": "Public",
      "domain": "v1.dodo.cloud",
    }, {
      "name": "Private",
      "domain": "p.v1.dodo.cloud",
    }],
  }});
  console.log(get().env);
  const setN = (nodes: AppNode[]) => {
    set({
      nodes: nodes,
      messages: v(nodes),
    })
  };
  function updateNodeData<T extends NodeType>(id: string, d: DeepPartial<(AppNode & (Pick<AppNode, "type"> | { type: T }))["data"]>): void {
    setN(get().nodes.map((n) => {
        if (n.id !== id) {
          return n;
        }
        const nd = {
          ...n,
          data: {
            ...n.data,
            ...d,
          },
        };
        return nd;
      })
    );
  };
  function updateNode<T extends NodeType>(id: string, d: DeepPartial<(AppNode & (Pick<AppNode, "type"> | { type: T }))>): void {
    setN(
      get().nodes.map((n) => {
        if (n.id !== id) {
          return n;
        }
        return {
          ...n,
          ...d,
        };
      })
    );
  };
  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 (c.sourceHandle === "env_var" && c.targetHandle === "env_var") {
      const sourceEnvVars = nodeEnvVarNames(sn);
      if (sourceEnvVars.length === 0) {
        throw new Error("MUST NOT REACH!");
      }
      const id = uuidv4();
      if (sourceEnvVars.length === 1) {
        updateNode(c.target, {
          ...tn,
          data: {
            ...tn.data,
            envVars: [
              ...(tn.data.envVars || []),
              {
                id: id,
                source: c.source,
                name: sourceEnvVars[0],
                isEditting: false,
              },
            ],
          },
        });
      } else {
        updateNode(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(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.sourceHandle === "volume") {
      updateNodeData<"volume">(c.source, {
        attachedTo: ((sn as VolumeNode).data.attachedTo || []).concat(c.source),
      });
    }
    if (c.targetHandle === "volume") {
      if (tn.type === "postgresql" || tn.type === "mongodb") {
        updateNodeData(c.target, {
          volumeId: 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,
        });
      }
    }
    if (tn.type === "app") {
      if (c.targetHandle === "repository") {
        updateNodeData<"app">(tn.id, {
          repository: {
            id: c.source,
            branch: "master",
            rootDir: "/",
          }
        });
      }
    }
  }
  return {
    projectId: undefined,
    projects: [],
    nodes: [],
    edges: [],
    categories: defaultCategories,
    messages: v([]),
    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),
      });
    },
    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 () => {
      return get().env;
        const resp = await fetch("/env");
        if (!resp.ok) {
          throw new Error("failed to fetch env config");
        }
        set({ env: envSchema.parse(await resp.json()) });
        return get().env;
    },
    setProject: (projectId) => set({ projectId }),
    setProjects: (projects) => set({ projects }),
  };
});