frontend/components/ai_assistant.py

"""
AI Assistant Components for Rubric AI
Contains: ai_assistant_pane, suggestion_card, chat_input
"""

from contextlib import contextmanager
from nicegui import ui
from .layout import material_icon
from frontend.utils.formatting import format_relative_timestamp
from frontend.config.stage_config import get_stage_config, get_stage_display_name


@contextmanager
def ai_assistant_pane():
    """
    Create the AI assistant pane (right side of the two-column layout).

    This pane contains:
    - Header with AI icon and title
    - Content area for suggestions/responses
    - Chat input at the bottom

    Usage:
        with ai_assistant_pane() as assistant:
            suggestion_card(
                message="I've drafted 3 groups for you.",
                suggestions=['Thesis', 'Evidence', 'Clarity'],
                on_add_all=handle_add_all
            )
    """
    with ui.element('div').classes('pane p-6'):
        # Header
        with ui.element('div').classes('ai-assistant-header'):
            with ui.element('div').classes('ai-icon-wrapper'):
                material_icon('auto_awesome')
            ui.html('<h3 class="ai-assistant-title">AI Assistant</h3>', sanitize=False)

        # Content area (will be filled by yielded content)
        with ui.element('div').classes('flex-grow flex flex-col') as content:
            yield content


class SuggestionCard:
    """
    A card showing AI suggestions with action buttons.

    Displays:
    - Message from the AI
    - Clickable suggestion chips
    - "Add all to Workspace" button
    - "Add Individually" button
    """

    def __init__(
        self,
        message: str,
        suggestions: list[str],
        on_add_all=None,
        on_add_individual=None,
        on_chip_click=None
    ):
        """
        Create a suggestion card.

        Args:
            message: AI message to display
            suggestions: List of suggested items
            on_add_all: Callback when "Add all" is clicked (receives list)
            on_add_individual: Callback to show individual add UI
            on_chip_click: Callback when a chip is clicked (receives chip text)
        """
        self.message = message
        self.suggestions = suggestions
        self.on_add_all = on_add_all
        self.on_add_individual = on_add_individual
        self.on_chip_click = on_chip_click
        self.element = None

        self._build()

    def _build(self):
        """Build the UI elements."""
        with ui.element('div').classes('ai-suggestion-card') as card:
            self.element = card

            # Message
            ui.html(f'<p class="ai-suggestion-text">{self.message}</p>', sanitize=False)

            # Suggestion chips
            with ui.element('div').classes('suggestion-chips'):
                for suggestion in self.suggestions:
                    chip = ui.button(suggestion).classes('suggestion-chip')
                    if self.on_chip_click:
                        chip.on('click', lambda s=suggestion: self.on_chip_click(s))

            # Action buttons
            with ui.element('div').classes('ai-action-buttons'):
                # Add all button (gradient)
                ui.button(
                    f'Add all {len(self.suggestions)} to Workspace',
                    on_click=lambda: self.on_add_all(self.suggestions) if self.on_add_all else None
                ).classes('btn btn-primary')

                # Add individually button
                ui.button(
                    'Add Individually +',
                    on_click=self.on_add_individual
                ).classes('btn btn-outline')


def suggestion_card(
    message: str,
    suggestions: list[str],
    on_add_all=None,
    on_add_individual=None,
    on_chip_click=None
) -> SuggestionCard:
    """
    Factory function to create a suggestion card.

    Args:
        message: AI message to display
        suggestions: List of suggested items
        on_add_all: Callback when "Add all" is clicked
        on_add_individual: Callback to show individual add UI
        on_chip_click: Callback when a chip is clicked

    Returns:
        SuggestionCard instance
    """
    return SuggestionCard(
        message=message,
        suggestions=suggestions,
        on_add_all=on_add_all,
        on_add_individual=on_add_individual,
        on_chip_click=on_chip_click
    )


def chat_input(placeholder: str = 'Ask your assistant...', on_send=None):
    """
    Create a chat input bar at the bottom of the assistant pane.

    Args:
        placeholder: Placeholder text for the input
        on_send: Async callback when message is sent (receives message text)

    Returns:
        The input element for reference
    """
    with ui.element('div').classes('chat-input-wrapper'):
        input_el = ui.input(placeholder=placeholder).classes('flex-grow').props('borderless dense')

        async def handle_send():
            message = input_el.value
            if message and on_send:
                input_el.set_value('')
                await on_send(message)

        ui.button(icon='send', on_click=handle_send).props('flat round dense').classes('text-primary')
        input_el.on('keydown.enter', handle_send)

    return input_el


class ChatHistory:
    """
    Manages a scrollable chat history display with auto-scroll.
    """

    def __init__(self, initial_message: str = None):
        """
        Create a chat history container.

        Args:
            initial_message: Optional initial AI message to display
        """
        self.container = None
        self.messages = []

        with ui.element('div').classes('ai-chat-scroll-container') as container:
            self.container = container
            if initial_message:
                self.add_ai_message(initial_message)

    def _scroll_to_bottom(self):
        """
        Scroll the chat container to the bottom if user is already near bottom.

        This respects manual scrolling - if user has scrolled up to read history,
        we don't force them back to the bottom. Only auto-scroll if they're
        already at or near the bottom (within 100px threshold).
        """
        if self.container:
            # Check if user is near bottom before scrolling
            ui.run_javascript(f'''
                const container = document.querySelector('.ai-chat-scroll-container');
                if (container) {{
                    const scrollThreshold = 100; // pixels from bottom
                    const isNearBottom = container.scrollHeight - container.scrollTop - container.clientHeight < scrollThreshold;

                    // Only auto-scroll if user is already near bottom
                    if (isNearBottom) {{
                        container.scrollTop = container.scrollHeight;
                    }}
                }}
            ''')

    def add_ai_message(self, message: str):
        """Add an AI message to the history and scroll to bottom."""
        with self.container:
            with ui.element('div').classes('mb-3'):
                with ui.element('div').classes('ai-suggestion-text text-left'):
                    ui.markdown(message)

        # Auto-scroll to bottom after adding message (if user is at bottom)
        self._scroll_to_bottom()

    def add_user_message(self, message: str):
        """Add a user message to the history and scroll to bottom."""
        with self.container:
            with ui.element('div').classes('flex gap-2 mb-3 justify-end'):
                with ui.element('div').classes('bg-primary text-white rounded-lg px-3 py-2 max-w-xs'):
                    ui.label(message).classes('text-sm')

        # Auto-scroll to bottom after adding message (if user is at bottom)
        self._scroll_to_bottom()

    def add_spinner(self):
        """Add a thinking spinner. Returns the spinner element for later removal."""
        with self.container:
            spinner = ui.spinner(type='dots').classes('self-start')

        # Scroll to show spinner (if user is at bottom)
        self._scroll_to_bottom()

        return spinner

    def add_system_message(self, message: str):
        """Add a system notification message to the history and scroll to bottom."""
        with self.container:
            with ui.element('div').classes('flex justify-center my-3'):
                with ui.element('div').classes(
                    'bg-slate-100 text-slate-600 rounded-full px-4 py-2 text-sm inline-flex items-center gap-2'
                ):
                    ui.label(message)

        # Auto-scroll to bottom (if user is at bottom)
        self._scroll_to_bottom()

    def clear(self):
        """Clear all messages from the chat history."""
        self.container.clear()
        self.messages = []


class HybridAssistant:
    """
    Hybrid AI assistant that combines suggestion cards with chat functionality.

    Shows suggestion cards for structured responses, but allows free-form
    chat input for custom requests.
    """

    def __init__(
        self,
        initial_message: str = None,
        on_chat_send=None
    ):
        """
        Create a hybrid assistant panel.

        Args:
            initial_message: Initial AI greeting message
            on_chat_send: Async callback when user sends a chat message
        """
        self.on_chat_send = on_chat_send
        self.suggestion_container = None
        self.chat_history = None
        self.input_el = None

        self._build(initial_message)

    def _build(self, initial_message: str):
        """Build the UI elements."""
        with ui.element('div').classes('pane p-6 h-full flex flex-col'):
            # Header
            with ui.element('div').classes('ai-assistant-header'):
                with ui.element('div').classes('ai-icon-wrapper'):
                    ui.html('<span class="material-symbols-outlined">auto_awesome</span>', sanitize=False)
                ui.html('<h3 class="ai-assistant-title">AI Assistant</h3>', sanitize=False)

            # Suggestion/Response area
            with ui.element('div').classes('flex-grow flex flex-col overflow-hidden') as container:
                self.suggestion_container = container

                # Initial message card
                if initial_message:
                    with ui.element('div').classes('ai-suggestion-card flex-grow'):
                        ui.html(f'<p class="ai-suggestion-text">{initial_message}</p>', sanitize=False)

            # Chat input at bottom
            self.input_el = chat_input(
                placeholder='Ask your assistant...',
                on_send=self._handle_chat
            )

    async def _handle_chat(self, message: str):
        """Handle chat message from user."""
        if self.on_chat_send:
            await self.on_chat_send(message)

    def show_suggestions(
        self,
        message: str,
        suggestions: list[str],
        on_add_all=None,
        on_add_individual=None
    ):
        """
        Display a suggestion card, replacing any previous content.

        Args:
            message: AI message
            suggestions: List of suggestions
            on_add_all: Callback when "Add all" is clicked
            on_add_individual: Callback for individual add
        """
        self.suggestion_container.clear()
        with self.suggestion_container:
            suggestion_card(
                message=message,
                suggestions=suggestions,
                on_add_all=on_add_all,
                on_add_individual=on_add_individual
            )

    def show_message(self, message: str):
        """
        Display a simple message card.

        Args:
            message: Message to display
        """
        self.suggestion_container.clear()
        with self.suggestion_container:
            with ui.element('div').classes('ai-suggestion-card flex-grow'):
                ui.html(f'<p class="ai-suggestion-text">{message}</p>', sanitize=False)

    def show_loading(self):
        """Show a loading spinner."""
        self.suggestion_container.clear()
        with self.suggestion_container:
            with ui.element('div').classes('ai-suggestion-card flex-grow flex items-center justify-center'):
                ui.spinner(type='dots', size='lg')


class ChatAssistant:
    """
    Pure chat-based AI assistant with scrollable message history.

    Provides a simple conversation interface where messages accumulate
    in a scrollable container, supporting back-and-forth dialogue.
    """

    def __init__(
        self,
        initial_message: str = None,
        on_send=None
    ):
        """
        Create a chat assistant panel.

        Args:
            initial_message: Initial AI greeting message
            on_send: Async callback when user sends a message (receives message text)
        """
        self.on_send = on_send
        self.chat_history = None
        self.input_el = None
        self.current_spinner = None

        self._build(initial_message)

    def _build(self, initial_message: str):
        """Build the UI elements."""
        with ui.element('div').classes('pane p-6 ai-assistant-pane'):
            # Header
            with ui.element('div').classes('ai-assistant-header'):
                with ui.element('div').classes('ai-icon-wrapper'):
                    ui.html('<span class="material-symbols-outlined">auto_awesome</span>', sanitize=False)
                ui.html('<h3 class="ai-assistant-title">AI Assistant</h3>', sanitize=False)

            # Chat history (scrollable)
            self.chat_history = ChatHistory(initial_message=initial_message)

            # Chat input at bottom
            self.input_el = chat_input(
                placeholder='Ask your assistant...',
                on_send=self._handle_send
            )

    async def _handle_send(self, message: str):
        """Handle message from user."""
        # Add user message to history
        self.add_user_message(message)

        # Call the callback if provided
        if self.on_send:
            await self.on_send(message)

    def add_ai_message(self, message: str):
        """Add an AI message to the chat history."""
        self.hide_loading()
        self.chat_history.add_ai_message(message)

    def add_user_message(self, message: str):
        """Add a user message to the chat history."""
        self.chat_history.add_user_message(message)

    def show_loading(self):
        """Show a loading spinner in the chat."""
        if not self.current_spinner:
            self.current_spinner = self.chat_history.add_spinner()

    def hide_loading(self):
        """Hide the loading spinner."""
        if self.current_spinner:
            try:
                self.current_spinner.delete()
            except (ValueError, AttributeError):
                # Spinner already removed or doesn't exist
                pass
            finally:
                self.current_spinner = None


class TabbedAssistant:
    """
    AI Assistant with tabbed interface.

    Tabs:
    - Assistant: Chat interface (existing ChatHistory + chat_input)
    - Recent Conversations: List of past chats
    - Configuration: Output mode selection and prompt customization
    - Prompt Preview: Read-only view of full context
    """

    def __init__(
        self,
        initial_message: str = None,
        on_send=None,
        conversations: list[dict] = None,
        on_load_conversation=None,
        on_delete_conversation=None,
        on_new_chat=None,
        stage_context: str = "stage1",
        on_navigate=None,
    ):
        """
        Create a tabbed assistant panel.

        Args:
            initial_message: Initial AI greeting message
            on_send: Async callback when user sends a message
            conversations: List of conversation dicts from Supabase
            on_load_conversation: Async callback when user clicks a conversation
            on_delete_conversation: Async callback when user deletes a conversation
            on_new_chat: Callback when user clicks "Start New Chat"
            stage_context: The current stage (stage1, stage2, etc.) - controls output modes
            on_navigate: Callback when navigation to different stage is needed (receives stage, conversation_id)
        """
        self.on_send = on_send
        self.on_load_conversation = on_load_conversation
        self.on_delete_conversation = on_delete_conversation
        self.on_new_chat = on_new_chat
        self.on_navigate = on_navigate
        self.chat_history = None
        self.input_el = None
        self.current_spinner = None
        self.stage_context = stage_context

        # Configuration state
        self.output_mode = 'standard'
        self.mode_radio = None
        self.mode_selector = None  # Dropdown in Assistant tab header
        self.prompt_description_label = None
        self.system_prompt_input = None
        self.prompt_preview_container = None

        # Conversation state
        self.current_conversation_id: str | None = None
        self.conversations = conversations or []
        self.conversations_list = None
        self.tabs = None
        self.assistant_tab = None

        self._build(initial_message)

    def _build(self, initial_message: str):
        """Build the tabbed UI."""
        with ui.element('div').classes('pane p-6 ai-assistant-pane'):
            # Header with mode selector
            with ui.element('div').classes('ai-assistant-header'):
                # Left side: icon and title
                with ui.element('div').classes('flex items-center gap-3'):
                    with ui.element('div').classes('ai-icon-wrapper'):
                        ui.html('<span class="material-symbols-outlined">auto_awesome</span>', sanitize=False)
                    ui.html('<h3 class="ai-assistant-title">AI Assistant</h3>', sanitize=False)

                # Right side: mode selector
                # Get stage-specific modes for dropdown options
                stage_config = get_stage_config(self.stage_context)
                modes = stage_config.get("modes", [])
                mode_options = [m["display_name"] for m in modes]
                default_mode = stage_config.get("default_mode", "standard")

                # Find default display name
                default_display = next(
                    (m["display_name"] for m in modes if m["key"] == default_mode),
                    mode_options[0] if mode_options else "Standard"
                )

                # Capsule/pill dropdown with icon
                self.mode_selector = ui.select(
                    options=mode_options,
                    value=default_display,
                    on_change=self._handle_mode_selector_change
                ).props('dense borderless').classes('mode-selector-dropdown')

                # Add sparkle icon to the left of the dropdown
                self.mode_selector.props(add='prepend')
                with self.mode_selector:
                    with self.mode_selector.add_slot('prepend'):
                        ui.html('<span class="material-symbols-outlined mode-selector-icon">auto_awesome</span>', sanitize=False)

            # Tabs
            with ui.tabs().classes('ai-assistant-tabs') as tabs:
                self.tabs = tabs
                self.assistant_tab = ui.tab('Assistant')
                conversations_tab = ui.tab('Recent Chats')
                config_tab = ui.tab('Configuration')
                preview_tab = ui.tab('Prompt Preview')

            # Tab Panels
            with ui.tab_panels(tabs, value=self.assistant_tab).classes('ai-assistant-tab-panels flex-grow'):
                # Tab 1: Assistant (Chat)
                with ui.tab_panel(self.assistant_tab).classes('flex flex-col w-full h-full'):
                    self._build_chat_tab(initial_message)

                # Tab 2: Recent Conversations
                with ui.tab_panel(conversations_tab).classes('flex flex-col h-full min-h-0'):
                    self._build_conversations_tab()

                # Tab 3: Configuration
                with ui.tab_panel(config_tab).classes('flex flex-col h-full'):
                    self._build_config_tab()

                # Tab 4: Prompt Preview
                with ui.tab_panel(preview_tab).classes('flex flex-col h-full'):
                    self._build_preview_tab()

    def _build_chat_tab(self, initial_message: str):
        """Build the chat tab content."""
        with ui.element('div').classes('flex flex-col w-full h-full'):
            # Chat history (scrollable)
            self.chat_history = ChatHistory(initial_message=initial_message)

            # Chat input at bottom
            self.input_el = chat_input(
                placeholder='Ask your assistant...',
                on_send=self._handle_send
            )

    async def _handle_send(self, message: str):
        """Handle message from user."""
        self.add_user_message(message)
        if self.on_send:
            await self.on_send(message)

    def _handle_mode_selector_change(self, e):
        """Handle mode change from dropdown selector (Assistant tab)."""
        display_name = e.value
        mode_config = self._stage_modes.get(display_name, {})
        new_mode = mode_config.get("key", "standard")

        # Only show notification if mode actually changed
        if new_mode != self.output_mode:
            self.output_mode = new_mode

            # Sync to Configuration tab radio
            if self.mode_radio:
                self.mode_radio.set_value(display_name)

            # Add system message to chat history
            if self.chat_history:
                self.chat_history.add_system_message(
                    f"Switched to {display_name} mode."
                )
        else:
            self.output_mode = new_mode

        # Update description and system prompt (same logic as config tab)
        if self.prompt_description_label:
            self.prompt_description_label.set_text(
                mode_config.get("description", "")
            )

        # Handle custom mode
        if self.system_prompt_input:
            if new_mode == "custom":
                self.system_prompt_input.props(remove='readonly')
            else:
                self.system_prompt_input.props(add='readonly')

            # Update system prompt display
            self.system_prompt_input.set_value(
                mode_config.get("system_prompt", "")
            )

    def _build_conversations_tab(self):
        """Build the recent conversations tab content."""
        with ui.element('div').classes('flex-grow overflow-y-auto min-h-0') as conv_list:
            self.conversations_list = conv_list
            self._render_conversations_list()

    def _render_conversations_list(self):
        """Render the list of conversations."""
        if self.conversations_list:
            self.conversations_list.clear()

        with self.conversations_list:
            if not self.conversations:
                ui.label('No previous conversations').classes('text-slate-500 text-sm p-4')
                return

            for conv in self.conversations:
                conv_id = conv.get('id')
                raw_title = conv.get('title') or 'Untitled conversation'
                conv_stage = conv.get('stage', 'stage1')
                stage_display = get_stage_display_name(conv_stage)

                # Prepend stage name to title
                title = f"{stage_display}: {raw_title}"

                updated_at = conv.get('updated_at', '')
                formatted_time = format_relative_timestamp(updated_at) if updated_at else ''

                with ui.element('div').classes(
                    'flex items-center gap-3 py-3 px-2 hover:bg-slate-50 cursor-pointer border-b border-slate-100 group'
                ):
                    # Clickable area for loading conversation
                    with ui.element('div').classes('flex items-center gap-3 flex-grow min-w-0').on(
                        'click', lambda e, cid=conv_id, cs=conv_stage: self._handle_load_conversation(cid, cs)
                    ):
                        # Chat icon
                        ui.icon('chat_bubble').classes('text-primary text-xl flex-shrink-0')

                        # Title with stage prefix
                        ui.label(title).classes('flex-grow text-sm font-medium text-slate-800 truncate')

                    # Timestamp and delete button
                    ui.label(formatted_time).classes('text-xs text-slate-500 whitespace-nowrap')
                    ui.button(
                        icon='delete_outline',
                        on_click=lambda e, cid=conv_id, t=raw_title: self._confirm_delete_conversation(cid, t)
                    ).props('flat round dense size=sm').classes(
                        'opacity-0 group-hover:opacity-100 text-slate-400 hover:text-red-500 ml-1'
                    )

    async def _handle_new_chat(self):
        """Handle starting a new chat."""
        self.current_conversation_id = None
        if self.chat_history:
            self.chat_history.clear()
        if self.on_new_chat:
            self.on_new_chat()
        # Switch to assistant tab
        if self.tabs and self.assistant_tab:
            self.tabs.set_value(self.assistant_tab)

    async def _handle_load_conversation(self, conversation_id: str, conversation_stage: str = None):
        """Handle loading a conversation, navigating to correct stage if needed."""
        # If conversation is from a different stage, navigate there first
        if conversation_stage and conversation_stage != self.stage_context:
            if self.on_navigate:
                self.on_navigate(conversation_stage, conversation_id)
            return

        # Same stage - load normally
        self.current_conversation_id = conversation_id
        if self.on_load_conversation:
            await self.on_load_conversation(conversation_id)
        # Switch to assistant tab
        if self.tabs and self.assistant_tab:
            self.tabs.set_value(self.assistant_tab)

    def _confirm_delete_conversation(self, conversation_id: str, title: str):
        """Show confirmation dialog before deleting a conversation."""
        with ui.dialog() as dialog, ui.card().classes('p-4'):
            ui.label(f'Delete "{title}"?').classes('text-lg font-medium mb-2')
            ui.label('This will permanently delete this conversation and all its messages.').classes(
                'text-sm text-slate-600 mb-4'
            )
            with ui.row().classes('gap-2 justify-end w-full'):
                ui.button('Cancel', on_click=dialog.close).props('flat')
                ui.button(
                    'Delete',
                    on_click=lambda: self._execute_delete(conversation_id, dialog)
                ).classes('bg-red-500 text-white')
        dialog.open()

    async def _execute_delete(self, conversation_id: str, dialog):
        """Execute the deletion after confirmation."""
        dialog.close()
        await self._handle_delete_conversation(conversation_id)

    async def _handle_delete_conversation(self, conversation_id: str):
        """Handle deleting a conversation."""
        if self.on_delete_conversation:
            success = await self.on_delete_conversation(conversation_id)
            if success:
                # Remove from local list and re-render
                self.conversations = [c for c in self.conversations if c.get('id') != conversation_id]
                self._render_conversations_list()
                # If we deleted the current conversation, start fresh
                if self.current_conversation_id == conversation_id:
                    await self._handle_new_chat()

    def update_conversations(self, conversations: list[dict]):
        """Update the conversations list."""
        self.conversations = conversations
        self._render_conversations_list()

    def add_conversation_to_list(self, conversation: dict):
        """Add a new conversation to the top of the list."""
        self.conversations.insert(0, conversation)
        self._render_conversations_list()

    def _build_config_tab(self):
        """Build the configuration tab content."""
        # Get stage-specific configuration
        stage_config = get_stage_config(self.stage_context)
        modes = stage_config.get("modes", [])
        default_mode = stage_config.get("default_mode", "standard")

        # Build options list for radio
        mode_options = [m["display_name"] for m in modes]

        # Find default display name
        default_display = next(
            (m["display_name"] for m in modes if m["key"] == default_mode),
            mode_options[0] if mode_options else "Standard"
        )

        # Store modes for lookup
        self._stage_modes = {m["display_name"]: m for m in modes}

        with ui.element('div').classes('px-4 pt-4 pb-2 h-full flex flex-col'):
            # Output Mode Selection
            with ui.element('div').classes('config-section'):
                ui.label('Output Mode').classes('config-label')
                self.mode_radio = ui.radio(
                    options=mode_options,
                    value=default_display,
                    on_change=self._handle_mode_change
                ).props('inline')

            # Prompt Description (read-only text display)
            with ui.element('div').classes('config-section'):
                ui.label('Prompt Description').classes('config-label')
                # Get initial description
                initial_mode = self._stage_modes.get(default_display, {})
                self.prompt_description_label = ui.label(
                    initial_mode.get("description", "Select an output mode.")
                ).classes('text-secondary text-sm')

            # System Prompt - flex-grow to fill remaining space
            with ui.element('div').classes('config-section flex-grow flex flex-col min-h-0'):
                ui.label('System Prompt').classes('config-label')
                self.system_prompt_input = ui.textarea(
                    placeholder='Enter system prompt...',
                    value=initial_mode.get("system_prompt", "")
                ).classes('config-textarea').props('readonly')
                self._system_prompt_readonly = True

    def _handle_mode_change(self, e):
        """Handle output mode radio button change (Configuration tab)."""
        display_name = e.value
        mode_config = self._stage_modes.get(display_name, {})
        new_mode = mode_config.get("key", "standard")

        # Only show notification if mode actually changed
        if new_mode != self.output_mode:
            self.output_mode = new_mode

            # Sync to Assistant tab dropdown
            if self.mode_selector:
                self.mode_selector.set_value(display_name)

            # Add system message to chat history
            if self.chat_history:
                self.chat_history.add_system_message(
                    f"Switched to {display_name} mode."
                )
        else:
            self.output_mode = new_mode

        # Update description
        self.prompt_description_label.set_text(
            mode_config.get("description", "")
        )

        # Handle custom mode
        if new_mode == "custom":
            self.system_prompt_input.props(remove='readonly')
        else:
            self.system_prompt_input.props(add='readonly')

        # Update system prompt display
        self.system_prompt_input.set_value(
            mode_config.get("system_prompt", "")
        )

    def _build_preview_tab(self):
        """Build the prompt preview tab content."""
        placeholder_prompt = """=== SYSTEM PROMPT ===
You are an educational assessment expert. Generate clear, measurable rubric criteria based on the provided prompt groups and task context.

=== USER CONTEXT ===
Project: [Project Title]
Task: [Task Name]

Prompt Groups:
1. [Group 1 Name]
   - Description: [Description]

2. [Group 2 Name]
   - Description: [Description]

=== INSTRUCTIONS ===
Generate a rubric with the following structure:
- Each prompt group becomes a criterion
- Include 4 performance levels (Exemplary, Proficient, Developing, Beginning)
- Provide clear, measurable descriptors for each level"""

        with ui.element('div').classes('prompt-preview-container') as container:
            self.prompt_preview_container = container
            ui.label(placeholder_prompt).classes('whitespace-pre-wrap')

    # Public API methods (same as ChatAssistant)
    def add_ai_message(self, message: str):
        """Add an AI message to the chat history."""
        self.hide_loading()
        if self.chat_history:
            self.chat_history.add_ai_message(message)

    def add_user_message(self, message: str):
        """Add a user message to the chat history."""
        if self.chat_history:
            self.chat_history.add_user_message(message)

    def show_loading(self):
        """Show a loading spinner in the chat."""
        if self.chat_history and not self.current_spinner:
            self.current_spinner = self.chat_history.add_spinner()

    def hide_loading(self):
        """Hide the loading spinner."""
        if self.current_spinner:
            try:
                self.current_spinner.delete()
            except (ValueError, AttributeError):
                # Spinner already removed or doesn't exist
                pass
            finally:
                self.current_spinner = None