/* Copyright 2026 Marimo. All rights reserved. */

import { useAtom } from "jotai";
import {
  BotMessageSquareIcon,
  RefreshCwIcon,
  StopCircleIcon,
} from "lucide-react";
import React, { memo, useEffect, useMemo, useRef, useState } from "react";
import useEvent from "react-use-event-hook";
import { JsonRpcError, useAcpClient } from "use-acp";
import {
  ConnectionStatus,
  PermissionRequest,
} from "@/components/chat/acp/common";
import {
  type AdditionalCompletions,
  PromptInput,
} from "@/components/editor/ai/add-cell-with-ai";
import { PanelEmptyState } from "@/components/editor/chrome/panels/empty-state";
import { Spinner } from "@/components/icons/spinner";
import { Button } from "@/components/ui/button";
import { TooltipProvider } from "@/components/ui/tooltip";
import { cn } from "@/utils/cn";
import { Logger } from "@/utils/Logger";
import { capitalize } from "@/utils/strings";
import { AgentDocs } from "./agent-docs";
import { AgentSelector } from "./agent-selector";
import { ModelSelector } from "./model-selector";
import ScrollToBottomButton from "./scroll-to-bottom-button";
import { SessionTabs } from "./session-tabs";
import {
  agentSessionStateAtom,
  type ExternalAgentId,
  getAgentWebSocketUrl,
  selectedTabAtom,
  updateSessionExternalAgentSessionId,
  updateSessionTitle,
} from "./state";
import { AgentThread } from "./thread";
import "./agent-panel.css";
import type { Completion } from "@codemirror/autocomplete";
import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
import type {
  ContentBlock,
  RequestPermissionResponse,
} from "@zed-industries/agent-client-protocol";
import {
  addContextCompletion,
  CONTEXT_TRIGGER,
} from "@/components/editor/ai/completion-utils";
import {
  Select,
  SelectContent,
  SelectGroup,
  SelectItem,
  SelectLabel,
  SelectTrigger,
} from "@/components/ui/select";
import { toast } from "@/components/ui/use-toast";
import { DelayMount } from "@/components/utils/delay-mount";
import { useRequestClient } from "@/core/network/requests";
import { cwdAtom, filenameAtom } from "@/core/saving/file-state";
import { store } from "@/core/state/jotai";
import { ErrorBanner } from "@/plugins/impl/common/error-banner";
import { Functions } from "@/utils/functions";
import { PathBuilder, Paths } from "@/utils/paths";
import {
  AddContextButton,
  AttachFileButton,
  FileAttachmentPill,
  SendButton,
} from "../chat-components";
import { useFileState } from "../chat-utils";
import { ReadyToChatBlock } from "./blocks";
import {
  convertFilesToResourceLinks,
  parseContextFromPrompt,
} from "./context-utils";
import { getAgentPrompt } from "./prompt";
import type {
  AgentConnectionState,
  AgentPendingPermission,
  AvailableCommands,
  ExternalAgentSessionId,
  NotificationEvent,
  SessionMode,
  SessionModelState,
} from "./types";

const logger = Logger.get("agents");

interface AgentTitleProps {
  currentAgentId?: ExternalAgentId;
}

const AgentTitle = memo<AgentTitleProps>(({ currentAgentId }) => (
  <span className="text-sm font-medium">
    {capitalize(currentAgentId ?? "")}
  </span>
));
AgentTitle.displayName = "AgentTitle";

interface ConnectionControlProps {
  connectionState: AgentConnectionState;
  onConnect: () => void;
  onDisconnect: () => void;
}

const ConnectionControl = memo<ConnectionControlProps>(
  ({ connectionState, onConnect, onDisconnect }) => {
    const isConnected = connectionState.status === "connected";

    return (
      <Button
        variant="outline"
        size="xs"
        onClick={isConnected ? onDisconnect : onConnect}
        disabled={connectionState.status === "connecting"}
      >
        {isConnected ? "Disconnect" : "Connect"}
      </Button>
    );
  },
);
ConnectionControl.displayName = "ConnectionControl";

interface HeaderInfoProps {
  currentAgentId?: ExternalAgentId;
  connectionStatus: string;
  shouldShowConnectionControl?: boolean;
}

const HeaderInfo = memo<HeaderInfoProps>(
  ({ currentAgentId, connectionStatus, shouldShowConnectionControl }) => (
    <div className="flex items-center gap-2">
      <BotMessageSquareIcon className="h-4 w-4 text-muted-foreground" />
      {currentAgentId && <AgentTitle currentAgentId={currentAgentId} />}
      {shouldShowConnectionControl && (
        <ConnectionStatus status={connectionStatus} />
      )}
    </div>
  ),
);
HeaderInfo.displayName = "HeaderInfo";

interface AgentPanelHeaderProps {
  connectionState: AgentConnectionState;
  currentAgentId?: ExternalAgentId;
  onConnect: () => void;
  onDisconnect: () => void;
  onRestartThread?: () => void;
  hasActiveSession?: boolean;
  shouldShowConnectionControl?: boolean;
}

const AgentPanelHeader = memo<AgentPanelHeaderProps>(
  ({
    connectionState,
    currentAgentId,
    onConnect,
    onDisconnect,
    onRestartThread,
    hasActiveSession,
    shouldShowConnectionControl,
  }) => (
    <div className="flex border-b px-3 py-2 justify-between shrink-0 items-center">
      <HeaderInfo
        currentAgentId={currentAgentId}
        connectionStatus={connectionState.status}
        shouldShowConnectionControl={shouldShowConnectionControl}
      />
      <div className="flex items-center gap-2">
        {hasActiveSession &&
          connectionState.status === "connected" &&
          onRestartThread && (
            <Button
              variant="outline"
              size="xs"
              onClick={onRestartThread}
              title="Restart thread (create new session)"
            >
              <RefreshCwIcon className="h-3 w-3 mr-1" />
              Restart
            </Button>
          )}

        {shouldShowConnectionControl && (
          <ConnectionControl
            connectionState={connectionState}
            onConnect={onConnect}
            onDisconnect={onDisconnect}
          />
        )}
      </div>
    </div>
  ),
);
AgentPanelHeader.displayName = "AgentPanelHeader";

interface EmptyStateProps {
  currentAgentId?: ExternalAgentId;
  connectionState: AgentConnectionState;
  onConnect: () => void;
  onDisconnect: () => void;
}

const EmptyState = memo<EmptyStateProps>(
  ({ currentAgentId, connectionState, onConnect, onDisconnect }) => {
    return (
      <div className="flex flex-col h-full">
        <AgentPanelHeader
          connectionState={connectionState}
          currentAgentId={currentAgentId}
          onConnect={onConnect}
          onDisconnect={onDisconnect}
          hasActiveSession={false}
        />
        <SessionTabs />
        <div className="flex-1 flex items-center justify-center p-6">
          <div className="max-w-md w-full space-y-6">
            <PanelEmptyState
              title="No Agent Sessions"
              description="Create a new session to start a conversation"
              action={<AgentSelector className="border-y rounded" />}
              icon={<BotMessageSquareIcon />}
            />
            {connectionState.status === "disconnected" && (
              <AgentDocs
                className="border-t pt-6"
                title="Connect to an agent"
                description={
                  <>
                    <span>
                      Start agents by running these commands in your terminal.
                    </span>
                    <br />
                    <span>
                      Authenticate with the agent before starting a session.
                    </span>
                  </>
                }
              />
            )}
          </div>
        </div>
      </div>
    );
  },
);
EmptyState.displayName = "EmptyState";

interface LoadingIndicatorProps {
  isLoading: boolean;
  isRequestingPermission: boolean;
  onStop: () => void;
}

const LoadingIndicator = memo<LoadingIndicatorProps>(
  ({ isLoading, isRequestingPermission, onStop }) => {
    if (!isLoading) {
      return null;
    }

    return (
      <div className="px-3 py-2 border-t bg-muted/30 flex-shrink-0">
        <div className="flex items-center justify-between text-xs text-muted-foreground">
          <div className="flex items-center gap-2">
            <Spinner size="small" className="text-primary" />
            {isRequestingPermission ? (
              <span>Waiting for permission to continue...</span>
            ) : (
              <span>Agent is working...</span>
            )}
          </div>
          <div className="flex items-center gap-2">
            <Button
              size="sm"
              variant="outline"
              onClick={onStop}
              className="h-6 px-2"
            >
              <StopCircleIcon className="h-3 w-3 mr-1" />
              <span className="text-xs">Stop</span>
            </Button>
          </div>
        </div>
      </div>
    );
  },
);
LoadingIndicator.displayName = "LoadingIndicator";

interface PromptAreaProps {
  isLoading: boolean;
  activeSessionId: ExternalAgentSessionId | null;
  promptValue: string;
  commands: AvailableCommands | undefined;
  onPromptValueChange: (value: string) => void;
  onPromptSubmit: (e: KeyboardEvent | undefined, prompt: string) => void;
  onAddFiles: (files: File[]) => void;
  onStop: () => void;
  fileInputRef: React.RefObject<HTMLInputElement | null>;
  sessionMode?: SessionMode;
  onModeChange?: (mode: string) => void;
  sessionModels?: SessionModelState | null;
  onModelChange?: (modelId: string) => void;
}

const PromptArea = memo<PromptAreaProps>(
  ({
    isLoading,
    activeSessionId,
    promptValue,
    commands,
    onPromptValueChange,
    onPromptSubmit,
    onAddFiles,
    onStop,
    fileInputRef,
    sessionMode,
    onModeChange,
    sessionModels,
    onModelChange,
  }) => {
    const inputRef = useRef<ReactCodeMirrorRef | null>(null);
    const promptCompletions: AdditionalCompletions | undefined = useMemo(() => {
      if (!commands) {
        return undefined;
      }
      // sentence has to begin with '/' to trigger autocomplete
      return {
        triggerCompletionRegex: /^\/(\w+)?/,
        completions: commands.map(
          (prompt): Completion => ({
            label: `/${prompt.name}`,
            info: prompt.description,
          }),
        ),
      };
    }, [commands]);

    const handleSendClick = useEvent(() => {
      if (promptValue.trim()) {
        onPromptSubmit(undefined, promptValue);
      }
    });

    const handleAddContext = useEvent(() => {
      // For now, just append @ to the current value
      addContextCompletion(inputRef);
    });

    return (
      <div className="border-t bg-background flex-shrink-0">
        <div
          className={cn(
            "px-3 py-2 min-h-[80px]",
            (isLoading || !activeSessionId) && "opacity-50 pointer-events-none",
          )}
        >
          <PromptInput
            inputRef={inputRef}
            value={promptValue}
            onChange={isLoading ? Functions.NOOP : onPromptValueChange}
            onSubmit={onPromptSubmit}
            additionalCompletions={promptCompletions}
            onClose={Functions.NOOP}
            onAddFiles={onAddFiles}
            placeholder={
              isLoading
                ? "Processing..."
                : `Ask anything, ${CONTEXT_TRIGGER} to include context about tables or dataframes`
            }
            className={isLoading ? "opacity-50 pointer-events-none" : ""}
            maxHeight="120px"
          />
        </div>
        <TooltipProvider>
          <div className="px-3 py-2 border-t border-border/20 flex flex-row items-center justify-between">
            <div className="flex items-center gap-2">
              {sessionMode && onModeChange && (
                <ModeSelector
                  sessionMode={sessionMode}
                  onModeChange={onModeChange}
                />
              )}
              {sessionModels && onModelChange && activeSessionId && (
                <ModelSelector
                  sessionModels={sessionModels}
                  onModelChange={onModelChange}
                  disabled={isLoading}
                />
              )}
            </div>
            <div className="flex flex-row">
              <AddContextButton
                handleAddContext={handleAddContext}
                isLoading={isLoading}
              />
              <AttachFileButton
                fileInputRef={fileInputRef}
                isLoading={isLoading}
                onAddFiles={onAddFiles}
              />
              <SendButton
                isLoading={isLoading}
                onStop={onStop}
                onSendClick={handleSendClick}
                isEmpty={!promptValue.trim()}
              />
            </div>
          </div>
        </TooltipProvider>
      </div>
    );
  },
);
PromptArea.displayName = "PromptArea";

interface ModeSelectorProps {
  sessionMode: SessionMode;
  onModeChange: (mode: string) => void;
}

const ModeSelector = memo<ModeSelectorProps>(
  ({ sessionMode, onModeChange }) => {
    const availableModes = sessionMode?.availableModes || [];
    const currentModeId = sessionMode?.currentModeId;
    if (availableModes.length === 0) {
      return null;
    }

    const modeOptions = availableModes.map((mode) => ({
      value: mode.id,
      label: mode.name,
      subtitle: mode.description ?? "",
    }));
    const currentMode = modeOptions.find((opt) => opt.value === currentModeId);

    return (
      <Select value={currentModeId} onValueChange={onModeChange}>
        <SelectTrigger className="h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1 capitalize">
          {currentMode?.label ?? currentModeId}
        </SelectTrigger>
        <SelectContent>
          <SelectGroup>
            <SelectLabel>Agent Mode</SelectLabel>
            {modeOptions.map((option) => (
              <SelectItem
                key={option.value}
                value={option.value}
                className="text-xs"
              >
                <div className="flex flex-col">
                  {option.label}
                  {option.subtitle && (
                    <div className="text-muted-foreground text-xs pt-1 block">
                      {option.subtitle}
                    </div>
                  )}
                </div>
              </SelectItem>
            ))}
          </SelectGroup>
        </SelectContent>
      </Select>
    );
  },
);
ModeSelector.displayName = "ModeSelector";

interface ChatContentProps {
  hasNotifications: boolean;
  agentId: ExternalAgentId | undefined;
  connectionState: AgentConnectionState;
  sessionId: ExternalAgentSessionId | null;
  notifications: NotificationEvent[];
  pendingPermission: AgentPendingPermission;
  onResolvePermission: (option: RequestPermissionResponse) => void;
  onRetryConnection?: () => void;
  onRetryLastAction?: () => void;
  onDismissError?: (errorId: string) => void;
}

const ChatContent = memo<ChatContentProps>(
  ({
    hasNotifications,
    agentId,
    connectionState,
    notifications,
    pendingPermission,
    onResolvePermission,
    onRetryConnection,
    onRetryLastAction,
    onDismissError: _onDismissError,
    sessionId,
  }) => {
    const [isScrolledToBottom, setIsScrolledToBottom] = useState(true);
    const scrollContainerRef = useRef<HTMLDivElement>(null);
    const isDisconnected = connectionState.status === "disconnected";

    // Scroll handler to determine if we're at the bottom of the chat
    const handleScroll = useEvent(() => {
      const container = scrollContainerRef.current;
      if (!container) {
        return;
      }

      const { scrollTop, scrollHeight, clientHeight } = container;
      const hasOverflow = scrollHeight > clientHeight;
      const isAtBottom = hasOverflow
        ? Math.abs(scrollHeight - clientHeight - scrollTop) < 5
        : true; // 5px threshold
      setIsScrolledToBottom(isAtBottom);
    });

    const scrollToBottom = useEvent(() => {
      const container = scrollContainerRef.current;
      if (!container) {
        return;
      }

      container.scrollTo({
        top: container.scrollHeight,
        behavior: "smooth",
      });
    });

    // Auto-scroll to bottom when new notifications arrive (if already at bottom)
    useEffect(() => {
      if (isScrolledToBottom && notifications.length > 0) {
        // Use setTimeout to ensure DOM is updated before scrolling
        const timeout = setTimeout(scrollToBottom, 100);
        return () => clearTimeout(timeout);
      }
    }, [notifications.length, isScrolledToBottom, scrollToBottom]);

    const renderThread = () => {
      if (hasNotifications) {
        return (
          <AgentThread
            isConnected={connectionState.status === "connected"}
            notifications={notifications}
            onRetryConnection={onRetryConnection}
            onRetryLastAction={onRetryLastAction}
          />
        );
      }

      const isConnected = connectionState.status === "connected";
      if (isConnected) {
        return <ReadyToChatBlock />;
      }

      return (
        <div className="flex items-center justify-center h-full min-h-[200px] flex-col">
          <PanelEmptyState
            title="Waiting for agent"
            description="Your AI agent will appear here when active"
            icon={<BotMessageSquareIcon />}
          />
          {isDisconnected && agentId && (
            <AgentDocs
              className="border-t pt-6 px-5"
              title="Make sure you're connected to an agent"
              description="Run this command in your terminal:"
              agents={[agentId]}
            />
          )}
          {isDisconnected && (
            <Button
              variant="outline"
              onClick={onRetryConnection}
              type="button"
              className="mt-4"
            >
              Retry
            </Button>
          )}
        </div>
      );
    };

    return (
      <div className="flex-1 flex flex-col overflow-hidden flex-shrink-0 relative">
        {pendingPermission && (
          <div className="p-3 border-b">
            <PermissionRequest
              permission={pendingPermission}
              onResolve={onResolvePermission}
            />
          </div>
        )}

        <div
          ref={scrollContainerRef}
          className="flex-1 bg-muted/20 w-full flex flex-col overflow-y-auto p-2"
          onScroll={handleScroll}
        >
          {sessionId && (
            <div className="text-xs text-muted-foreground mb-2 px-2">
              Session ID: {sessionId}
            </div>
          )}
          {renderThread()}
        </div>

        <ScrollToBottomButton
          isVisible={!isScrolledToBottom && hasNotifications}
          onScrollToBottom={scrollToBottom}
        />
      </div>
    );
  },
);
ChatContent.displayName = "ChatContent";

const NO_WS_SET = "_skip_auto_connect_";
const AUTH_REQUIRED_CODE = -32_000;

function getDataMessage(data: unknown): string | undefined {
  if (data != null && typeof data === "object" && "message" in data) {
    const msg = (data as Record<string, unknown>).message;
    return typeof msg === "string" ? msg : undefined;
  }
  return undefined;
}

function getCwd(): string {
  const cwd = store.get(cwdAtom);
  if (cwd) {
    return cwd;
  }
  const filename = store.get(filenameAtom);
  if (!filename) {
    throw new Error(
      "Please save the notebook and refresh the browser to use the agent",
    );
  }
  return Paths.dirname(filename);
}

function getAbsoluteFilename(): string {
  const filename = store.get(filenameAtom);
  if (!filename) {
    throw new Error(
      "Please save the notebook and refresh the browser to use the agent",
    );
  }
  const cwd = store.get(cwdAtom);
  if (cwd) {
    const builder = PathBuilder.guessDeliminator(cwd);
    return builder.join(cwd, String(Paths.basename(filename)));
  }
  return filename;
}

const AgentPanel: React.FC = () => {
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<Error | string | null>(null);
  const [promptValue, setPromptValue] = useState("");
  const { files, addFiles, clearFiles, removeFile } = useFileState();
  const [sessionModels, setSessionModels] = useState<SessionModelState | null>(
    null,
  );
  const fileInputRef = useRef<HTMLInputElement>(null);

  const [selectedTab] = useAtom(selectedTabAtom);
  const [sessionState, setSessionState] = useAtom(agentSessionStateAtom);

  const wsUrl = selectedTab
    ? getAgentWebSocketUrl(selectedTab.agentId)
    : NO_WS_SET;
  const { sendUpdateFile, sendFileDetails } = useRequestClient();
  const creatingOrResumingSession = useRef(false);

  const acpClient = useAcpClient({
    wsUrl,
    clientOptions: {
      readTextFile: (request) => {
        logger.debug("Agent requesting file read", {
          path: request.path,
        });
        return sendFileDetails({ path: request.path }).then((response) => ({
          content: response.contents || "",
        }));
      },
      writeTextFile: (request) => {
        logger.debug("Agent requesting file write", {
          path: request.path,
          contentLength: request.content.length,
        });
        return sendUpdateFile({
          path: request.path,
          contents: request.content,
        }).then(() => ({}));
      },
    },
    autoConnect: false, // We'll manage connection manually based on active session
  });

  const {
    connect,
    disconnect,
    setActiveSessionId,
    connectionState,
    notifications,
    pendingPermission,
    availableCommands,
    resolvePermission,
    sessionMode,
    activeSessionId,
    agent,
    clearNotifications,
  } = acpClient;

  useEffect(() => {
    if (!agent) {
      return;
    }

    const initAndAuth = async () => {
      const response = await agent.initialize({
        protocolVersion: 1,
        clientCapabilities: {
          fs: {
            readTextFile: true,
            writeTextFile: true,
          },
        },
      });

      // We try to authenticate with the agent if it supports it.
      // The user must then restart the session
      const authMethods = response?.authMethods;
      if (authMethods && authMethods.length > 0) {
        await agent.authenticate({ methodId: authMethods[0].id });
      }
    };

    initAndAuth().catch((error) => {
      logger.error("Failed to initialize/authenticate agent", { error });
    });
  }, [agent]);

  // Auto-connect to agent when we have an active session, but only once per session
  useEffect(() => {
    setActiveSessionId(null);

    if (wsUrl === NO_WS_SET) {
      return;
    }

    logger.debug("Auto-connecting to agent", {
      sessionId: activeSessionId,
    });
    void connect().catch((error) => {
      logger.error("Failed to connect to agent", { error });
    });

    return () => {
      // We don't want to disconnect so users can switch between different
      // panels without losing their session
    };
    // oxlint-disable-next-line react-hooks/exhaustive-deps
  }, [wsUrl]);

  const handleNewSession = useEvent(async () => {
    if (!agent) {
      return;
    }

    creatingOrResumingSession.current = true;

    try {
      // If there is an active session, we should stop it
      if (activeSessionId) {
        await agent.cancel({ sessionId: activeSessionId }).catch((error) => {
          logger.error("Failed to cancel active session", { error });
        });
        clearNotifications(activeSessionId);
        setActiveSessionId(null);
      }

      // Get the selected model from the current session state
      const currentModel = selectedTab?.selectedModel ?? null;
      logger.debug("Creating new agent session", { model: currentModel });
      const newSession = await agent.newSession({
        cwd: getCwd(),
        mcpServers: [],
        _meta: currentModel ? { model: currentModel } : undefined,
      });

      // Capture models from the response
      if (newSession.models) {
        logger.debug("Session models received", { models: newSession.models });
        setSessionModels(newSession.models);
      }

      setSessionState((prev) =>
        updateSessionExternalAgentSessionId(
          prev,
          newSession.sessionId as ExternalAgentSessionId,
        ),
      );
    } finally {
      creatingOrResumingSession.current = false;
    }
  });

  const handleResumeSession = useEvent(
    async (previousSessionId: ExternalAgentSessionId) => {
      if (!agent) {
        return;
      }
      logger.debug("Resuming agent session", {
        sessionId: previousSessionId,
      });
      if (!agent.loadSession) {
        throw new Error("Agent does not support loading sessions");
      }
      creatingOrResumingSession.current = true;
      try {
        const loadedSession = await agent.loadSession({
          sessionId: previousSessionId,
          cwd: getCwd(),
          mcpServers: [],
        });

        // Capture models from the response if available
        if (loadedSession?.models) {
          logger.debug("Session models received", {
            models: loadedSession.models,
          });
          setSessionModels(loadedSession.models);
        }

        setSessionState((prev) =>
          updateSessionExternalAgentSessionId(prev, previousSessionId),
        );
      } finally {
        creatingOrResumingSession.current = false;
      }
    },
  );

  // Create or resume a session when successfully connected
  const isConnected = connectionState.status === "connected";
  const tabLastActiveSessionId = selectedTab?.externalAgentSessionId;
  useEffect(() => {
    // No need to do anything if we're not connected, don't have an agent, or don't have a selected tab
    if (!isConnected || !selectedTab || !agent) {
      return;
    }

    // Already have an active session
    if (activeSessionId && tabLastActiveSessionId) {
      return;
    }

    // Prevent race conditions
    if (creatingOrResumingSession.current) {
      return;
    }

    // If there is an available session, resume it, otherwise create a new one
    const createOrResumeSession = async () => {
      const availableSession = tabLastActiveSessionId ?? activeSessionId;
      try {
        if (availableSession) {
          try {
            await handleResumeSession(availableSession);
          } catch (error) {
            logger.error("Failed to resume session", {
              sessionId: availableSession,
              error,
            });
            await handleNewSession();
          }
        } else {
          await handleNewSession();
        }
        setError(null);
      } catch (error) {
        logger.error("Failed to create or resume session:", error);
        setError(error instanceof Error ? error : String(error));
      }
    };

    createOrResumeSession();
    // oxlint-disable-next-line react-hooks/exhaustive-deps
  }, [isConnected, agent, tabLastActiveSessionId, activeSessionId]);

  // Handler for prompt submission
  const handlePromptSubmit = useEvent(
    async (_e: KeyboardEvent | undefined, prompt: string) => {
      if (!activeSessionId || !agent || isLoading) {
        return;
      }

      logger.debug("Submitting prompt to agent", {
        sessionId: activeSessionId,
      });
      setIsLoading(true);
      setPromptValue("");
      clearFiles();

      // Update session title with first message if it's still the default
      if (selectedTab?.title.startsWith("New ")) {
        setSessionState((prev) => updateSessionTitle(prev, prompt));
      }

      let absoluteFilename: string;
      try {
        absoluteFilename = getAbsoluteFilename();
      } catch {
        toast({
          title: "Notebook must be named",
          description: "Please name the notebook to use the agent",
          variant: "danger",
        });
        return;
      }

      const promptBlocks: ContentBlock[] = [{ type: "text", text: prompt }];

      // Parse context from the prompt
      const { contextBlocks, attachmentBlocks } =
        await parseContextFromPrompt(prompt);
      promptBlocks.push(...contextBlocks, ...attachmentBlocks);

      // Add manually uploaded files as resource links
      if (files && files.length > 0) {
        const fileResourceLinks = await convertFilesToResourceLinks(files);
        promptBlocks.push(...fileResourceLinks);
      }

      const hasGivenRules = notifications.some(
        (notification) =>
          notification.type === "session_notification" &&
          notification.data.update.sessionUpdate === "user_message_chunk",
      );
      if (!hasGivenRules) {
        promptBlocks.push(
          {
            type: "resource_link",
            uri: absoluteFilename,
            mimeType: "text/x-python",
            name: absoluteFilename,
          },
          {
            type: "resource",
            resource: {
              uri: "marimo_rules.md",
              mimeType: "text/plain",
              text: getAgentPrompt(absoluteFilename),
            },
          },
        );
      }

      try {
        await agent.prompt({
          sessionId: activeSessionId,
          prompt: promptBlocks,
        });
      } catch (error) {
        logger.error("Failed to send prompt", { error });
      } finally {
        setIsLoading(false);
      }
    },
  );

  // Handler for stopping the current operation
  const handleStop = useEvent(async () => {
    if (!activeSessionId || !agent) {
      return;
    }
    await agent.cancel({ sessionId: activeSessionId });
    setIsLoading(false);
  });

  // Handler for manual connect
  const handleManualConnect = useEvent(() => {
    logger.debug("Manual connect requested", {
      currentStatus: connectionState.status,
    });
    connect();
  });

  // Handler for manual disconnect
  const handleManualDisconnect = useEvent(() => {
    logger.debug("Manual disconnect requested", {
      sessionId: activeSessionId,
      currentStatus: connectionState.status,
    });
    disconnect();
  });

  const handleModelChange = useEvent((modelId: string) => {
    logger.debug("Model change requested", {
      modelId,
      sessionId: activeSessionId,
    });

    if (!agent || !activeSessionId) {
      toast({
        title: "Cannot change model",
        description: "Please connect to an agent with an active session first",
        variant: "danger",
      });
      return;
    }

    // Call agent.setSessionModel to notify the agent
    void agent.setSessionModel?.({
      sessionId: activeSessionId,
      modelId,
    });

    // Update local state
    setSessionModels((prev) =>
      prev ? { ...prev, currentModelId: modelId } : null,
    );
  });

  const handleModeChange = useEvent((mode: string) => {
    logger.debug("Mode change requested", {
      sessionId: activeSessionId,
      mode,
    });
    if (!agent) {
      toast({
        title: "Agent not connected",
        description: "Please connect to an agent to change the mode",
        variant: "danger",
      });
      return;
    }
    if (!agent.setSessionMode) {
      toast({
        title: "Mode change not supported",
        description: "The agent does not support mode changes",
        variant: "danger",
      });
      return;
    }
    void agent.setSessionMode?.({
      sessionId: activeSessionId as string,
      modeId: mode,
    });
  });

  const hasNotifications = notifications.length > 0;
  const hasActiveSessions = sessionState.sessions.length > 0;

  if (!hasActiveSessions) {
    return (
      <EmptyState
        currentAgentId={selectedTab?.agentId}
        connectionState={connectionState}
        onConnect={handleManualConnect}
        onDisconnect={handleManualDisconnect}
      />
    );
  }

  const renderBody = () => {
    if (error) {
      const isAuthError =
        error instanceof JsonRpcError && error.code === AUTH_REQUIRED_CODE;
      const dataMessage =
        error instanceof JsonRpcError ? getDataMessage(error.data) : undefined;
      const displayError = dataMessage ? new Error(dataMessage) : error;

      return (
        <ErrorBanner
          className="w-3/4 mx-auto mt-10"
          error={displayError}
          action={
            isAuthError ? (
              <Button
                variant="linkDestructive"
                size="sm"
                onClick={() => {
                  setError(null);
                  handleNewSession();
                }}
              >
                Restart session
              </Button>
            ) : (
              <Button
                variant="linkDestructive"
                size="sm"
                onClick={() => setError(null)}
              >
                Dismiss
              </Button>
            )
          }
        />
      );
    }

    const isConnecting = connectionState.status === "connecting";
    const delay = 200; // ms
    if (isConnecting) {
      return (
        <DelayMount milliseconds={delay}>
          <div className="flex items-center justify-center h-full min-h-[200px] flex-col">
            <Spinner size="medium" className="text-primary" />
            <span className="text-sm text-muted-foreground">
              Connecting to the agent...
            </span>
          </div>
        </DelayMount>
      );
    }

    const isLoadingSession =
      tabLastActiveSessionId == null && connectionState.status === "connected";
    if (isLoadingSession) {
      return (
        <DelayMount milliseconds={delay}>
          <div className="flex items-center justify-center h-full min-h-[200px] flex-col">
            <Spinner size="medium" className="text-primary" />
            <span className="text-sm text-muted-foreground">
              Creating a new session...
            </span>
          </div>
        </DelayMount>
      );
    }

    return (
      <>
        <ChatContent
          key={activeSessionId}
          agentId={selectedTab?.agentId}
          sessionId={selectedTab?.externalAgentSessionId ?? null}
          hasNotifications={hasNotifications}
          connectionState={connectionState}
          notifications={notifications}
          pendingPermission={pendingPermission}
          onResolvePermission={(option) => {
            logger.debug("Resolving permission request", {
              sessionId: activeSessionId,
              option,
            });
            resolvePermission(option);
          }}
          onRetryConnection={handleManualConnect}
        />

        <LoadingIndicator
          isLoading={isLoading}
          isRequestingPermission={!!pendingPermission}
          onStop={handleStop}
        />

        {files && files.length > 0 && (
          <div className="flex flex-row gap-1 flex-wrap p-3 border-t">
            {files.map((file) => (
              <FileAttachmentPill
                file={file}
                key={file.name}
                onRemove={() => removeFile(file)}
              />
            ))}
          </div>
        )}

        <PromptArea
          isLoading={isLoading}
          activeSessionId={activeSessionId}
          promptValue={promptValue}
          onPromptValueChange={setPromptValue}
          onPromptSubmit={handlePromptSubmit}
          onAddFiles={addFiles}
          onStop={handleStop}
          fileInputRef={fileInputRef}
          commands={availableCommands}
          sessionMode={sessionMode}
          onModeChange={handleModeChange}
          sessionModels={sessionModels}
          onModelChange={handleModelChange}
        />
      </>
    );
  };

  return (
    <div className="flex flex-col flex-1 overflow-hidden mo-agent-panel">
      <AgentPanelHeader
        connectionState={connectionState}
        currentAgentId={selectedTab?.agentId}
        onConnect={handleManualConnect}
        onDisconnect={handleManualDisconnect}
        onRestartThread={handleNewSession}
        hasActiveSession={true}
        shouldShowConnectionControl={wsUrl !== NO_WS_SET}
      />
      <SessionTabs />

      {renderBody()}
    </div>
  );
};

export default AgentPanel;