paseo/packages/app/src/components/workspace-setup-dialog.tsx

import { KeyboardTranslateView } from "@/keyboard/shift";
import type { CreateWorkspaceRequestOptions } from "@getpaseo/client/internal/daemon-client";
import type { AgentSnapshotPayload } from "@getpaseo/protocol/messages";
import { getHostRuntimeStore } from "@/runtime/host-runtime";
import { useCallback, useEffect, useMemo, useState } from "react";
import { Text, View } from "react-native";
import { StyleSheet } from "react-native-unistyles";
import { useTranslation } from "react-i18next";
import { AdaptiveModalSheet, type SheetHeader } from "@/components/adaptive-modal-sheet";
import { FileDropZone } from "@/components/file-drop/file-drop-zone";
import { Composer } from "@/composer";
import { ProjectIconView } from "@/components/project-icon-view";
import { ICON_SIZE } from "@/styles/theme";
import { useToast } from "@/contexts/toast-context";
import { useAgentInputDraft } from "@/composer/draft/input-draft";
import { useProjectIcon } from "@/projects/icons";
import { useHostRuntimeClient, useHostRuntimeIsConnected } from "@/runtime/host-runtime";
import {
  normalizeWorkspaceDescriptor,
  useSessionStore,
  type WorkspaceDescriptor,
} from "@/stores/session-store";
import { useWorkspaceSetupStore } from "@/stores/workspace-setup-store";
import { normalizeAgentSnapshot } from "@/utils/agent-snapshots";
import { applyLegacyDaemonWorkspaceOwnership } from "@/workspace/legacy-daemon-workspaces";
import { encodeImages } from "@/utils/encode-images";
import { toErrorMessage } from "@/utils/error-messages";
import {
  resolveComposerAttachmentSubmitFormat,
  splitComposerAttachmentsForSubmit,
} from "@/composer/attachments/submit";
import type {
  CreateAgentRequestOptions,
  DaemonClient,
} from "@getpaseo/client/internal/daemon-client";
import { projectIconPlaceholderLabelFromDisplayName } from "@/utils/project-display-name";
import { navigateToAgent } from "@/utils/navigate-to-agent";
import { navigateToWorkspace } from "@/stores/navigation-active-workspace-store";
import type { MessagePayload } from "@/composer/types";

function toProjectIconDataUri(icon: { mimeType: string; data: string } | null): string | null {
  if (!icon) {
    return null;
  }
  return `data:${icon.mimeType};base64,${icon.data}`;
}

const SNAP_POINTS: string[] = ["82%", "94%"];

function resolveWorkspaceTitle({
  workspace,
  displayName,
  sourceDirectory,
}: {
  workspace: { name?: string | null; projectDisplayName?: string | null } | null;
  displayName: string;
  sourceDirectory: string;
}): string {
  return (
    workspace?.name ||
    workspace?.projectDisplayName ||
    displayName ||
    sourceDirectory.split(/[\\/]/).findLast(Boolean) ||
    sourceDirectory
  );
}

function buildChatDraftComposerArgs({
  serverId,
  workspaceDirectory,
  sourceDirectory,
  pendingWorkspaceSetup,
}: {
  serverId: string;
  workspaceDirectory: string | undefined;
  sourceDirectory: string;
  pendingWorkspaceSetup: { creationMethod: string } | null;
}) {
  return {
    initialServerId: serverId || null,
    isVisible: pendingWorkspaceSetup !== null,
    lockedWorkingDir: workspaceDirectory || sourceDirectory || undefined,
  };
}

async function callWorkspaceCreation({
  creationMethod,
  connectedClient,
  creationId,
  worktreeSlug,
  input,
}: {
  creationMethod: "create_worktree" | "open_project";
  connectedClient: DaemonClient;
  creationId: string;
  worktreeSlug: string;
  input: { cwd: string; agent?: CreateWorkspaceRequestOptions["agent"] };
}) {
  return connectedClient.createWorkspace({
    idempotencyKey: creationId,
    agent: input.agent,
    source:
      creationMethod === "create_worktree"
        ? { kind: "worktree", cwd: input.cwd, worktreeSlug }
        : { kind: "directory", path: input.cwd },
  });
}

function failureMessageForCreationMethod(
  method: "create_worktree" | "open_project",
  t: ReturnType<typeof useTranslation>["t"],
) {
  return method === "create_worktree"
    ? t("workspaceSetup.errors.failedCreateWorktree")
    : t("workspaceSetup.errors.failedOpenProject");
}

function buildCreateAgentOptions({
  composerState,
  text,
  attachments,
  encodedImages,
  workspaceDirectory,
  workspaceId,
  provider,
}: {
  composerState: {
    effectiveModeId: string;
    effectiveModelId: string | null;
    effectiveThinkingOptionId: string | null;
  };
  text: string;
  attachments: NonNullable<CreateAgentRequestOptions["attachments"]>;
  encodedImages: NonNullable<CreateAgentRequestOptions["images"]> | null;
  workspaceDirectory: string;
  workspaceId: string;
  provider: CreateAgentRequestOptions["provider"];
}): CreateAgentRequestOptions {
  return {
    provider,
    cwd: workspaceDirectory,
    workspaceId,
    ...(composerState.effectiveModeId ? { modeId: composerState.effectiveModeId } : {}),
    ...(composerState.effectiveModelId ? { model: composerState.effectiveModelId } : {}),
    ...(composerState.effectiveThinkingOptionId
      ? { thinkingOptionId: composerState.effectiveThinkingOptionId }
      : {}),
    ...(text.trim() ? { initialPrompt: text.trim() } : {}),
    ...(encodedImages && encodedImages.length > 0 ? { images: encodedImages } : {}),
    ...(attachments.length > 0 ? { attachments } : {}),
  };
}

export function WorkspaceSetupDialog() {
  const { t } = useTranslation();
  const toast = useToast();
  const pendingWorkspaceSetup = useWorkspaceSetupStore((state) => state.pendingWorkspaceSetup);
  const clearWorkspaceSetup = useWorkspaceSetupStore((state) => state.clearWorkspaceSetup);
  const mergeWorkspaces = useCallback(
    (targetServerId: string, workspaces: Iterable<WorkspaceDescriptor>) => {
      getHostRuntimeStore().acceptWorkspaceSnapshots(targetServerId, Array.from(workspaces));
    },
    [],
  );
  const setHasHydratedWorkspaces = useSessionStore((state) => state.setHasHydratedWorkspaces);
  const [errorMessage, setErrorMessage] = useState<string | null>(null);
  const [createdWorkspace, setCreatedWorkspace] = useState<ReturnType<
    typeof normalizeWorkspaceDescriptor
  > | null>(null);
  const [pendingAction, setPendingAction] = useState<"chat" | null>(null);

  const serverId = pendingWorkspaceSetup?.serverId ?? "";
  const supportsForgeSearch = useSessionStore(
    (state) => state.sessions[serverId]?.serverInfo?.features?.forgeSearch === true,
  );
  const sourceDirectory = pendingWorkspaceSetup?.sourceDirectory ?? "";
  const displayName = pendingWorkspaceSetup?.displayName?.trim() ?? "";
  const workspace = createdWorkspace;
  const client = useHostRuntimeClient(serverId);
  const isConnected = useHostRuntimeIsConnected(serverId);
  const chatDraft = useAgentInputDraft({
    draftKey: `workspace-setup:${serverId}:${sourceDirectory}`,
    composer: buildChatDraftComposerArgs({
      serverId,
      workspaceDirectory: workspace?.workspaceDirectory,
      sourceDirectory,
      pendingWorkspaceSetup,
    }),
  });
  const composerState = chatDraft.composerState;
  if (!composerState && pendingWorkspaceSetup) {
    throw new Error(t("workspaceSetup.errors.composerStateRequired"));
  }

  const { icon: projectIcon } = useProjectIcon({
    serverId,
    cwd: sourceDirectory,
  });
  const iconDataUri = toProjectIconDataUri(projectIcon);

  useEffect(() => {
    setErrorMessage(null);
    setCreatedWorkspace(null);
    setPendingAction(null);
  }, [pendingWorkspaceSetup?.creationId]);

  const handleClose = useCallback(() => {
    clearWorkspaceSetup();
  }, [clearWorkspaceSetup]);

  const navigateAfterCreation = useCallback(
    (
      workspaceId: string,
      target: { kind: "agent"; agentId: string } | { kind: "terminal"; terminalId: string },
    ) => {
      if (!pendingWorkspaceSetup) {
        return;
      }

      clearWorkspaceSetup();
      if (target.kind === "agent") {
        navigateToAgent({
          serverId: pendingWorkspaceSetup.serverId,
          agentId: target.agentId,
        });
        return;
      }

      navigateToWorkspace({
        serverId: pendingWorkspaceSetup.serverId,
        workspaceId,
        target,
      });
    },
    [clearWorkspaceSetup, pendingWorkspaceSetup],
  );

  const withConnectedClient = useCallback(() => {
    if (!client || !isConnected) {
      throw new Error(t("workspaceSetup.errors.hostDisconnected"));
    }
    return client;
  }, [client, isConnected, t]);

  const ensureWorkspace = useCallback(
    async (input: {
      cwd: string;
      attachments: MessagePayload["attachments"];
      agent?: CreateWorkspaceRequestOptions["agent"];
      onAgentCreated?: (agent: AgentSnapshotPayload) => void;
    }) => {
      if (!pendingWorkspaceSetup) {
        throw new Error(t("workspaceSetup.errors.pendingRequired"));
      }

      if (createdWorkspace && !input.agent) {
        return createdWorkspace;
      }

      const connectedClient = withConnectedClient();
      const payload = await callWorkspaceCreation({
        creationMethod: pendingWorkspaceSetup.creationMethod,
        creationId: pendingWorkspaceSetup.creationId,
        worktreeSlug: pendingWorkspaceSetup.worktreeSlug,
        connectedClient,
        input,
      });

      if (payload.error || !payload.workspace) {
        throw new Error(
          payload.error ?? failureMessageForCreationMethod(pendingWorkspaceSetup.creationMethod, t),
        );
      }

      if (payload.agent) input.onAgentCreated?.(payload.agent);
      const normalizedWorkspace = normalizeWorkspaceDescriptor(payload.workspace);
      mergeWorkspaces(pendingWorkspaceSetup.serverId, [normalizedWorkspace]);
      if (pendingWorkspaceSetup.creationMethod === "open_project") {
        setHasHydratedWorkspaces(pendingWorkspaceSetup.serverId, true);
      }
      setCreatedWorkspace(normalizedWorkspace);
      return normalizedWorkspace;
    },
    [
      createdWorkspace,
      mergeWorkspaces,
      pendingWorkspaceSetup,
      setHasHydratedWorkspaces,
      t,
      withConnectedClient,
    ],
  );

  const getIsStillActive = useCallback(() => {
    const current = useWorkspaceSetupStore.getState().pendingWorkspaceSetup;
    return (
      current?.serverId === pendingWorkspaceSetup?.serverId &&
      current?.sourceDirectory === pendingWorkspaceSetup?.sourceDirectory &&
      current?.creationMethod === pendingWorkspaceSetup?.creationMethod
    );
  }, [
    pendingWorkspaceSetup?.creationMethod,
    pendingWorkspaceSetup?.serverId,
    pendingWorkspaceSetup?.sourceDirectory,
  ]);

  const handleCreateChatAgent = useCallback(
    async ({ text, attachments, cwd }: MessagePayload) => {
      try {
        setPendingAction("chat");
        setErrorMessage(null);

        if (!composerState) {
          throw new Error(t("workspaceSetup.errors.composerStateRequired"));
        }
        if (!composerState.selectedProvider) {
          throw new Error(t("workspaceSetup.errors.selectModel"));
        }

        const wirePayload = splitComposerAttachmentsForSubmit(attachments, {
          format: resolveComposerAttachmentSubmitFormat({
            supportsForgeAttachments: supportsForgeSearch,
          }),
        });
        const encodedImages = await encodeImages(wirePayload.images);
        if (!pendingWorkspaceSetup) throw new Error(t("workspaceSetup.errors.pendingRequired"));
        let createdAgent: AgentSnapshotPayload | undefined;
        const { workspaceId: _workspaceId, ...agentInput } = buildCreateAgentOptions({
          composerState,
          text,
          attachments: wirePayload.attachments,
          encodedImages: encodedImages ?? null,
          workspaceDirectory: cwd,
          workspaceId: "",
          provider: composerState.selectedProvider,
        });
        const ensuredWorkspace = await ensureWorkspace({
          cwd,
          attachments,
          agent: {
            ...agentInput,
            clientMessageId: `${pendingWorkspaceSetup.creationId}:initial-message`,
          },
          onAgentCreated: (value) => {
            createdAgent = value;
          },
        });
        if (!createdAgent) throw new Error("The daemon did not create the requested agent");
        const agent = createdAgent;

        if (!getIsStillActive()) {
          return;
        }

        getHostRuntimeStore().acceptAgentSnapshot(
          serverId,
          applyLegacyDaemonWorkspaceOwnership({
            serverId,
            agent: normalizeAgentSnapshot(agent, serverId),
          }),
        );
        navigateAfterCreation(ensuredWorkspace.id, { kind: "agent", agentId: agent.id });
      } catch (error) {
        const message = toErrorMessage(error);
        setErrorMessage(message);
        toast.error(message);
      } finally {
        if (getIsStillActive()) {
          setPendingAction(null);
        }
      }
    },
    [
      composerState,
      pendingWorkspaceSetup,
      getIsStillActive,
      navigateAfterCreation,
      serverId,
      ensureWorkspace,
      t,
      toast,
      supportsForgeSearch,
    ],
  );

  const workspaceTitle = resolveWorkspaceTitle({ workspace, displayName, sourceDirectory });

  const placeholderLabel = projectIconPlaceholderLabelFromDisplayName(workspaceTitle);
  const placeholderInitial = placeholderLabel.charAt(0).toUpperCase();

  const agentControlsWithDisabled = useMemo(
    () =>
      composerState
        ? {
            ...composerState.agentControls,
            disabled: pendingAction !== null,
          }
        : undefined,
    [composerState, pendingAction],
  );

  const subtitleContent = useMemo(
    () => (
      <View style={styles.subtitleRow}>
        <ProjectIconView
          iconDataUri={iconDataUri}
          initial={placeholderInitial}
          projectViewKey={sourceDirectory}
          size={ICON_SIZE.md}
          textStyle={styles.projectIconFallbackText}
        />
        <Text style={styles.projectTitle} numberOfLines={1}>
          {workspaceTitle}
        </Text>
      </View>
    ),
    [iconDataUri, placeholderInitial, sourceDirectory, workspaceTitle],
  );

  const sheetHeader = useMemo<SheetHeader>(
    () => ({ title: t("workspaceSetup.title"), subtitle: subtitleContent }),
    [subtitleContent, t],
  );

  if (!pendingWorkspaceSetup || !sourceDirectory) {
    return null;
  }

  return (
    <AdaptiveModalSheet
      header={sheetHeader}
      visible={true}
      onClose={handleClose}
      snapPoints={SNAP_POINTS}
      testID="workspace-setup-dialog"
      desktopMaxWidth={640}
    >
      <FileDropZone style={styles.section}>
        <KeyboardTranslateView>
          <Composer
            agentId={`workspace-setup:${serverId}:${sourceDirectory}`}
            serverId={serverId}
            isPaneFocused={true}
            onSubmitMessage={handleCreateChatAgent}
            isSubmitLoading={pendingAction === "chat"}
            blurOnSubmit={true}
            textSource={chatDraft.textSource}
            onChangeText={chatDraft.editText}
            textReplacement={chatDraft.textReplacement}
            attachments={chatDraft.attachments}
            onChangeAttachments={chatDraft.setAttachments}
            cwd={sourceDirectory}
            clearDraft={chatDraft.clear}
            autoFocus
            commandDraft={composerState?.commandDraft}
            agentControls={agentControlsWithDisabled}
            inputWrapperStyle={styles.composerInputWrapper}
          />
        </KeyboardTranslateView>
      </FileDropZone>

      {errorMessage ? <Text style={styles.errorText}>{errorMessage}</Text> : null}
    </AdaptiveModalSheet>
  );
}

const styles = StyleSheet.create((theme) => ({
  subtitleRow: {
    flexDirection: "row",
    alignItems: "center",
    gap: theme.spacing[2],
  },
  projectIconFallbackText: {
    fontSize: 9,
  },
  projectTitle: {
    fontSize: theme.fontSize.base,
    color: theme.colors.foregroundMuted,
  },
  section: {
    gap: theme.spacing[3],
    marginHorizontal: -theme.spacing[6],
    marginVertical: -theme.spacing[2],
  },
  errorText: {
    fontSize: theme.fontSize.base,
    color: theme.colors.destructive,
    lineHeight: 20,
  },
  composerInputWrapper: {
    backgroundColor: theme.colors.surface2,
  },
}));