frontend/components/workspace.py

"""
Workspace Components for Rubric AI
Contains: workspace_pane, workspace_item, button_group, AssessmentItemPicker
"""

from contextlib import contextmanager
from typing import Callable, List, Dict, Any
from nicegui import ui
from .layout import material_icon


@contextmanager
def workspace_pane(title: str, subtitle: str = None, action_button: tuple = None):
    """
    Create a workspace pane (left side of the two-column layout).

    Usage:
        with workspace_pane('Workspace', 'Stage 1: Define Prompt Groups'):
            # Your workspace content here

    Args:
        title: Pane title (e.g., 'Workspace')
        subtitle: Pane subtitle (e.g., 'Stage 1: Define Prompt Groups')
        action_button: Optional tuple of (label, on_click, icon) for header action button
    """
    with ui.element('div').classes('pane'):
        # Header
        with ui.element('div').classes('pane-header'):
            with ui.element('div').classes('flex justify-between items-center w-full'):
                with ui.element('div'):
                    ui.html(f'<h2 class="pane-title">{title}</h2>', sanitize=False)
                    if subtitle:
                        ui.html(f'<p class="pane-subtitle">{subtitle}</p>', sanitize=False)
                if action_button:
                    label, on_click, *rest = action_button
                    icon = rest[0] if rest else 'add'
                    ui.button(label, on_click=on_click).classes('btn btn-primary').props(f'icon={icon}')

        # Content area
        with ui.element('div').classes('pane-content') as content:
            yield content


class WorkspaceItem:
    """
    A draggable workspace list item with edit and delete capabilities.

    This creates an item that looks like:
    [drag_handle] [editable_text_input] [more_vert_menu]
    """

    def __init__(
        self,
        name: str,
        on_delete=None,
        on_rename=None,
        on_menu_click=None,
        container=None
    ):
        """
        Create a workspace item.

        Args:
            name: Initial name/text for the item
            on_delete: Callback when item is deleted (receives name)
            on_rename: Callback when item is renamed (receives old_name, new_name)
            on_menu_click: Callback when menu button is clicked
            container: Parent container (for removal)
        """
        self.name = name
        self.on_delete = on_delete
        self.on_rename = on_rename
        self.on_menu_click = on_menu_click
        self.container = container
        self.element = None
        self.input_element = None

        self._build()

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

            # Drag handle
            ui.html('<span class="material-symbols-outlined drag-handle">drag_indicator</span>', sanitize=False)

            # Editable input
            self.input_element = ui.input(value=self.name).classes(
                'flex-grow bg-transparent border-none outline-none'
            ).props('borderless dense')

            # Bind rename callback
            if self.on_rename:
                def handle_rename(e):
                    new_name = e.value
                    if new_name != self.name:
                        self.on_rename(self.name, new_name)
                        self.name = new_name

                self.input_element.on('blur', handle_rename)

            # More menu button
            with ui.button(icon='more_vert').props('flat round dense').classes('item-actions'):
                with ui.menu() as menu:
                    ui.menu_item('Rename', lambda: self.input_element.run_method('focus'))
                    ui.menu_item('Delete', self._handle_delete)

    def _handle_delete(self):
        """Handle deletion of this item."""
        if self.on_delete:
            self.on_delete(self.name)

        # Remove from UI
        if self.element:
            self.element.delete()

    def delete(self):
        """Programmatically delete this item."""
        self._handle_delete()


def workspace_item(
    name: str,
    on_delete=None,
    on_rename=None
) -> WorkspaceItem:
    """
    Factory function to create a workspace item.

    Args:
        name: Initial name/text for the item
        on_delete: Callback when item is deleted (receives name)
        on_rename: Callback when item is renamed (receives old_name, new_name)

    Returns:
        WorkspaceItem instance
    """
    return WorkspaceItem(name=name, on_delete=on_delete, on_rename=on_rename)


@contextmanager
def workspace_list():
    """
    Create a container for workspace items.

    Usage:
        with workspace_list() as item_list:
            workspace_item('Thesis & Argument', on_delete=handle_delete)
            workspace_item('Use of Evidence', on_delete=handle_delete)
    """
    with ui.element('div').classes('workspace-list') as container:
        yield container


def button_group(
    back_label: str = 'Back',
    next_label: str = 'Next',
    on_back=None,
    on_next=None,
    next_icon: str = 'arrow_forward',
    show_back: bool = True
):
    """
    Create a button group for navigation (Back / Next buttons).

    Args:
        back_label: Label for back button
        next_label: Label for next button
        on_back: Callback for back button
        on_next: Callback for next button
        next_icon: Icon for next button (default: arrow_forward)
        show_back: Whether to show the back button
    """
    with ui.element('div').classes('button-group'):
        if show_back:
            ui.button(back_label, on_click=on_back).classes('btn btn-secondary')
        else:
            # Spacer to push next button to the right
            ui.element('div')

        with ui.button(on_click=on_next).classes('btn btn-primary'):
            ui.label(next_label)
            if next_icon:
                ui.html(f'<span class="material-symbols-outlined">{next_icon}</span>', sanitize=False)


def add_item_input(placeholder: str = 'Add new item...', on_add=None):
    """
    Create an input field with add button for adding new items.

    Args:
        placeholder: Placeholder text for the input
        on_add: Callback when add button is clicked (receives input value)
    """
    with ui.element('div').classes('flex items-center gap-2 mt-4 pt-4 border-t border-slate-200'):
        input_el = ui.input(placeholder=placeholder).classes('flex-grow').props('dense')

        def handle_add():
            if input_el.value and on_add:
                on_add(input_el.value)
                input_el.set_value('')

        ui.button('Add', on_click=handle_add).classes('btn btn-secondary')
        input_el.on('keydown.enter', handle_add)

    return input_el


class AssessmentItemPicker:
    """
    A modal dialog for selecting assessment items from a task's content_stream.

    Displays a searchable list of assessment items with checkboxes,
    allowing users to select multiple items to add to their prompt groups.
    """

    def __init__(
        self,
        task_data: Dict[str, Any],
        on_add_group: Callable[[str, List[Dict[str, Any]]], None] = None,
        used_item_ids: List[str] = None
    ):
        """
        Create an assessment item picker modal.

        Args:
            task_data: The task_data JSONB from the tasks table
            on_add_group: Callback when a group is added (receives group_name and list of selected items)
            used_item_ids: List of item IDs already used in other groups (will be greyed out)
        """
        self.task_data = task_data
        self.on_add_group = on_add_group
        self.used_item_ids = set(used_item_ids or [])
        self.assessment_items = self._extract_assessment_items()
        self.selected_items: Dict[str, bool] = {}  # id -> selected
        self.group_name = ""
        self.dialog = None
        self.items_container = None
        self.group_name_input = None
        self.checkboxes: Dict[str, ui.checkbox] = {}

    def _extract_assessment_items(self) -> List[Dict[str, Any]]:
        """Extract assessment items from task_data content_stream (no tags)."""
        items = []
        content_stream = self.task_data.get('content_stream', [])

        for item in content_stream:
            if item.get('type') == 'assessment_item':
                items.append({
                    'id': item.get('id', str(item.get('sequence_id', ''))),
                    'prompt_text': item.get('prompt_text', ''),
                    'question_type': item.get('question_type', 'unknown'),
                    'is_scorable': item.get('is_scorable', False),
                    'sequence_id': item.get('sequence_id', 0),
                    # Additional fields for table display
                    'options': item.get('options', []),
                    'correct_answer_text': item.get('correct_answer_text', '')
                })

        # Sort by sequence_id
        items.sort(key=lambda x: x.get('sequence_id', 0))
        return items

    def _truncate_text(self, text: str, max_length: int = 100) -> str:
        """Truncate text and add ellipsis if too long."""
        if len(text) <= max_length:
            return text
        return text[:max_length].rsplit(' ', 1)[0] + '...'

    def _format_display_text(self, item: Dict[str, Any]) -> str:
        """Format the prompt text for display."""
        prompt = item.get('prompt_text', '')
        # Clean up newlines and extra whitespace
        prompt = ' '.join(prompt.split())
        return self._truncate_text(prompt, 80)

    def _get_question_type_label(self, question_type: str) -> str:
        """Get a human-readable label for question type."""
        labels = {
            'drawing': 'Drawing',
            'multiple_choice': 'Multiple Choice',
            'open_response': 'Open Response',
            'table_completion': 'Table'
        }
        return labels.get(question_type, question_type.title())


    def _is_correct_option(self, option: str, correct_answer_text: str) -> bool:
        """Check if an option matches the correct answer text."""
        if not correct_answer_text:
            return False
        # Options have format like "A. Answer text" or "A) Answer text"
        # correct_answer_text is just the answer without the letter prefix
        # Strip the letter prefix (e.g., "A. ", "B) ") from option
        import re
        cleaned_option = re.sub(r'^[A-Z][.\)]\s*', '', option)
        return cleaned_option.strip() == correct_answer_text.strip()

    def _format_markdown_bold(self, text: str) -> str:
        """Convert **text** markdown to HTML bold tags."""
        import re
        # Replace **text** with <b>text</b>
        return re.sub(r'\*\*(.+?)\*\*', r'<b>\1</b>', text)

    def _render_items(self):
        """Render the list of assessment items as a table (no tags)."""
        self.items_container.clear()
        self.checkboxes.clear()

        with self.items_container:
            if not self.assessment_items:
                ui.label('No assessment items found.').classes('text-secondary text-center py-4')
                return

            # Table header - 3 columns only
            with ui.element('div').classes('grid grid-cols-9 gap-2 px-3 py-2 bg-slate-100 border-b font-semibold text-xs text-slate-600 uppercase tracking-wide'):
                ui.element('div').classes('col-span-1')  # Checkbox column
                ui.label('Item').classes('col-span-5')
                ui.label('Response Options').classes('col-span-3')

            # Table rows
            for item in self.assessment_items:
                item_id = item['id']
                is_used = item_id in self.used_item_ids
                is_selected = self.selected_items.get(item_id, False)

                # Determine row classes based on state
                row_classes = 'grid grid-cols-9 gap-2 px-3 py-3 border-b border-slate-100 items-start '
                if is_used:
                    row_classes += 'opacity-50 cursor-not-allowed bg-slate-50'
                elif is_selected:
                    row_classes += 'cursor-pointer hover:bg-slate-50 bg-primary/10'
                else:
                    row_classes += 'cursor-pointer hover:bg-slate-50'

                with ui.element('div').classes(row_classes) as row:
                    # Checkbox column
                    with ui.element('div').classes('col-span-1 flex items-start pt-1'):
                        def make_checkbox_handler(iid):
                            def handler(e):
                                self._toggle_item(iid, e.value)
                            return handler

                        cb = ui.checkbox(
                            value=is_selected,
                            on_change=make_checkbox_handler(item_id) if not is_used else None
                        ).props('dense' + (' disable' if is_used else ''))
                        self.checkboxes[item_id] = cb

                    # Make the whole row clickable (but not if used)
                    if not is_used:
                        def make_row_handler(iid):
                            def handler(e):
                                self._toggle_item_from_row(iid)
                            return handler
                        row.on('click', make_row_handler(item_id))

                    # ITEM column - full text, wrapped
                    with ui.element('div').classes('col-span-5'):
                        # Item number badge
                        with ui.element('div').classes('flex items-center gap-2 mb-1'):
                            ui.label(f"Item {item['id']}").classes(
                                'text-xs font-semibold ' + ('text-slate-400' if is_used else 'text-primary')
                            )
                            ui.element('span').classes(
                                'text-xs px-2 py-0.5 rounded-full bg-slate-200 text-slate-600'
                            ).text = self._get_question_type_label(item['question_type'])
                            if is_used:
                                ui.element('span').classes(
                                    'text-xs px-2 py-0.5 rounded-full bg-amber-100 text-amber-700'
                                ).text = 'In use'
                        # Full prompt text (wrapped, no truncation)
                        prompt_text = item.get('prompt_text', '')
                        ui.label(prompt_text).classes(
                            'text-sm leading-relaxed break-words ' + ('text-slate-400' if is_used else 'text-slate-700')
                        )

                    # RESPONSE OPTIONS column
                    with ui.element('div').classes('col-span-3'):
                        options = item.get('options', [])
                        correct_answer = item.get('correct_answer_text', '')
                        if options:
                            # Multiple choice - show all options with correct highlighted
                            for option in options:
                                is_correct = self._is_correct_option(option, correct_answer)
                                option_classes = 'text-xs leading-relaxed mb-1 '
                                if is_used:
                                    option_classes += 'text-slate-400'
                                elif is_correct:
                                    option_classes += 'text-green-600 font-medium'
                                else:
                                    option_classes += 'text-slate-600'
                                ui.label(option).classes(option_classes)
                        elif correct_answer:
                            # Non-MC item with correct_answer_text - show with "Ideal Response:" label
                            # Light blue color (#64b5f6) and handle markdown bold
                            ui.label('Ideal Response:').classes(
                                'text-xs font-semibold mb-1 ' +
                                ('text-slate-400' if is_used else '')
                            ).style('' if is_used else 'color: #64b5f6;')
                            # Format markdown bold and render as HTML
                            formatted_answer = self._format_markdown_bold(correct_answer)
                            ui.html(
                                f'<span class="text-xs leading-relaxed" style="color: {"#94a3b8" if is_used else "#64b5f6"};">{formatted_answer}</span>',
                                sanitize=False
                            )
                        else:
                            ui.label('—').classes('text-xs text-slate-400')


    # ========== OLD _render_items (COMMENTED OUT) ==========
    # def _render_items_old(self):
    #     """Render the list of assessment items."""
    #     self.items_container.clear()
    #     self.checkboxes.clear()
    #
    #     with self.items_container:
    #         if not self.assessment_items:
    #             ui.label('No assessment items found.').classes('text-secondary text-center py-4')
    #             return
    #
    #         for item in self.assessment_items:
    #             item_id = item['id']
    #             is_used = item_id in self.used_item_ids
    #             is_selected = self.selected_items.get(item_id, False)
    #
    #             # Determine row classes based on state
    #             row_classes = 'flex items-start gap-3 p-3 rounded-lg '
    #             if is_used:
    #                 row_classes += 'opacity-50 cursor-not-allowed bg-slate-100'
    #             elif is_selected:
    #                 row_classes += 'cursor-pointer hover:bg-slate-50 bg-primary/10 border border-primary/30'
    #             else:
    #                 row_classes += 'cursor-pointer hover:bg-slate-50 border border-transparent'
    #
    #             with ui.element('div').classes(row_classes) as row:
    #                 # Checkbox - disabled if item is already used
    #                 def make_checkbox_handler(iid):
    #                     def handler(e):
    #                         self._toggle_item(iid, e.value)
    #                     return handler
    #
    #                 cb = ui.checkbox(
    #                     value=is_selected,
    #                     on_change=make_checkbox_handler(item_id) if not is_used else None
    #                 ).props('dense' + (' disable' if is_used else ''))
    #                 self.checkboxes[item_id] = cb
    #
    #                 # Make the whole row clickable (but not if used)
    #                 if not is_used:
    #                     def make_row_handler(iid):
    #                         def handler(e):
    #                             self._toggle_item_from_row(iid)
    #                         return handler
    #                     row.on('click', make_row_handler(item_id))
    #
    #                 # Item content
    #                 with ui.element('div').classes('flex-grow min-w-0'):
    #                     # Question type badge and item number
    #                     with ui.element('div').classes('flex items-center gap-2 mb-1'):
    #                         ui.label(f"Item {item['id']}").classes(
    #                             'text-xs font-semibold ' + ('text-slate-400' if is_used else 'text-primary')
    #                         )
    #                         ui.element('span').classes(
    #                             'text-xs px-2 py-0.5 rounded-full bg-slate-200 text-slate-600'
    #                         ).text = self._get_question_type_label(item['question_type'])
    #                         if item.get('is_scorable'):
    #                             ui.element('span').classes(
    #                                 'text-xs px-2 py-0.5 rounded-full bg-green-100 text-green-700'
    #                             ).text = 'Scorable'
    #                         # Show "In use" badge for used items
    #                         if is_used:
    #                             ui.element('span').classes(
    #                                 'text-xs px-2 py-0.5 rounded-full bg-amber-100 text-amber-700'
    #                             ).text = 'In use'
    #
    #                     # Prompt text
    #                     ui.label(self._format_display_text(item)).classes(
    #                         'text-sm leading-relaxed ' + ('text-slate-400' if is_used else 'text-slate-700')
    #                     )
    # ========== END OLD _render_items ==========

    def _toggle_item(self, item_id: str, value: bool):
        """Toggle selection state of an item."""
        print(f"DEBUG _toggle_item: item_id={item_id}, value={value}")
        self.selected_items[item_id] = value
        print(f"DEBUG selected_items now: {self.selected_items}")

    def _toggle_item_from_row(self, item_id: str):
        """Toggle item when row is clicked."""
        current = self.selected_items.get(item_id, False)
        new_value = not current
        print(f"DEBUG _toggle_item_from_row: item_id={item_id}, current={current}, new_value={new_value}")
        self.selected_items[item_id] = new_value
        print(f"DEBUG selected_items now: {self.selected_items}")
        # Update the checkbox visually
        if item_id in self.checkboxes:
            self.checkboxes[item_id].set_value(new_value)

    async def _handle_add_selected(self):
        """Handle adding selected items to a prompt group."""
        group_name = self.group_name_input.value.strip() if self.group_name_input else ""

        if not group_name:
            ui.notify('Please enter a name for the prompt group.', color='warning')
            return

        selected = [
            item.copy() for item in self.assessment_items
            if self.selected_items.get(item['id'], False)
        ]

        if not selected:
            ui.notify('Please select at least one item.', color='warning')
            return

        if self.on_add_group:
            # Handle both sync and async callbacks
            import asyncio
            result = self.on_add_group(group_name, selected)
            if asyncio.iscoroutine(result):
                await result

        self.dialog.close()

    def _get_selected_count(self) -> int:
        """Get count of selected items."""
        return sum(1 for v in self.selected_items.values() if v)

    def show(self):
        """Show the modal dialog."""
        # Use max-w-5xl for wider table layout
        with ui.dialog() as self.dialog, ui.card().classes('w-full max-w-5xl'):
            # Header
            with ui.element('div').classes('flex justify-between items-center mb-4'):
                ui.label('Add Prompt Group').classes('text-xl font-bold text-slate-800')
                ui.button(icon='close', on_click=self.dialog.close).props('flat round dense')

            # Group name input
            self.group_name_input = ui.input(
                label='Prompt Group Name',
                placeholder='e.g., Model Building, Written Explanation...'
            ).classes('w-full mb-4').props('outlined')

            # Items list label
            ui.label('Select assessment items for this group:').classes('text-sm text-slate-600 mb-2')

            # Items list (scrollable) - increased height for table view
            with ui.element('div').classes('max-h-96 overflow-y-auto border rounded-lg') as container:
                self.items_container = container
                self._render_items()

            # Footer with buttons
            with ui.element('div').classes('flex justify-end gap-3 mt-4 pt-4 border-t'):
                ui.button('Cancel', on_click=self.dialog.close).classes('btn btn-secondary')
                ui.button('Add Selected', on_click=self._handle_add_selected).classes('btn btn-primary')

        self.dialog.open()

    def close(self):
        """Close the modal dialog."""
        if self.dialog:
            self.dialog.close()