kandev/apps/web/components/task/chat/chat-input-container.tsx

"use client";

import { forwardRef, useCallback, useState } from "react";
import { IconAlertTriangle, IconPlayerPlay, IconRefresh } from "@tabler/icons-react";
import { Button } from "@kandev/ui/button";
import { Tooltip, TooltipContent, TooltipTrigger } from "@kandev/ui/tooltip";
import { NewSessionDialog } from "@/components/task/new-session-dialog";
import { useAppStore } from "@/components/state-provider";
import type { ContextFile } from "@/lib/state/context-files-store";
import type { Message } from "@/lib/types/http";
import type { DiffComment } from "@/lib/diff/types";
import { getWebSocketClient } from "@/lib/ws/connection";
import { useChatInputContainer } from "./use-chat-input-container";
import {
  ChatInputBody,
  type ChatInputContextAreaProps,
  type ChatInputEditorAreaProps,
} from "./chat-input-body";
import type { ContextItem } from "@/lib/types/context";
import { useUtilityAgentGenerator } from "@/hooks/use-utility-agent-generator";
import { useIsUtilityConfigured } from "@/hooks/use-is-utility-configured";

// Re-export ImageAttachment type for consumers
export type { ImageAttachment } from "./image-attachment-preview";

// Type for message attachments sent to backend
export type MessageAttachment = {
  type: "image" | "audio" | "resource";
  data: string;
  mime_type: string;
  name?: string;
};

export type ChatInputContainerHandle = {
  focusInput: () => void;
  getTextareaElement: () => HTMLElement | null;
  getValue: () => string;
  getSelectionStart: () => number;
  insertText: (text: string, from: number, to: number) => void;
  clear: () => void;
  getAttachments: () => MessageAttachment[];
};

type ChatInputContainerProps = {
  onSubmit: (
    message: string,
    reviewComments?: DiffComment[],
    attachments?: MessageAttachment[],
    inlineMentions?: ContextFile[],
    inlineTaskMentions?: import("@/hooks/use-inline-mention").TaskMentionData[],
  ) => void;
  sessionId: string | null;
  taskId: string | null;
  taskTitle?: string;
  taskDescription: string;
  planModeEnabled: boolean;
  planModeAvailable?: boolean;
  mcpServers?: string[];
  onPlanModeChange: (enabled: boolean) => void;
  isAgentBusy: boolean;
  isStarting: boolean;
  /** True only while a containerized executor is bootstrapping (Docker
   * prepare, Sprites sandbox spin-up). Distinct from the brief STARTING
   * state every session — including local quick-chat — passes through;
   * see useSessionState. */
  isPreparingEnvironment?: boolean;
  isMoving?: boolean;
  isSending: boolean;
  onCancel: () => void;
  placeholder?: string;
  pendingClarification?: Message | null;
  onClarificationResolved?: () => void;
  showRequestChangesTooltip?: boolean;
  onRequestChangesTooltipDismiss?: () => void;
  pendingCommentsByFile?: Record<string, DiffComment[]>;
  hasContextComments?: boolean;
  submitKey?: "enter" | "cmd_enter";
  hasAgentCommands?: boolean;
  isFailed?: boolean;
  needsRecovery?: boolean;
  executorUnavailable?: boolean;
  executorUnavailableReason?: string;
  contextItems?: ContextItem[];
  planContextEnabled?: boolean;
  contextFiles?: ContextFile[];
  onToggleContextFile?: (file: ContextFile) => void;
  onAddContextFile?: (file: ContextFile) => void;
  onImplementPlan?: (fresh: boolean) => void;
  hideSessionsDropdown?: boolean;
  minimalToolbar?: boolean;
  /** Hide the plan mode toggle button (for ephemeral/quick chat sessions) */
  hidePlanMode?: boolean;
};

async function requestSessionRecover(
  taskId: string,
  sessionId: string,
  action: "resume" | "fresh_start",
): Promise<boolean> {
  const client = getWebSocketClient();
  if (!client) return false;
  try {
    await client.request(
      "session.recover",
      { task_id: taskId, session_id: sessionId, action },
      30000,
    );
    return true;
  } catch {
    return false;
  }
}

function FailedSessionBanner({
  showDialog,
  onShowDialog,
  taskId,
  sessionId,
  message = "This agent has stopped.",
  detail,
  resumeLabel = "Resume",
  resumingLabel = "Resuming...",
}: {
  showDialog: boolean;
  onShowDialog: (open: boolean) => void;
  taskId: string | null;
  sessionId: string | null;
  message?: string;
  detail?: string;
  resumeLabel?: string;
  resumingLabel?: string;
}) {
  const [isResuming, setIsResuming] = useState(false);
  const [isStartingFresh, setIsStartingFresh] = useState(false);

  const agentProfileId = useAppStore((s) =>
    sessionId ? (s.taskSessions.items[sessionId]?.agent_profile_id ?? "") : "",
  );
  const profileExists = useAppStore(
    (s) =>
      agentProfileId !== "" &&
      s.agentProfiles.items.some((p: { id: string }) => p.id === agentProfileId),
  );

  const handleRecover = useCallback(
    async (action: "resume" | "fresh_start") => {
      if (!sessionId || !taskId) return;
      const setBusy = action === "resume" ? setIsResuming : setIsStartingFresh;
      setBusy(true);
      const ok = await requestSessionRecover(taskId, sessionId, action);
      if (!ok) setBusy(false);
    },
    [sessionId, taskId],
  );

  const handleResume = useCallback(() => handleRecover("resume"), [handleRecover]);
  const handleFreshStart = useCallback(() => {
    if (!profileExists) {
      onShowDialog(true);
      return;
    }
    void handleRecover("fresh_start");
  }, [profileExists, onShowDialog, handleRecover]);

  return (
    <>
      <div className="rounded border border-border overflow-hidden">
        <div className="flex items-center justify-between gap-3 px-4 py-3">
          <div className="flex min-w-0 items-center gap-2 text-sm text-muted-foreground">
            <IconAlertTriangle className="h-4 w-4 text-orange-500 shrink-0" />
            <span className="truncate">{message}</span>
            {detail && <span className="shrink-0 text-xs text-muted-foreground">({detail})</span>}
          </div>
          <div className="flex items-center gap-2">
            {sessionId && taskId && (
              <Tooltip>
                <TooltipTrigger asChild>
                  <span className="inline-flex" data-testid="failed-session-resume-wrapper">
                    <Button
                      variant="default"
                      size="sm"
                      data-testid="recovery-resume-button"
                      className="shrink-0 gap-1.5 cursor-pointer"
                      onClick={handleResume}
                      disabled={isResuming || !profileExists}
                    >
                      <IconPlayerPlay className="h-3.5 w-3.5" />
                      {isResuming ? resumingLabel : resumeLabel}
                    </Button>
                  </span>
                </TooltipTrigger>
                {!profileExists && <TooltipContent>Agent profile no longer exists</TooltipContent>}
              </Tooltip>
            )}
            <Button
              variant="outline"
              size="sm"
              className="shrink-0 gap-1.5 cursor-pointer"
              onClick={handleFreshStart}
              disabled={isStartingFresh}
              data-testid="recovery-fresh-button"
            >
              <IconRefresh className="h-3.5 w-3.5" />
              {isStartingFresh ? "Starting..." : "Start fresh session"}
            </Button>
          </div>
        </div>
      </div>
      {taskId && <NewSessionDialog open={showDialog} onOpenChange={onShowDialog} taskId={taskId} />}
    </>
  );
}

type ContainerState = ReturnType<typeof useChatInputContainer>;
type NormalizedChatInputProps = ChatInputContainerProps & {
  isFailed: boolean;
  hasAgentCommands: boolean;
  submitKey: "enter" | "cmd_enter";
  planContextEnabled: boolean;
  contextFiles: ContextFile[];
  contextItems: ContextItem[];
  showRequestChangesTooltip: boolean;
};

function normalizeChatInputProps(p: ChatInputContainerProps): NormalizedChatInputProps {
  return {
    ...p,
    isFailed: p.isFailed ?? false,
    hasAgentCommands: p.hasAgentCommands ?? false,
    submitKey: p.submitKey ?? "cmd_enter",
    planContextEnabled: p.planContextEnabled ?? false,
    contextFiles: p.contextFiles ?? [],
    contextItems: p.contextItems ?? [],
    showRequestChangesTooltip: p.showRequestChangesTooltip ?? false,
  };
}

function buildContextAreaProps(
  s: ContainerState,
  p: ChatInputContainerProps,
): ChatInputContextAreaProps {
  return {
    hasContextZone: s.hasContextZone,
    allItems: s.allItems,
    sessionId: p.sessionId,
  };
}

type EnhancePromptExtras = {
  onEnhancePrompt?: () => void;
  isEnhancingPrompt?: boolean;
  isUtilityConfigured?: boolean;
  onVoiceTranscript?: (text: string) => void;
  onVoiceAutoSend?: () => void;
};

function buildEditorAreaProps(
  s: ContainerState,
  p: NormalizedChatInputProps,
  extras: EnhancePromptExtras = {},
): ChatInputEditorAreaProps {
  return {
    inputRef: s.inputRef,
    value: s.value,
    handleChange: s.handleChange,
    handleSubmitWithReset: s.handleSubmitWithReset,
    inputPlaceholder: s.inputPlaceholder,
    isDisabled: s.isDisabled,
    submitDisabled: s.submitDisabled,
    submitDisabledReason: s.submitDisabledReason,
    hasClarification: s.hasClarification,
    planModeEnabled: p.planModeEnabled,
    planModeAvailable: p.planModeAvailable ?? true,
    mcpServers: p.mcpServers ?? [],
    submitKey: p.submitKey,
    setIsInputFocused: s.setIsInputFocused,
    sessionId: p.sessionId,
    taskId: p.taskId,
    onAddContextFile: p.onAddContextFile,
    onToggleContextFile: p.onToggleContextFile,
    planContextEnabled: p.planContextEnabled,
    handleAgentCommand: s.handleAgentCommand,
    addFiles: s.addFiles,
    fileInputRef: s.fileInputRef,
    showRequestChangesTooltip: p.showRequestChangesTooltip,
    isAgentBusy: p.isAgentBusy,
    onPlanModeChange: p.onPlanModeChange,
    taskTitle: p.taskTitle,
    taskDescription: p.taskDescription,
    isSending: p.isSending,
    onCancel: p.onCancel,
    contextCount: s.allItems.length,
    contextPopoverOpen: s.contextPopoverOpen,
    setContextPopoverOpen: s.setContextPopoverOpen,
    contextFiles: p.contextFiles,
    onImplementPlan: p.onImplementPlan,
    onEnhancePrompt: extras.onEnhancePrompt,
    isEnhancingPrompt: extras.isEnhancingPrompt,
    isUtilityConfigured: extras.isUtilityConfigured,
    onVoiceTranscript: extras.onVoiceTranscript,
    onVoiceAutoSend: extras.onVoiceAutoSend,
    hideSessionsDropdown: p.hideSessionsDropdown,
    minimalToolbar: p.minimalToolbar,
    hidePlanMode: p.hidePlanMode,
  };
}

function buildStoppedBannerProps(p: ChatInputContainerProps) {
  if (!p.executorUnavailable) return {};
  return {
    message: "Executor environment is unavailable.",
    detail: p.executorUnavailableReason,
    resumeLabel: "Restart",
    resumingLabel: "Restarting...",
  };
}

export const ChatInputContainer = forwardRef<ChatInputContainerHandle, ChatInputContainerProps>(
  function ChatInputContainer(props, ref) {
    const { sessionId, taskId, taskTitle, taskDescription, isAgentBusy, isStarting, isSending } =
      props;
    const p = normalizeChatInputProps(props);
    const isMoving = props.isMoving ?? false;
    const executorUnavailable = props.executorUnavailable ?? false;
    const isBusyVisual = isStarting || isMoving;

    const s = useChatInputContainer({
      ref,
      sessionId,
      isSending,
      isStarting,
      isPreparingEnvironment: props.isPreparingEnvironment ?? false,
      isMoving,
      isFailed: p.isFailed,
      needsRecovery: props.needsRecovery ?? false,
      executorUnavailable,
      isAgentBusy,
      hasAgentCommands: p.hasAgentCommands,
      placeholder: props.placeholder,
      contextItems: p.contextItems,
      pendingClarification: props.pendingClarification,
      onClarificationResolved: props.onClarificationResolved,
      pendingCommentsByFile: props.pendingCommentsByFile,
      hasContextComments: props.hasContextComments ?? false,
      showRequestChangesTooltip: p.showRequestChangesTooltip,
      onRequestChangesTooltipDismiss: props.onRequestChangesTooltipDismiss,
      onSubmit: props.onSubmit,
    });

    const isUtilityConfigured = useIsUtilityConfigured();
    const { enhancePrompt, isEnhancingPrompt } = useUtilityAgentGenerator({
      sessionId,
      taskTitle,
      taskDescription,
    });

    const handleEnhancePrompt = useCallback(() => {
      const currentValue = s.value?.trim();
      if (!currentValue) return;
      enhancePrompt(currentValue, (enhanced) => {
        // Use setValue to directly update TipTap editor (handleChange only updates React state)
        s.inputRef.current?.setValue(enhanced);
      });
    }, [s, enhancePrompt]);

    const handleVoiceTranscript = useCallback(
      (text: string) => {
        const editor = s.inputRef.current;
        if (!editor) return;
        const trimmed = text.trim();
        if (!trimmed) return;
        const cursor = editor.getSelectionStart();
        const current = editor.getValue();
        // Prepend a space when inserting after existing non-whitespace content
        // so transcripts flow naturally without running into the previous word.
        const charBefore = cursor > 0 ? current.charAt(cursor - 1) : "";
        const needsLeadingSpace = charBefore !== "" && !/\s/.test(charBefore);
        const insert = needsLeadingSpace ? ` ${trimmed}` : trimmed;
        editor.insertText(insert, cursor, cursor);
      },
      [s.inputRef],
    );

    // Auto-send fires the same submit path as the regular send button. Guards
    // against firing while the input is in a disabled state (e.g. the agent
    // is currently booting) — the button is hidden in that case anyway, but
    // defence-in-depth so a stale keyboard shortcut press doesn't trigger.
    const { submitDisabled: voiceSubmitDisabled, handleSubmitWithReset: voiceSubmit } = s;
    const handleVoiceAutoSend = useCallback(() => {
      if (voiceSubmitDisabled) return;
      voiceSubmit();
    }, [voiceSubmitDisabled, voiceSubmit]);

    if (p.isFailed || executorUnavailable) {
      return (
        <FailedSessionBanner
          showDialog={s.showNewSessionDialog}
          onShowDialog={s.setShowNewSessionDialog}
          taskId={taskId}
          sessionId={sessionId}
          {...buildStoppedBannerProps(props)}
        />
      );
    }

    return (
      <ChatInputBody
        containerRef={s.containerRef}
        height={s.height}
        resizeHandleProps={s.resizeHandleProps}
        isStarting={isBusyVisual}
        isAgentBusy={isAgentBusy}
        hasClarification={s.hasClarification}
        showRequestChangesTooltip={p.showRequestChangesTooltip}
        hasPendingComments={s.hasPendingComments}
        planModeEnabled={props.planModeEnabled}
        showFocusHint={s.showFocusHint}
        needsRecovery={(props.needsRecovery ?? false) || executorUnavailable}
        addFiles={s.addFiles}
        contextAreaProps={buildContextAreaProps(s, p)}
        editorAreaProps={buildEditorAreaProps(s, p, {
          onEnhancePrompt: handleEnhancePrompt,
          isEnhancingPrompt,
          isUtilityConfigured,
          onVoiceTranscript: handleVoiceTranscript,
          onVoiceAutoSend: handleVoiceAutoSend,
        })}
      />
    );
  },
);