ChatMessageList.tsx

import React, {useCallback, useEffect, useMemo, useRef, useState} from "react"

import type {SimpleChatMessage} from "@agenta/shared/types"
import {
    extractTextFromContent,
    extractDisplayTextFromMessage,
    updateTextInContent,
    addImageToContent,
    addFileToContent,
    removeAttachmentFromContent,
    getAttachments,
} from "@agenta/shared/utils"
import {Copy, MinusCircle, Plus} from "@phosphor-icons/react"
import {Button, Tooltip} from "antd"
import {useAtom} from "jotai"

import {CollapseToggleButton} from "../../components/presentational/buttons"
import {ViewModeDropdown} from "../../drill-in/core/ViewModeDropdown"
import {messageViewModeAtom} from "../../drill-in/state/messageViewModeAtom"
import {getViewOptions, toMessageViewMode, type ViewMode} from "../../drill-in/utils/getViewOptions"
import {message, modal} from "../../utils/appMessageContext"
import {cn, flexLayouts, gapClasses} from "../../utils/styles"
import {createSnippetPdfAttachment} from "../utils/snippetAttachment"

import AttachmentButton from "./AttachmentButton"
import ChatMessageEditor from "./ChatMessageEditor"
import MessageAttachments from "./MessageAttachments"
import ToolMessageHeader from "./ToolMessageHeader"

type ChatViewMode = Extract<ViewMode, "text" | "markdown" | "json" | "yaml">

const ChatMessageItem: React.FC<{
    msg: SimpleChatMessage
    index: number
    editorId: string
    disabled?: boolean
    messageClassName?: string
    placeholder: string
    isMinimized: boolean
    showControls: boolean
    showRemoveButton?: boolean
    showCopyButton: boolean
    allowFileUpload: boolean
    enableTokens: boolean
    templateFormat?: "mustache" | "curly" | "fstring" | "jinja2"
    tokens?: string[]
    loadingFallback: "skeleton" | "none" | "static"
    maxPasteChars?: number
    /** Restrict the view-mode dropdown to a subset (e.g. ["text", "markdown"]
     *  for config messages where JSON/YAML modes are noise). When omitted,
     *  the dropdown shows whatever getViewOptions returns for the content. */
    viewModes?: ChatViewMode[]
    ImagePreview?: React.ComponentType<{
        src: string
        alt: string
        size: number
        isValidPreview: boolean
    }>
    onRoleChange: (index: number, role: string) => void
    onTextChange: (index: number, text: string) => void
    onRemove: (index: number) => void
    onAddImage: (index: number, url: string) => void
    onAddFile: (index: number, data: string, name: string, format: string) => void
    onRemoveAttachment: (msgIndex: number, attachmentIndex: number) => void
    onToggleMinimize: (index: number) => void
}> = ({
    msg,
    index,
    editorId,
    disabled,
    messageClassName,
    placeholder,
    isMinimized,
    showControls,
    showRemoveButton,
    showCopyButton,
    allowFileUpload,
    enableTokens,
    templateFormat,
    tokens,
    loadingFallback,
    maxPasteChars,
    viewModes,
    ImagePreview,
    onRoleChange,
    onTextChange,
    onRemove,
    onAddImage,
    onAddFile,
    onRemoveAttachment,
    onToggleMinimize,
}) => {
    const containerRef = useRef<HTMLDivElement>(null)
    // Shared + persisted across all message editors (see messageViewModeAtom).
    // The atom is typed `ViewMode` (can hold "form"), so coerce to a mode this
    // editor can actually render before deriving any mode-dependent state.
    const [viewMode, setViewMode] = useAtom(messageViewModeAtom)
    const chatViewMode = toMessageViewMode(viewMode)
    const isCodeMode = chatViewMode === "json" || chatViewMode === "yaml"
    const editorLanguage: "json" | "yaml" = chatViewMode === "yaml" ? "yaml" : "json"

    const isToolResponse = msg.role === "tool"
    const hasToolCalls = Boolean(msg.tool_calls && msg.tool_calls.length > 0)
    const textContent = hasToolCalls
        ? extractDisplayTextFromMessage(msg)
        : extractTextFromContent(msg.content ?? null)
    const attachments = getAttachments(msg.content ?? null)
    const hasAttachmentsFlag = attachments.length > 0

    const viewOptions = useMemo(() => {
        const all = getViewOptions(textContent) as {value: ChatViewMode; label: string}[]
        if (!viewModes || viewModes.length === 0) return all
        const allowed = new Set(viewModes)
        return all.filter((opt) => allowed.has(opt.value))
    }, [textContent, viewModes])

    const handleCreateSnippetFromPaste = useCallback(
        ({
            pastedText,
            maxPasteChars,
            overBy,
        }: {
            pastedText: string
            maxPasteChars: number
            overBy: number
        }) => {
            if (!allowFileUpload || !modal) {
                return false
            }

            const limitSummary =
                overBy > 0
                    ? `This paste is ${overBy.toLocaleString()} characters over the ${maxPasteChars.toLocaleString()}-character limit.`
                    : `This paste exceeds the ${maxPasteChars.toLocaleString()}-character limit.`

            modal.confirm({
                title: "That's too long to paste",
                content: `${limitSummary} To keep the editor responsive, you can attach the pasted content as a snippet instead.`,
                okText: "Create Snippet",
                cancelText: "Dismiss",
                centered: true,
                onOk: async () => {
                    try {
                        const {fileData, filename, mimeType} =
                            await createSnippetPdfAttachment(pastedText)
                        onAddFile(index, fileData, filename, mimeType)
                        message?.success(`Attached ${filename} as a snippet.`)
                    } catch (error) {
                        message?.error(
                            error instanceof Error
                                ? error.message
                                : "Failed to create snippet attachment.",
                        )
                        throw error
                    }
                },
            })

            return true
        },
        [allowFileUpload, index, onAddFile],
    )

    return (
        <div
            className={cn(
                flexLayouts.column,
                // Collapsed = role row + one line of content, clipped vertically so
                // formatting (indentation, markdown, JSON) is preserved instead of
                // collapsed to a single run-on line. Height is a clean line-height
                // multiple so it never bleeds a half-line.
                isMinimized &&
                    "[&_.agenta-editor-wrapper]:!max-h-[1lh] [&_.agenta-editor-wrapper]:overflow-hidden",
            )}
            ref={containerRef}
        >
            <ChatMessageEditor
                id={editorId}
                key={`${editorId}-${viewMode}`}
                role={msg.role}
                text={textContent}
                disabled={disabled}
                className={cn(messageClassName)}
                placeholder={placeholder}
                onChangeRole={(role) => onRoleChange(index, role)}
                onChangeText={(text) => onTextChange(index, text)}
                isJSON={isCodeMode}
                language={editorLanguage}
                markdownView={chatViewMode === "text"}
                enableTokens={enableTokens && !isCodeMode}
                templateFormat={templateFormat}
                tokens={tokens}
                loadingFallback={loadingFallback}
                maxPasteChars={maxPasteChars}
                onPasteLimitExceeded={({pastedText, maxPasteChars, overBy}) =>
                    handleCreateSnippetFromPaste({pastedText, maxPasteChars, overBy})
                }
                headerBottom={
                    isToolResponse && (msg.name || msg.tool_call_id) ? (
                        <ToolMessageHeader name={msg.name} toolCallId={msg.tool_call_id} />
                    ) : undefined
                }
                headerRight={
                    <div className={cn(flexLayouts.rowCenter, gapClasses.xs)}>
                        <div
                            className={cn(
                                flexLayouts.rowCenter,
                                gapClasses.xs,
                                "invisible group-hover/item:visible",
                            )}
                        >
                            <ViewModeDropdown<ChatViewMode>
                                value={chatViewMode}
                                options={viewOptions}
                                onChange={setViewMode}
                            />
                            {allowFileUpload && !disabled && (
                                <AttachmentButton
                                    onAddImage={(url) => onAddImage(index, url)}
                                    onAddFile={(data, name, format) =>
                                        onAddFile(index, data, name, format)
                                    }
                                    disabled={disabled}
                                />
                            )}
                            {showCopyButton && (
                                <Tooltip title="Copy">
                                    <Button
                                        type="text"
                                        size="small"
                                        icon={<Copy size={14} />}
                                        onClick={() => {
                                            navigator.clipboard.writeText(textContent)
                                        }}
                                    />
                                </Tooltip>
                            )}
                            {(showRemoveButton ?? showControls) && !disabled && (
                                <Tooltip title="Remove">
                                    <Button
                                        type="text"
                                        size="small"
                                        icon={<MinusCircle size={14} />}
                                        onClick={() => onRemove(index)}
                                    />
                                </Tooltip>
                            )}
                        </div>
                        <CollapseToggleButton
                            collapsed={isMinimized}
                            onToggle={() => onToggleMinimize(index)}
                            className="!transition-opacity !duration-0 !delay-200 group-hover/item:!delay-0 opacity-50 group-hover/item:opacity-100"
                        />
                    </div>
                }
                footer={
                    hasAttachmentsFlag ? (
                        <MessageAttachments
                            content={msg.content!}
                            onRemove={(attachmentIndex) =>
                                onRemoveAttachment(index, attachmentIndex)
                            }
                            disabled={disabled}
                            ImagePreview={ImagePreview}
                        />
                    ) : undefined
                }
            />
        </div>
    )
}

export interface ChatMessageListProps {
    /** Array of chat messages to display */
    messages: SimpleChatMessage[]
    /** Callback when messages change */
    onChange: (messages: SimpleChatMessage[]) => void
    /** Whether the list is disabled */
    disabled?: boolean
    /** Additional class name for the container */
    className?: string
    /** Additional class name for each message editor */
    messageClassName?: string
    /** Placeholder text for empty messages */
    placeholder?: string
    /** Whether to show add/remove controls (add message button + per-message remove) */
    showControls?: boolean
    /** Whether to show per-message remove button (independent of showControls) */
    showRemoveButton?: boolean
    /** Whether to show per-message copy button */
    showCopyButton?: boolean
    /** Whether to allow file uploads */
    allowFileUpload?: boolean
    /** Whether to enable variable token highlighting */
    enableTokens?: boolean
    /** Template format for variable syntax highlighting */
    templateFormat?: "mustache" | "curly" | "fstring" | "jinja2"
    /** Available template variables for token highlighting */
    tokens?: string[]
    /** Optional image preview component */
    ImagePreview?: React.ComponentType<{
        src: string
        alt: string
        size: number
        isValidPreview: boolean
    }>
    /** Whether messages should start minimized */
    defaultMinimized?: boolean
    /** Suspense fallback mode for editor plugins */
    loadingFallback?: "skeleton" | "none" | "static"
    /** Block paste operations that would make a message exceed this many characters. */
    maxPasteChars?: number
    /** Restrict the per-message view-mode dropdown to a subset. Pass
     *  ["text", "markdown"] for plain-text config messages where JSON/YAML
     *  modes are noise. When omitted, all four modes are offered. */
    viewModes?: ChatViewMode[]
}

/**
 * A list of chat message editors for editing multiple messages.
 * This is a simpler alternative to ChatInputs that uses the same visual style
 * as the Playground message editors.
 */
let _keyCounter = 0
function generateKey(): string {
    return `__id-${++_keyCounter}-${Date.now()}`
}

export const ChatMessageList: React.FC<ChatMessageListProps> = ({
    messages,
    onChange,
    disabled,
    className,
    messageClassName,
    placeholder = "Enter message...",
    showControls = true,
    showRemoveButton,
    showCopyButton = false,
    allowFileUpload = true,
    enableTokens = false,
    templateFormat,
    tokens,
    ImagePreview,
    defaultMinimized = false,
    loadingFallback = "skeleton",
    maxPasteChars,
    viewModes,
}) => {
    const listInstanceIdRef = useRef(generateKey())
    // Maintain stable React keys for each message position.
    // This prevents React from reusing the wrong component instance
    // when messages are added or removed from the middle of the list.
    const stableKeysRef = useRef<string[]>([])

    const stableKeys = useMemo(() => {
        const prev = stableKeysRef.current
        const next: string[] = []

        for (let i = 0; i < messages.length; i++) {
            // Reuse existing key if we have one at this position, otherwise generate new
            if (i < prev.length) {
                next.push(prev[i])
            } else {
                next.push(messages[i].id || generateKey())
            }
        }

        stableKeysRef.current = next
        return next
    }, [messages])

    const [minimizedMessages, setMinimizedMessages] = useState<Record<string, boolean>>(() =>
        defaultMinimized ? Object.fromEntries(stableKeys.map((key) => [key, true])) : {},
    )

    useEffect(() => {
        if (!defaultMinimized) return

        setMinimizedMessages((prev) => {
            const next: Record<string, boolean> = {}

            for (const key of stableKeys) {
                next[key] = prev[key] ?? true
            }

            return next
        })
    }, [defaultMinimized, stableKeys])

    const handleRoleChange = useCallback(
        (index: number, role: string) => {
            const updated = [...messages]
            updated[index] = {...updated[index], role}
            onChange(updated)
        },
        [messages, onChange],
    )

    const handleTextChange = useCallback(
        (index: number, newText: string) => {
            const updated = [...messages]
            const currentContent = updated[index].content ?? ""
            updated[index] = {
                ...updated[index],
                content: updateTextInContent(currentContent, newText),
            }
            onChange(updated)
        },
        [messages, onChange],
    )

    const handleAddMessage = useCallback(() => {
        onChange([...messages, {role: "user", content: ""}])
    }, [messages, onChange])

    const handleRemoveMessage = useCallback(
        (index: number) => {
            // Remove the stable key at the deleted index so remaining messages
            // keep their original keys and React preserves the correct component instances
            stableKeysRef.current = stableKeysRef.current.filter((_, i) => i !== index)
            const updated = messages.filter((_, i) => i !== index)
            onChange(updated)
        },
        [messages, onChange],
    )

    const handleAddImage = useCallback(
        (index: number, imageUrl: string) => {
            const updated = [...messages]
            updated[index] = {
                ...updated[index],
                content: addImageToContent(updated[index].content ?? "", imageUrl),
            }
            onChange(updated)
        },
        [messages, onChange],
    )

    const handleAddFile = useCallback(
        (index: number, fileData: string, filename: string, format: string) => {
            const updated = [...messages]
            updated[index] = {
                ...updated[index],
                content: addFileToContent(updated[index].content ?? "", fileData, filename, format),
            }
            onChange(updated)
        },
        [messages, onChange],
    )

    const handleRemoveAttachment = useCallback(
        (msgIndex: number, attachmentIndex: number) => {
            const updated = [...messages]
            updated[msgIndex] = {
                ...updated[msgIndex],
                content: removeAttachmentFromContent(
                    updated[msgIndex].content ?? "",
                    attachmentIndex,
                ),
            }
            onChange(updated)
        },
        [messages, onChange],
    )

    return (
        <div className={cn(flexLayouts.column, gapClasses.sm, className)}>
            {messages.map((msg, index) => {
                const rowKey = stableKeys[index]
                // Scope editor ids to the list instance so markdown-view state,
                // Lexical namespaces, and other editor-local caches never bleed
                // across separate prompt/message lists that share the same row index.
                const editorId = `chat-msg-${listInstanceIdRef.current}-${rowKey}`

                return (
                    <ChatMessageItem
                        key={rowKey}
                        msg={msg}
                        index={index}
                        editorId={editorId}
                        disabled={disabled}
                        messageClassName={messageClassName}
                        placeholder={placeholder}
                        isMinimized={minimizedMessages[rowKey] ?? false}
                        showControls={showControls}
                        showRemoveButton={showRemoveButton}
                        showCopyButton={showCopyButton}
                        allowFileUpload={allowFileUpload}
                        enableTokens={enableTokens}
                        templateFormat={templateFormat}
                        tokens={tokens}
                        loadingFallback={loadingFallback}
                        maxPasteChars={maxPasteChars}
                        viewModes={viewModes}
                        ImagePreview={ImagePreview}
                        onRoleChange={handleRoleChange}
                        onTextChange={handleTextChange}
                        onRemove={handleRemoveMessage}
                        onAddImage={handleAddImage}
                        onAddFile={handleAddFile}
                        onRemoveAttachment={handleRemoveAttachment}
                        onToggleMinimize={(i) => {
                            const key = stableKeys[i]
                            setMinimizedMessages((prev) => ({...prev, [key]: !prev[key]}))
                        }}
                    />
                )
            })}
            {showControls && !disabled && (
                <Button
                    variant="outlined"
                    color="default"
                    size="small"
                    icon={<Plus size={14} />}
                    onClick={handleAddMessage}
                    className="self-start"
                >
                    Message
                </Button>
            )}
        </div>
    )
}

export default ChatMessageList