# dataset_editor_interface.py
"""
Dataset Editor Interface for Enhanced Datasets.

Provides a standalone interface for editing test datasets with full CRUD operations.
"""

import gradio as gr
from typing import List, Dict, Tuple, Optional, Any
from datetime import datetime
import uuid

from src.core.verification_models import TestMessage, TestDataset
from src.interface.enhanced_dataset_interface import EnhancedDatasetInterfaceController


def create_dataset_editor_interface() -> gr.Blocks:
    """
    Create the dataset editor interface.
    
    Returns:
        Gradio Blocks component for dataset editing
    """
    with gr.Blocks(title="Edit Enhanced Datasets") as interface:
        # Initialize controller
        dataset_controller = EnhancedDatasetInterfaceController()
        
        # Pre-load dataset choices for initialization
        try:
            initial_choices, _, _, _ = dataset_controller.initialize_interface()
        except Exception:
            initial_choices = []
        
        # Application state
        current_dataset_state = gr.State(value=None)
        editing_dataset_state = gr.State(value=None)
        
        # Header
        gr.Markdown("# ✏️ Edit Enhanced Datasets")
        gr.Markdown("Manage and customize test datasets for verification. Add, modify, or delete test cases to create datasets tailored to your testing needs.")
        
        # Dataset selection interface
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("## 📋 Select Dataset")
                
                dataset_selector = gr.Dropdown(
                    choices=initial_choices,
                    value=initial_choices[0] if initial_choices else None,
                    label="Available Datasets",
                    info="Choose a dataset to edit",
                    interactive=True
                )
                
                with gr.Row():
                    edit_dataset_btn = gr.Button("✏️ Edit Dataset", variant="primary", scale=2)
                    create_dataset_btn = gr.Button("➕ Create New Dataset", variant="secondary", scale=1)
            
            with gr.Column(scale=1):
                gr.Markdown("## 📊 Dataset Information")
                dataset_info_display = gr.Markdown(
                    "Select a dataset to view details",
                    label="Dataset Details"
                )
        
        # Create new dataset section (initially hidden)
        create_section = gr.Row(visible=False)
        with create_section:
            with gr.Column():
                gr.Markdown("## ➕ Create New Dataset")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        new_dataset_name = gr.Textbox(
                            label="Dataset Name",
                            placeholder="Enter dataset name...",
                            interactive=True
                        )
                    with gr.Column(scale=2):
                        new_dataset_description = gr.Textbox(
                            label="Description",
                            placeholder="Describe the dataset purpose...",
                            interactive=True
                        )
                    with gr.Column(scale=1):
                        template_selector = gr.Dropdown(
                            choices=["", "empty", "mixed_sample"],
                            label="Template",
                            value="",
                            info="Optional template"
                        )
                
                with gr.Row():
                    create_confirm_btn = gr.Button("✅ Create Dataset", variant="primary")
                    create_cancel_btn = gr.Button("❌ Cancel", variant="secondary")
        
        # Edit Dataset interface (initially hidden)
        edit_section = gr.Row(visible=False)
        with edit_section:
            with gr.Column():
                gr.Markdown("## ✏️ Edit Dataset")
                
                # Dataset actions
                with gr.Row():
                    save_dataset_btn = gr.Button("💾 Save Changes", variant="primary")
                    close_edit_btn = gr.Button("✖️ Close Editor", variant="secondary")
                
                # Test cases display
                test_cases_display = gr.HTML(
                    value="<p>Loading test cases...</p>",
                    label="Test Cases"
                )
                
                gr.Markdown("### ➕ Add New Test Case")
                with gr.Row():
                    with gr.Column(scale=3):
                        new_message_text = gr.Textbox(
                            label="Message Text",
                            placeholder="Enter the test message...",
                            lines=3,
                            interactive=True
                        )
                    with gr.Column(scale=1):
                        new_message_classification = gr.Radio(
                            choices=["green", "yellow", "red"],
                            label="Expected Classification",
                            value="yellow"
                        )
                
                with gr.Row():
                    add_test_case_btn = gr.Button("➕ Add Test Case", variant="primary")
                    clear_form_btn = gr.Button("🗑️ Clear Form", variant="secondary")
        
        # Status message
        status_message = gr.Markdown("", visible=True)
        
        # Event handlers
        def on_dataset_selection_change(dataset_selection):
            """Handle dataset selection change."""
            try:
                dataset_info, dataset_obj = dataset_controller.get_dataset_info(dataset_selection)
                return (
                    dataset_info,     # dataset_info_display
                    dataset_obj       # current_dataset_state
                )
            except Exception as e:
                return (
                    f"❌ Error loading dataset info: {str(e)}",  # dataset_info_display
                    None  # current_dataset_state
                )
        
        def on_edit_dataset(current_dataset):
            """Handle edit dataset button."""
            try:
                if not current_dataset:
                    return (
                        gr.Row(visible=False),  # edit_section
                        "",  # test_cases_display
                        current_dataset,  # editing_dataset_state
                        "❌ No dataset selected for editing"  # status_message
                    )
                
                # Render test cases HTML
                test_cases_html = dataset_controller.render_test_cases_display(current_dataset)
                
                return (
                    gr.Row(visible=True),  # edit_section
                    test_cases_html,  # test_cases_display
                    current_dataset,  # editing_dataset_state
                    f"✏️ Editing dataset: '{current_dataset.name}'"  # status_message
                )
            except Exception as e:
                return (
                    gr.Row(visible=False),  # edit_section
                    "",  # test_cases_display
                    None,  # editing_dataset_state
                    f"❌ Error accessing dataset for editing: {str(e)}"  # status_message
                )
        
        def on_close_edit():
            """Close the edit section."""
            return (
                gr.Row(visible=False),  # edit_section
                "",  # new_message_text (clear)
                "yellow",  # new_message_classification (reset)
                "✅ Editor closed"  # status_message
            )
        
        def on_show_create():
            """Show create dataset section."""
            return (
                gr.Row(visible=True),  # create_section
                "➕ Create a new dataset"  # status_message
            )
        
        def on_cancel_create():
            """Cancel create dataset."""
            return (
                gr.Row(visible=False),  # create_section
                "",  # new_dataset_name (clear)
                "",  # new_dataset_description (clear)
                "",  # template_selector (reset)
                "❌ Dataset creation cancelled"  # status_message
            )
        
        def on_create_dataset(name, description, template):
            """Create a new dataset."""
            try:
                success, message, new_dataset = dataset_controller.create_new_dataset(
                    name, description, template if template else None
                )
                
                if success:
                    # Refresh dataset choices
                    updated_choices, _, _, _ = dataset_controller.initialize_interface()
                    
                    return (
                        gr.Row(visible=False),  # create_section
                        gr.update(choices=updated_choices, value=f"{new_dataset.name} ({new_dataset.message_count} messages)"),  # dataset_selector
                        "",  # new_dataset_name (clear)
                        "",  # new_dataset_description (clear)
                        "",  # template_selector (reset)
                        new_dataset,  # current_dataset_state
                        message  # status_message
                    )
                else:
                    return (
                        gr.Row(visible=True),  # create_section (keep visible)
                        gr.update(),  # dataset_selector (no change)
                        name,  # new_dataset_name (keep)
                        description,  # new_dataset_description (keep)
                        template,  # template_selector (keep)
                        None,  # current_dataset_state
                        message  # status_message
                    )
            except Exception as e:
                return (
                    gr.Row(visible=True),  # create_section (keep visible)
                    gr.update(),  # dataset_selector (no change)
                    name,  # new_dataset_name (keep)
                    description,  # new_dataset_description (keep)
                    template,  # template_selector (keep)
                    None,  # current_dataset_state
                    f"❌ Error creating dataset: {str(e)}"  # status_message
                )
        
        def on_add_test_case(editing_dataset, message_text, classification):
            """Add a new test case to the dataset."""
            try:
                if not editing_dataset:
                    return (
                        "",  # test_cases_display
                        editing_dataset,  # editing_dataset_state
                        message_text,  # new_message_text (keep)
                        classification,  # new_message_classification (keep)
                        "❌ No dataset being edited"  # status_message
                    )
                
                if not message_text or not message_text.strip():
                    return (
                        dataset_controller.render_test_cases_display(editing_dataset),  # test_cases_display
                        editing_dataset,  # editing_dataset_state
                        message_text,  # new_message_text (keep)
                        classification,  # new_message_classification (keep)
                        "❌ Message text is required"  # status_message
                    )
                
                # Add test case
                success, message, updated_dataset = dataset_controller.add_test_case(
                    editing_dataset, message_text.strip(), classification
                )
                
                if success:
                    # Render updated test cases
                    test_cases_html = dataset_controller.render_test_cases_display(updated_dataset)
                    return (
                        test_cases_html,  # test_cases_display
                        updated_dataset,  # editing_dataset_state
                        "",  # new_message_text (clear)
                        "yellow",  # new_message_classification (reset)
                        message  # status_message
                    )
                else:
                    return (
                        dataset_controller.render_test_cases_display(editing_dataset),  # test_cases_display
                        editing_dataset,  # editing_dataset_state
                        message_text,  # new_message_text (keep)
                        classification,  # new_message_classification (keep)
                        message  # status_message
                    )
            except Exception as e:
                return (
                    "",  # test_cases_display
                    editing_dataset,  # editing_dataset_state
                    message_text,  # new_message_text (keep)
                    classification,  # new_message_classification (keep)
                    f"❌ Error adding test case: {str(e)}"  # status_message
                )
        
        def on_clear_form():
            """Clear the add test case form."""
            return (
                "",  # new_message_text
                "yellow",  # new_message_classification
                "🗑️ Form cleared"  # status_message
            )
        
        def on_save_dataset(editing_dataset):
            """Save the current dataset."""
            try:
                if not editing_dataset:
                    return "❌ No dataset being edited"
                
                success, message = dataset_controller.save_dataset(editing_dataset)
                return message
            except Exception as e:
                return f"❌ Error saving dataset: {str(e)}"
        
        # Bind event handlers
        dataset_selector.change(
            on_dataset_selection_change,
            inputs=[dataset_selector],
            outputs=[dataset_info_display, current_dataset_state]
        )
        
        edit_dataset_btn.click(
            on_edit_dataset,
            inputs=[current_dataset_state],
            outputs=[edit_section, test_cases_display, editing_dataset_state, status_message]
        )
        
        create_dataset_btn.click(
            on_show_create,
            outputs=[create_section, status_message]
        )
        
        create_cancel_btn.click(
            on_cancel_create,
            outputs=[create_section, new_dataset_name, new_dataset_description, template_selector, status_message]
        )
        
        create_confirm_btn.click(
            on_create_dataset,
            inputs=[new_dataset_name, new_dataset_description, template_selector],
            outputs=[create_section, dataset_selector, new_dataset_name, new_dataset_description, template_selector, current_dataset_state, status_message]
        )
        
        close_edit_btn.click(
            on_close_edit,
            outputs=[edit_section, new_message_text, new_message_classification, status_message]
        )
        
        add_test_case_btn.click(
            on_add_test_case,
            inputs=[editing_dataset_state, new_message_text, new_message_classification],
            outputs=[test_cases_display, editing_dataset_state, new_message_text, new_message_classification, status_message]
        )
        
        clear_form_btn.click(
            on_clear_form,
            outputs=[new_message_text, new_message_classification, status_message]
        )
        
        save_dataset_btn.click(
            on_save_dataset,
            inputs=[editing_dataset_state],
            outputs=[status_message]
        )
    
    return interface