supabase/apps/studio/components/interfaces/BranchManagement/EditBranchModal.tsx

import { zodResolver } from '@hookform/resolvers/zod'
import { useDebounce } from '@uidotdev/usehooks'
import { useParams } from 'common'
import { Check, Loader2 } from 'lucide-react'
import Image from 'next/image'
import { useCallback, useEffect, useState } from 'react'
import { useForm, useWatch } from 'react-hook-form'
import { toast } from 'sonner'
import {
  Button,
  cn,
  Dialog,
  DialogContent,
  DialogFooter,
  DialogHeader,
  DialogSection,
  DialogSectionSeparator,
  DialogTitle,
  Form,
  FormControl,
  FormField,
  Input,
} from 'ui'
import { FormItemLayout } from 'ui-patterns/form/FormItemLayout/FormItemLayout'
import { ShimmeringLoader } from 'ui-patterns/ShimmeringLoader'
import * as z from 'zod'

import { ConnectToGitHub } from './ConnectToGitHub'
import { AlertError } from '@/components/ui/AlertError'
import { InlineLink } from '@/components/ui/InlineLink'
import { useBranchUpdateMutation } from '@/data/branches/branch-update-mutation'
import { Branch, useBranchesQuery } from '@/data/branches/branches-query'
import { useGitHubAuthorizationQuery } from '@/data/integrations/github-authorization-query'
import { useCheckGithubBranchValidity } from '@/data/integrations/github-branch-check-query'
import { useGitHubConnectionsQuery } from '@/data/integrations/github-connections-query'
import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import { BASE_PATH } from '@/lib/constants'

interface EditBranchModalProps {
  branch?: Branch
  visible: boolean
  onClose: () => void
}

export const EditBranchModal = ({ branch, visible, onClose }: EditBranchModalProps) => {
  const { ref } = useParams()
  const { data: projectDetails } = useSelectedProjectQuery()
  const { data: selectedOrg } = useSelectedOrganizationQuery()

  const [isGitBranchValid, setIsGitBranchValid] = useState(true)

  const isBranch = projectDetails?.parent_project_ref !== undefined
  const projectRef =
    projectDetails !== undefined ? (isBranch ? projectDetails.parent_project_ref : ref) : undefined

  const {
    data: githubAuthorization,
    error: authorizationError,
    isPending: isLoadingAuthorization,
    isSuccess: isSuccessAuthorization,
    isError: isErrorAuthorization,
  } = useGitHubAuthorizationQuery()

  const {
    data: connections,
    error: connectionsError,
    isPending: isLoadingConnections,
    isSuccess: isSuccessConnections,
    isError: isErrorConnections,
  } = useGitHubConnectionsQuery({
    organizationId: selectedOrg?.id,
  })

  const isLoading = isLoadingAuthorization || isLoadingConnections
  const isSuccess = isSuccessAuthorization && isSuccessConnections
  const isError = isErrorAuthorization || isErrorConnections
  const error = authorizationError || connectionsError

  const { data: branches } = useBranchesQuery({ projectRef })
  const { mutate: checkGithubBranchValidity, isPending: isChecking } = useCheckGithubBranchValidity(
    { onError: () => {} }
  )

  const { mutate: updateBranch, isPending: isUpdating } = useBranchUpdateMutation({
    onSuccess: (data) => {
      toast.success(`Successfully updated branch "${data.name}"`)
      onClose()
    },
    onError: (error) => {
      toast.error(`Failed to update branch: ${error.message}`)
    },
  })

  const githubConnection = connections?.find((connection) => connection.project.ref === projectRef)
  const [repoOwner, repoName] = githubConnection?.repository.name.split('/') ?? []

  const formId = 'edit-branch-form'
  const FormSchema = z.object({
    branchName: z
      .string()
      .min(1, 'Branch name cannot be empty')
      .refine(
        (val) => /^[a-zA-Z0-9\-_]+$/.test(val),
        'Branch name can only contain alphanumeric characters, hyphens, and underscores.'
      )
      .refine(
        (val) =>
          // Allow the current branch name during edit
          val === branch?.name || (branches ?? []).every((b) => b.name !== val),
        'A branch with this name already exists'
      ),
    gitBranchName: z.string().optional(),
  })

  const form = useForm<z.infer<typeof FormSchema>>({
    mode: 'onChange',
    reValidateMode: 'onChange',
    resolver: zodResolver(FormSchema),
    defaultValues: { branchName: '', gitBranchName: '' },
  })
  const gitBranchName = useWatch({ control: form.control, name: 'gitBranchName' })
  const debouncedGitBranchName = useDebounce(gitBranchName, 500)

  const isFormValid = form.formState.isValid && (!gitBranchName || isGitBranchValid)
  const canSubmit = isFormValid && !isUpdating && !isChecking

  const onSubmit = (data: z.infer<typeof FormSchema>) => {
    if (!projectRef) return console.error('Project ref is required')
    if (!branch?.project_ref) return console.error('Branch ref is required')

    const payload: {
      branchRef: string
      projectRef: string
      branchName: string
      gitBranch?: string
    } = {
      branchRef: branch.project_ref,
      projectRef,
      branchName: data.branchName,
    }

    // Only add gitBranch to the payload if it is present and valid
    // If gitBranchName is empty or invalid, gitBranch remains undefined in the payload
    if (data.gitBranchName && isGitBranchValid) {
      payload.gitBranch = data.gitBranchName
    }

    updateBranch(payload)
  }

  const validateGitBranchName = useCallback(
    (branchName: string) => {
      if (!githubConnection)
        return console.error(
          '[EditBranchModal > validateGitBranchName] GitHub Connection is missing'
        )

      const repositoryId = githubConnection.repository.id
      const requested = branchName
      checkGithubBranchValidity(
        { repositoryId, branchName },
        {
          onSuccess: () => {
            if (form.getValues('gitBranchName') !== requested) return

            // Check if another branch is already linked to this git branch
            const existingBranch = (branches ?? []).find(
              (b) => b.git_branch === branchName && b.id !== branch?.id
            )
            if (existingBranch) {
              setIsGitBranchValid(false)
              form.setError('gitBranchName', {
                message: `Branch "${existingBranch.name}" is already linked to git branch "${branchName}"`,
              })
              return
            }

            setIsGitBranchValid(true)
            form.clearErrors('gitBranchName')
          },
          onError: (error) => {
            if (form.getValues('gitBranchName') !== requested) return
            setIsGitBranchValid(false)
            form.setError('gitBranchName', {
              message:
                error?.message ??
                `Unable to find branch "${branchName}" in ${repoOwner}/${repoName}`,
            })
          },
        }
      )
    },
    [githubConnection, form, checkGithubBranchValidity, repoOwner, repoName, branches, branch]
  )

  // Pre-fill form when the modal becomes visible and branch data is available
  useEffect(() => {
    if (visible && branch) {
      form.reset({
        branchName: branch.name ?? '',
        gitBranchName: branch.git_branch ?? '',
      })
    }
  }, [branch, visible, form])

  useEffect(() => {
    if (!githubConnection || !debouncedGitBranchName) {
      return form.clearErrors('gitBranchName')
    }

    form.clearErrors('gitBranchName')
    validateGitBranchName(debouncedGitBranchName)
  }, [debouncedGitBranchName, validateGitBranchName, form, githubConnection])

  return (
    <Dialog open={visible} onOpenChange={(open) => !open && onClose()}>
      <DialogContent size="large" hideClose>
        <DialogHeader padding="small">
          <DialogTitle>Edit branch "{branch?.name}"</DialogTitle>
        </DialogHeader>
        <DialogSectionSeparator />
        <Form {...form}>
          <form id={formId} onSubmit={form.handleSubmit(onSubmit)}>
            <DialogSection padding="medium" className="space-y-4">
              <FormField
                control={form.control}
                name="branchName"
                render={({ field }) => (
                  <FormItemLayout label="Preview branch name">
                    <FormControl>
                      <Input
                        {...field}
                        placeholder="e.g. staging, dev-feature-x"
                        autoComplete="off"
                      />
                    </FormControl>
                  </FormItemLayout>
                )}
              />

              {isLoading && (
                <div className="flex flex-col gap-y-2">
                  <ShimmeringLoader />
                  <ShimmeringLoader className="w-1/2" />
                </div>
              )}

              {isError && (
                <AlertError
                  error={error}
                  subject="Failed to retrieve GitHub connection information"
                />
              )}

              {isSuccess &&
                (!githubAuthorization || !githubConnection ? (
                  <ConnectToGitHub />
                ) : (
                  <FormField
                    control={form.control}
                    name="gitBranchName"
                    render={({ field }) => (
                      <FormItemLayout
                        label={
                          <div className="flex items-center justify-between w-full gap-4">
                            <span className="flex-1">Sync with Git branch</span>
                            <div className="flex items-center gap-2 text-sm">
                              <Image
                                className={cn('dark:invert')}
                                src={`${BASE_PATH}/img/icons/github-icon.svg`}
                                width={16}
                                height={16}
                                alt={`GitHub icon`}
                              />
                              <InlineLink href={`https://github.com/${repoOwner}/${repoName}`}>
                                {repoOwner}/{repoName}
                              </InlineLink>
                            </div>
                          </div>
                        }
                        labelOptional="Optional"
                        description="Automatically deploy changes on every commit"
                      >
                        <div className="relative">
                          <FormControl>
                            <Input
                              {...field}
                              placeholder="e.g. main, feat/some-feature"
                              autoComplete="off"
                              onChange={(e) => {
                                field.onChange(e)
                                setIsGitBranchValid(false)
                              }}
                            />
                          </FormControl>
                          <div className="absolute top-2.5 right-3 flex items-center gap-2">
                            {field.value ? (
                              isChecking ? (
                                <Loader2 size={14} className="animate-spin" />
                              ) : isGitBranchValid ? (
                                <Check size={14} className="text-primary" strokeWidth={2} />
                              ) : null
                            ) : null}
                          </div>
                        </div>
                      </FormItemLayout>
                    )}
                  />
                ))}
            </DialogSection>

            <DialogFooter padding="medium">
              <Button disabled={isUpdating} onClick={onClose}>
                Cancel
              </Button>
              <Button
                form={formId}
                disabled={
                  (!!gitBranchName && !isSuccessConnections) ||
                  isUpdating ||
                  !canSubmit ||
                  isChecking
                }
                loading={isUpdating}
                variant="primary"
                type="submit"
              >
                Update branch
              </Button>
            </DialogFooter>
          </form>
        </Form>
      </DialogContent>
    </Dialog>
  )
}