Briven/apps/studio/components/interfaces/Settings/Infrastructure/InfrastructureConfiguration/InstanceConfiguration.tsx

import { PermissionAction } from '@supabase/shared-types/out/constants'
import {
  Background,
  ColorMode,
  Edge,
  ReactFlow,
  ReactFlowProvider,
  useNodesInitialized,
  useReactFlow,
} from '@xyflow/react'
import { partition } from 'lodash'
import { ChevronDown, Globe2, Loader2, Network } from 'lucide-react'
import { useTheme } from 'next-themes'
import Link from 'next/link'
import { useEffect, useMemo, useState } from 'react'

import '@xyflow/react/dist/style.css'

import { useParams } from 'common'
import { useRouter } from 'next/router'
import {
  Button,
  cn,
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuSeparator,
  DropdownMenuTrigger,
} from 'ui'

import DropAllReplicasConfirmationModal from './DropAllReplicasConfirmationModal'
import { DropReplicaConfirmationModal } from './DropReplicaConfirmationModal'
import { SmoothstepEdge } from './Edge'
import { REPLICA_STATUS } from './InstanceConfiguration.constants'
import { addRegionNodes, generateNodes, getDagreGraphLayout } from './InstanceConfiguration.utils'
import { LoadBalancerNode, PrimaryNode, RegionNode, ReplicaNode } from './InstanceNode'
import MapView from './MapView'
import { RestartReplicaConfirmationModal } from './RestartReplicaConfirmationModal'
import AlertError from '@/components/ui/AlertError'
import { ButtonTooltip } from '@/components/ui/ButtonTooltip'
import { useLoadBalancersQuery } from '@/data/read-replicas/load-balancers-query'
import { Database, useReadReplicasQuery } from '@/data/read-replicas/replicas-query'
import {
  ReplicaInitializationStatus,
  useReadReplicasStatusesQuery,
} from '@/data/read-replicas/replicas-status-query'
import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions'
import { useIsFeatureEnabled } from '@/hooks/misc/useIsFeatureEnabled'
import {
  useIsAwsCloudProvider,
  useIsOrioleDb,
  useSelectedProjectQuery,
} from '@/hooks/misc/useSelectedProject'
import { useStaticEffectEvent } from '@/hooks/useStaticEffectEvent'
import { timeout } from '@/lib/helpers'

interface InstanceConfigurationUIProps {
  diagramOnly?: boolean
}

const InstanceConfigurationUI = ({ diagramOnly = false }: InstanceConfigurationUIProps) => {
  const router = useRouter()
  const reactFlow = useReactFlow()
  const isOrioleDb = useIsOrioleDb()
  const { resolvedTheme } = useTheme()
  const { ref: projectRef } = useParams()
  const { isPending: isLoadingProject } = useSelectedProjectQuery()

  const isAws = useIsAwsCloudProvider()
  const { infrastructureReadReplicas } = useIsFeatureEnabled(['infrastructure:read_replicas'])
  const newReplicaURL = `/project/${projectRef}/database/replication?type=Read+Replica`

  const [view, setView] = useState<'flow' | 'map'>('flow')
  const [showDeleteAllModal, setShowDeleteAllModal] = useState(false)
  const [refetchInterval, setRefetchInterval] = useState<number | false>(10000)
  const [selectedReplicaToDrop, setSelectedReplicaToDrop] = useState<Database>()
  const [selectedReplicaToRestart, setSelectedReplicaToRestart] = useState<Database>()

  const { can: canManageReplicas } = useAsyncCheckPermissions(PermissionAction.CREATE, 'projects')

  const {
    data: loadBalancers,
    refetch: refetchLoadBalancers,
    isSuccess: isSuccessLoadBalancers,
  } = useLoadBalancersQuery({ projectRef })
  const {
    data,
    error,
    refetch: refetchReplicas,
    isPending: isLoading,
    isError,
    isSuccess: isSuccessReplicas,
  } = useReadReplicasQuery({ projectRef })
  const [[primary], replicas] = useMemo(
    () => partition(data ?? [], (db) => db.identifier === projectRef),
    [data, projectRef]
  )
  const numReplicas = useMemo(() => data?.length ?? 0, [data])

  const { data: replicasStatuses, isSuccess: isSuccessReplicasStatuses } =
    useReadReplicasStatusesQuery(
      { projectRef },
      {
        refetchInterval: refetchInterval,
        refetchOnWindowFocus: false,
      }
    )

  useEffect(() => {
    if (!isSuccessReplicasStatuses) return
    const refetch = async () => {
      const fixedStatues = [
        REPLICA_STATUS.ACTIVE_HEALTHY,
        REPLICA_STATUS.ACTIVE_UNHEALTHY,
        REPLICA_STATUS.INIT_READ_REPLICA_FAILED,
      ]
      const replicasInTransition = replicasStatuses.filter((db) => {
        const { status } = db.replicaInitializationStatus || {}
        return (
          !fixedStatues.includes(db.status) || status === ReplicaInitializationStatus.InProgress
        )
      })
      const hasTransientStatus = replicasInTransition.length > 0

      // If any replica's status has changed, refetch databases
      if (replicasStatuses.length !== numReplicas) {
        await refetchReplicas()
        setTimeout(() => refetchLoadBalancers(), 2000)
      }

      // If all replicas are active healthy, stop fetching statuses
      if (!hasTransientStatus) {
        setRefetchInterval(false)
      }
    }
    refetch()
  }, [
    numReplicas,
    isSuccessReplicasStatuses,
    refetchLoadBalancers,
    refetchReplicas,
    replicasStatuses,
  ])

  const backgroundPatternColor =
    resolvedTheme === 'dark' ? 'rgba(255, 255, 255, 0.3)' : 'rgba(0, 0, 0, 0.4)'

  const nodes = useMemo(
    () =>
      isSuccessReplicas && isSuccessLoadBalancers && primary !== undefined
        ? generateNodes({
            primary,
            replicas,
            loadBalancers: loadBalancers ?? [],
            onSelectRestartReplica: setSelectedReplicaToRestart,
            onSelectDropReplica: setSelectedReplicaToDrop,
          })
        : [],
    [isSuccessReplicas, isSuccessLoadBalancers, primary, replicas, loadBalancers]
  )

  const edges: Edge[] = useMemo(
    () =>
      isSuccessReplicas && isSuccessLoadBalancers
        ? [
            ...((loadBalancers ?? []).length > 0
              ? [
                  {
                    id: `load-balancer-${primary.identifier}`,
                    source: 'load-balancer',
                    target: primary.identifier,
                    type: 'smoothstep',
                    animated: true,
                    className: 'cursor-default!',
                  },
                ]
              : []),
            ...replicas.map((database) => {
              return {
                id: `${primary.identifier}-${database.identifier}`,
                source: primary.identifier,
                target: database.identifier,
                type: 'smoothstep',
                animated: true,
                className: 'cursor-default!',
                data: {
                  status: database.status,
                  identifier: database.identifier,
                  connectionString: database.connectionString,
                },
              }
            }),
          ]
        : [],
    [isSuccessLoadBalancers, isSuccessReplicas, loadBalancers, primary?.identifier, replicas]
  )

  const nodeTypes = useMemo(
    () => ({
      PRIMARY: PrimaryNode,
      READ_REPLICA: ReplicaNode,
      REGION: RegionNode,
      LOAD_BALANCER: LoadBalancerNode,
    }),
    []
  )

  const edgeTypes = useMemo(
    () => ({
      smoothstep: SmoothstepEdge,
    }),
    []
  )

  const nodesInitialized = useNodesInitialized()
  const [hasMeasuredLayout, setHasMeasuredLayout] = useState(false)

  const setReactFlow = async ({ measured }: { measured: boolean }) => {
    // Merge in React Flow's measured dimensions (if any) so dagre can use real
    // heights instead of the first-paint fallbacks.
    const measuredNodes = nodes.map((node) => {
      const existing = reactFlow.getNode(node.id)
      return existing?.measured ? { ...node, measured: existing.measured } : node
    })
    const graph = getDagreGraphLayout(measuredNodes, edges)
    const { nodes: updatedNodes } = addRegionNodes(graph.nodes, graph.edges)
    reactFlow.setNodes(updatedNodes)
    reactFlow.setEdges(graph.edges)

    // [Joshen] Odd fix to ensure that react flow snaps back to center when adding nodes
    await timeout(1)
    reactFlow.fitView({ maxZoom: 0.9, minZoom: 0.9 })
    if (measured) setHasMeasuredLayout(true)
  }

  // First pass: lay out using fallback heights for any not-yet-measured nodes.
  // The diagram is kept invisible until the measured pass below has run, so the
  // user never sees the fallback positions.
  // [Joshen] Just FYI this block is oddly triggering whenever we refocus on the viewport
  // even if I change the dependency array to just data. Not blocker, just an area to optimize
  useEffect(() => {
    if (isSuccessReplicas && isSuccessLoadBalancers && nodes.length > 0 && view === 'flow') {
      setReactFlow({ measured: false })
    }
  }, [isSuccessReplicas, isSuccessLoadBalancers, nodes, edges, view])

  // Second pass: once React Flow has measured the nodes, re-run the layout so
  // dagre uses real heights. Only `nodesInitialized` going true should trigger
  // this — the first-pass effect above handles node/view changes.
  const runMeasuredLayout = useStaticEffectEvent(() => {
    if (nodesInitialized && nodes.length > 0 && view === 'flow') {
      setReactFlow({ measured: true })
    }
  })
  useEffect(() => {
    runMeasuredLayout()
  }, [nodesInitialized, runMeasuredLayout])

  return (
    <div className={cn('nowheel', diagramOnly ? 'h-full' : 'border-y')}>
      <div
        className={`${diagramOnly ? 'h-full' : 'h-[500px]'} w-full relative ${
          isSuccessReplicas && !isLoadingProject ? '' : 'flex items-center justify-center px-28'
        }`}
      >
        {(isLoading || isLoadingProject) && (
          <Loader2 className="animate-spin text-foreground-light" />
        )}
        {isError && <AlertError error={error} subject="Failed to retrieve replicas" />}
        {isSuccessReplicas && !isLoadingProject && (
          <>
            {!diagramOnly && infrastructureReadReplicas && (
              <div className="z-10 absolute top-4 right-4 flex items-center justify-center gap-x-2">
                <div className="flex items-center justify-center">
                  <ButtonTooltip
                    asChild
                    type="default"
                    disabled={!canManageReplicas || isOrioleDb}
                    className={cn(replicas.length > 0 ? 'rounded-r-none' : '')}
                    tooltip={{
                      content: {
                        side: 'bottom',
                        text: !canManageReplicas
                          ? 'You need additional permissions to deploy replicas'
                          : isOrioleDb
                            ? 'Read replicas are not supported with OrioleDB'
                            : undefined,
                      },
                    }}
                  >
                    <Link href={newReplicaURL}>Deploy a new replica</Link>
                  </ButtonTooltip>
                  {replicas.length > 0 && (
                    <DropdownMenu>
                      <DropdownMenuTrigger asChild>
                        <Button
                          type="default"
                          icon={<ChevronDown size={16} />}
                          className="px-1 rounded-l-none border-l-0"
                        />
                      </DropdownMenuTrigger>
                      <DropdownMenuContent align="end" className="w-52 *:space-x-2">
                        <DropdownMenuItem asChild>
                          <Link href={`/project/${projectRef}/settings/compute-and-disk`}>
                            Resize databases
                          </Link>
                        </DropdownMenuItem>
                        <DropdownMenuSeparator />
                        <DropdownMenuItem onClick={() => setShowDeleteAllModal(true)}>
                          <div>Remove all replicas</div>
                        </DropdownMenuItem>
                      </DropdownMenuContent>
                    </DropdownMenu>
                  )}
                </div>
                {isAws && (
                  <div className="flex items-center justify-center">
                    <Button
                      type="default"
                      icon={<Network size={15} />}
                      className={`rounded-r-none transition ${
                        view === 'flow' ? 'opacity-100' : 'opacity-50'
                      }`}
                      onClick={() => setView('flow')}
                    />
                    <Button
                      type="default"
                      icon={<Globe2 size={15} />}
                      className={`rounded-l-none transition ${
                        view === 'map' ? 'opacity-100' : 'opacity-50'
                      }`}
                      onClick={() => setView('map')}
                    />
                  </div>
                )}
              </div>
            )}
            {view === 'flow' ? (
              <ReactFlow
                // FIXME: https://github.com/xyflow/xyflow/issues/4876
                colorMode={'' as unknown as ColorMode}
                fitView
                fitViewOptions={{ minZoom: 0.9, maxZoom: 0.9 }}
                // Keep the diagram invisible (but laid out, so nodes can be
                // measured) until the measured-height layout pass has run.
                className={cn(
                  'instance-configuration transition-opacity duration-150',
                  hasMeasuredLayout ? 'opacity-100' : 'opacity-0'
                )}
                zoomOnPinch={false}
                zoomOnScroll={false}
                nodesDraggable={false}
                nodesConnectable={false}
                zoomOnDoubleClick={false}
                edgesFocusable={false}
                edgesReconnectable={false}
                defaultNodes={[]}
                defaultEdges={[]}
                nodeTypes={nodeTypes}
                edgeTypes={edgeTypes}
                proOptions={{ hideAttribution: true }}
              >
                <Background color={backgroundPatternColor} />
              </ReactFlow>
            ) : (
              <MapView
                onSelectDeployNewReplica={() => router.push(newReplicaURL)}
                onSelectRestartReplica={setSelectedReplicaToRestart}
                onSelectDropReplica={setSelectedReplicaToDrop}
              />
            )}
          </>
        )}
      </div>

      {!diagramOnly && (
        <>
          <DropReplicaConfirmationModal
            selectedReplica={selectedReplicaToDrop}
            onSuccess={() => setRefetchInterval(5000)}
            onCancel={() => setSelectedReplicaToDrop(undefined)}
          />

          <DropAllReplicasConfirmationModal
            visible={showDeleteAllModal}
            onSuccess={() => setRefetchInterval(5000)}
            onCancel={() => setShowDeleteAllModal(false)}
          />

          <RestartReplicaConfirmationModal
            selectedReplica={selectedReplicaToRestart}
            onSuccess={() => setRefetchInterval(5000)}
            onCancel={() => setSelectedReplicaToRestart(undefined)}
          />
        </>
      )}
    </div>
  )
}

interface InstanceConfigurationProps {
  diagramOnly?: boolean
}

export const InstanceConfiguration = ({ diagramOnly = false }: InstanceConfigurationProps) => {
  return (
    <ReactFlowProvider>
      <InstanceConfigurationUI diagramOnly={diagramOnly} />
    </ReactFlowProvider>
  )
}