ui/components.py

"""
UI components for the Magic Story Creator application
"""

import gradio as gr
from config import constants


def create_header():
    """Create the main header section of the application"""
    return gr.HTML(
        '<div class="image-header-wrapper">'
        '<h1 class="main-header">Magic Story Creator </h1>'
        '<img src="/gradio_api/file=assets/images/wand.png" class="image-next-header" alt="Magic Wand">'
        "</div>"
    )


def create_instructions_accordion():
    """Create the instructions accordion"""
    with gr.Accordion(
        "How to Create Your Magical Story", elem_classes="gr-accordion", open=False
    ) as accordion:
        gr.HTML("""
        <div>
            <div style="display: flex; align-items: center; margin-bottom: 15px;">
                <div style="margin-right: 15px;">
                    <img src="/gradio_api/file=assets/images/wand.png" style="width: 80px; height: 80px;" alt="Magic Wand">
                </div>
                <div>
                    <h3 class="sub-header">Welcome to Magic Story Creator!</h3>
                    <p class="text">This magical tool helps you create personalized stories for children. Follow these simple steps:</p>
                </div>
            </div>
            
            <div style="display: flex; flex-wrap: wrap; gap: 20px;">
                <div class ="box">
                    <h4 class = "box-header">Step 1: Choose Your Adventure</h4>
                    <p>Select a story type and tone that your child will enjoy.</p>
                    <p><strong>Example:</strong> Fantasy with Enthusiastic tone</p>
                </div>
                <div class ="box">
                    <h4 class = "box-header">Step 2: Child's Details</h4>
                    <p>Tell us about the child who will be enjoying the story.</p>
                    <p><strong>Example:</strong> Age 6, English language, interests in dinosaurs and space</p>
                </div>
                
                <div class ="box">
                    <h4 class = "box-header">Step 3: Story Details</h4>
                    <p>Add a subject for your story and how long you'd like it to be.</p>
                    <p><strong>Example:</strong> "Quantum Computing" for 3 minutes of reading time</p>
                </div>

                <div class ="box">
                    <h4 class = "box-header">Step 4: Create Magic!</h4>
                    <p>Click "Create My Magical Story" and watch the magic happen! Once your story appears, click "Create Illustrated Chapters" to add pictures.</p>
                </div>
            </div>
        </div>
        """)
    return accordion


def create_story_tab_header():
    """Create the header for the story tab"""
    return gr.HTML("""
        <div class="image-header-wrapper">
            <img src="/gradio_api/file=assets/images/wizard.png" class="image-next-header" alt="Magic Wand">
            <div class="sub-title" style="margin-right: 15px;">
                Let's create a magical story just for you!
            </div>
        </div>
    """)


def create_step_heading(number, title):
    """Create a step heading"""
    return gr.HTML(f"""
        <div style="text-align: center; margin-bottom: 15px;">
            <span class="step-header">Step {number}: {title}</span>
        </div>
    """)


def create_story_options():
    """Create story type and tone selection options"""
    with gr.Row():
        # Story Type
        with gr.Column(scale=1):
            story_type = gr.Radio(
                choices=constants.STORY_TYPES,
                value=constants.DEFAULT_STORY_TYPE,
                interactive=True,
                label="Story Type",
                elem_id="story-type-buttons",
            )

        # Tone
        with gr.Column(scale=1):
            tone = gr.Radio(
                choices=constants.TONE_TYPES,
                label="Story Tone",
                value=constants.DEFAULT_TONE_TYPE,
                interactive=True,
                elem_id="story-tone-buttons",
            )

    return story_type, tone


def create_reader_details():
    """Create the reader details section"""
    with gr.Row():
        # Kid Age
        with gr.Column(scale=1):
            kid_age = gr.Slider(
                minimum=constants.KID_AGE_MIN,
                maximum=constants.KID_AGE_MAX,
                value=constants.DEFAULT_KID_AGE,
                step=1,
                label="Kid Age",
                elem_id="age-slider",
            )

        # Kid Language
        with gr.Column(scale=1):
            kid_language = gr.Textbox(
                label="Kid Language",
                value=constants.DEFAULT_LANGUAGE,
                elem_id="language-input",
            )

    with gr.Row():
        # Kid Interests with icon
        with gr.Column():
            kid_interests = gr.Textbox(
                label="Kid Interests",
                placeholder="Robots, Space, Dinosaurs...",
                elem_id="interests-input",
            )

    return kid_age, kid_language, kid_interests


def create_story_details():
    """Create the story details section"""
    with gr.Row():
        # Subject
        with gr.Column(scale=1):
            subject = gr.Textbox(
                placeholder="MCP servers, Quantum Computing...",
                label="Subject (What is the story about?)",
                elem_id="subject-input",
            )

        # Reading time
        with gr.Column(scale=1):
            reading_time = gr.Slider(
                minimum=constants.READING_TIME_MIN,
                maximum=constants.READING_TIME_MAX,
                value=constants.DEFAULT_READING_TIME,
                step=1,
                label="Reading Time (minutes)",
                elem_id="reading-time-slider",
            )

    with gr.Row():
        # PDF Upload with icon
        with gr.Column(scale=1):
            pdf_upload = gr.File(
                file_types=[".pdf"],
                label="PDF Upload (Optional)",
                elem_id="pdf-upload",
            )

        # AI Model with icon (hidden behind a collapsible for simplicity)
        with gr.Column(scale=1):
            with gr.Accordion("Advanced: Choose AI Model", open=False):
                model_selector = gr.Dropdown(
                    constants.MODEL_OPTIONS,
                    label="AI Model",
                    value=constants.DEFAULT_MODEL,
                    elem_id="model-selector",
                )

    return subject, reading_time, pdf_upload, model_selector


def create_story_action_buttons():
    """Create the action buttons for story generation"""
    with gr.Row():
        generate_button = gr.Button(
            "✨ Create My Magical Story! ✨",
            variant="primary",
        )
        clear_button = gr.Button(
            "🧹 Start Over",
            variant="secondary",
        )

    return generate_button, clear_button


def create_loading_container():
    """Create the loading animation container"""
    with gr.Row(visible=False) as container:
        gr.HTML("""
            <div class="loading-animation">
                📚 <span style="color: #5d9df5;">Creating your magical story...</span> 📚
            </div>
        """)
    return container


def story_ready_header():
    """Return the HTML for the story ready header"""
    return gr.HTML("""
            <div style="text-align: center; margin-bottom: 15px;">
                <span class="step-header">Your Magical Story is Ready!</span>
            </div>
        """)


def create_story_output_container():
    """Create the container for displaying the generated story"""
    story_title = gr.Textbox(
        label="Story Title",
        interactive=False,
        lines=1,
        elem_id="story-title",
    )

    story_text = gr.Textbox(
        label="Generated Story",
        interactive=False,
        lines=15,
        elem_id="story-text",
    )

    with gr.Row():
        process_chapters_button = gr.Button(
            "🎨 Create Illustrated Chapters! 🎨",
            variant="primary",
            interactive=False,
            elem_classes="chapters-button",
        )

    return story_title, story_text, process_chapters_button


def create_chapter_loading_container():
    """Create the loading animation container for chapters"""
    with gr.Row(visible=False, elem_id="chapter-loading-container") as container:
        gr.HTML("""
            <div id="progress-info"">
                <span id="progress-stage">...</span>
            </div>
        """)
    return container


def create_empty_chapters_placeholder():
    """Create the placeholder for when no chapters are available"""
    return gr.HTML("""
        <div class="image-header-wrapper empty-placeholder" id="empty-chapters-placeholder">
            <img src="/gradio_api/file=assets/images/empty_bk.png" style="width: 150px; margin-bottom: 20px;" alt="Empty Book">
            <h1 class="sub-header" style="margin-left: 16px; padding-bottom: 36px;">
               Your Illustrated Story Awaits!
            </h1>
        </div>
            <p class="text empty-placeholder" style="font-size: 1.2em; text-align: center;">First, create your story and then click the "Create Illustrated Chapters!" button to see your story come to life with pictures!</p>
    """)


def create_chapter_loading_placeholder():
    return gr.HTML("""
                                <div style="text-align: center; padding: 30px; background-color: #f0f7ff; border-radius: 15px;">
                                    <div class="loading-animation" style="font-size: 36px;">✨📚✨</div>
                                    <h3 style="color: #5d9df5; margin-top: 20px;">Creating Your Magical Story Chapters...</h3>
                                    <p style="color: #666; margin-top: 10px;">Our wizard is crafting beautiful chapters and illustrations for your story...</p>
                                    <div class="small-loader" style="text-align: center; padding: 10px;">
                                    <div class="spinner" style="display: inline-block; width: 20px; height: 20px; 
                                        border: 3px solid rgba(0, 0, 0, 0.1); border-radius: 50%; 
                                        border-top-color: #3498db; animation: spin 1s linear infinite;"></div>
                                    <style>
                                        @keyframes spin {
                                            to { transform: rotate(360deg); }
                                        }
                                    </style>
                                </div>
                                </div>
            """)


def create_chapter_error_display(error_message):
    """Create an error display for chapter generation issues"""
    return gr.HTML(f"""
        <div style="text-align: center; padding: 20px; background-color: #ffeded; border-radius: 15px; border: 2px solid #ff5757; margin-top: 20px;">
            <h3 style="color: #ff5757; margin-bottom: 10px;">Oops! Something went wrong</h3>
            <p style="color: #333;">{error_message}</p>
        </div>
    """)


def create_story_title_display(title):
    """Create a display for the story title"""
    return gr.HTML(f"""
        <div style="text-align: center; margin: 20px 0;">
            <h1 class="step-header">
                {title}
            </h1>
        </div>
    """)


def create_chapter_navigation():
    """Create the chapter navigation display"""
    return gr.HTML("""
        <span style="display:flex; gap:8px; font-size: 1.2em;">
            <p style="color: var( --accent-pink-light);margin: 0;">Your Adventure Chapters</p>
            <p class="text"">Click on each chapter to reveal its magical contents!</p>
        </span>
    """)


def create_chapter_accordion(index, chapter):
    """Create an accordion for a story chapter"""
    with gr.Accordion(
        label=f"Chapter {index}: {chapter.get('title', 'Untitled')}",
        open=False,
        elem_classes="gr-accordion",
    ):
        with gr.Blocks(elem_classes="chapter-content-box"):
            # Chapter content with magical styling
            gr.HTML(f"""
                <div style="padding: 10px; background-color: #f8fcff; border-radius: 10px; border-left: 5px solid #5d9df5;">
                    <div style="font-family: 'Comic Sans MS', cursive, sans-serif; font-size: 1.1em; line-height: 1.6; color: #333;">
                        {chapter.get("content", "")}
                    </div>
                </div>
            """)
            # Image prompt
            gr.HTML(f"""
                <div style="margin-top: 15px; padding: 10px; background-color: #fff9f5; border-radius: 10px; border-left: 5px solid #ff9057;">
                    <h4 style="color: #ff9057; margin-top: 0;">Image Description:</h4>
                    <p style="font-style: italic; color: #666;">{chapter.get("image_prompt", "")}</p>
                </div>
            """)

            # Display image if available with nice styling
            image_b64 = chapter.get("image_b64")

            with gr.Row(equal_height=True):
                with gr.Column(scale=1):
                    if image_b64:
                        gr.HTML(f"""
                            <div style="text-align: center;">
                                <div style="display: inline-block; padding: 10px; background-color: white; border-radius: 15px; box-shadow: 0 4px 8px rgba(0,0,0,0.1);">
                                    <img src="data:image/png;base64,{chapter["image_b64"]}" style="border-radius: 10px; max-width: 100%;" alt="Chapter Illustration"/>
                                </div>
                            </div>
                        """)
                    else:
                        gr.HTML("""
                            <div style="text-align: center; padding: 30px; background-color: #f0f7ff; border-radius: 15px;">
                                <div class="loading-animation" style="font-size: 30px;">🖌️</div>
                                <p style="color: #5d9df5; font-weight: bold; margin-top: 15px;">{status}</p>
                            </div>
                        """)

                with gr.Column(scale=1):
                    chapter_content_state = gr.State(value=chapter.get("content", ""))
                    audio_statuss = gr.Markdown("", visible=False, elem_classes="text")

                    read_aloud_btn = gr.Button(
                        f"🔊 Read Chapter {index} Aloud",
                        size="lg",
                        variant="primary",
                        elem_classes="audio-button",
                    )

                    read_aloud_audio = gr.Audio(
                        label="Generated Audio",
                        interactive=False,
                        autoplay=True,
                        show_download_button=True,
                        show_share_button=False,
                        elem_classes="audio-output",
                        waveform_options=gr.WaveformOptions(
                            waveform_color="#ff9057",
                            waveform_progress_color="#ff6a3b",
                            skip_length=2,
                            show_controls=True,
                        ),
                    )

    return read_aloud_btn, read_aloud_audio, chapter_content_state, audio_statuss


def create_story_melody_section():
    """Create the section for generating a story melody"""
    gr.HTML("""
        <div class="image-header-wrapper"  style="margin: 30px 0 15px 0;">
            <img src="/gradio_api/file=assets/images/headphones.png" style="width: 132px; height: 132px; margin-bottom: 20px;" alt="Musical headphone">
            <h1 class="sub-header" style="margin-left: 16px; padding-bottom: 36px;">
                Create a Magical Soundtrack for Your Story!
            </h1>
        </div>
        <p class="text" style="font-size: 1.2em; text-align: center;">Click the button below to generate a special melody that matches the mood of your adventure!</p>
    """)

    with gr.Blocks(elem_classes="melody-box"):
        with gr.Row():
            with gr.Column(scale=1, min_width=200):
                generate_melody_button = gr.Button(
                    "🎵 Create Story Soundtrack 🎵",
                    variant="primary",
                    elem_classes="melody-button",
                )
                melody_status = gr.Markdown("", visible=False)

            with gr.Column(scale=2):
                melody_output = gr.Audio(
                    label="Your Story's Magical Soundtrack",
                    interactive=False,
                    autoplay=True,
                    show_download_button=True,
                    show_share_button=False,
                    waveform_options=gr.WaveformOptions(
                        waveform_color="#9c5fff",
                        waveform_progress_color="#6a2fff",
                        skip_length=2,
                        show_controls=True,
                    ),
                )

    return generate_melody_button, melody_status, melody_output


# TODO split this further later
def create_3d_model_viewer(story: str = "simple yellow duck"):
    """
    Create a 3D model viewer component that displays a mesh from a base64 string.
    For now, this uses a placeholder HTML with a message and a download link for the mesh.
    In a real app, you would use a JS 3D viewer (e.g., Three.js) embedded in the HTML.
    """

    gr.HTML("""
        <div class="image-header-wrapper"  style="margin: 30px 0 15px 0;">
            <h1 class="sub-header" style="margin-left: 16px; padding-bottom: 36px;">
                Create a Magical 3D Model from Your Story!
            </h1>
        </div>
        <p class="text" style="font-size: 1.2em; text-align: center;">Click the button below to generate a 3D model from your story!</p>
    """)

    with gr.Blocks(elem_classes="melody-box"):
        with gr.Row():
            with gr.Column(scale=1, min_width=200):
                generate_model_button = gr.Button(
                    "Create 3D Model 🎨",
                    variant="primary",
                    elem_classes="melody-button",
                )

            with gr.Column(scale=2):
                model_viewer = gr.Model3D(
                    label="3D Model Viewer", clear_color=[0.0, 0.0, 0.0, 0.0]
                )

            model_status = gr.Markdown("", visible=False)

    # Add status message and model display

    return (
        generate_model_button,
        model_viewer,
        model_status,
    )


def create_readme_display(section_titles=None):
    """
    Create a component that displays the project's README file.
    Optionally filter to show only specific sections by title.

    Args:
        section_titles (str or list, optional): If provided, only display these sections of the README.
            Can be a single string (e.g., "Getting Started") or a list of strings
            (e.g., ["Getting Started", "Features", "Installation"]).
    """
    import re
    from pathlib import Path

    gr.set_static_paths(paths=[Path.cwd().absolute() / "/assets/images"])

    # Handle single string for backward compatibility
    if isinstance(section_titles, str):
        section_titles = [section_titles]
    elif section_titles is None:
        section_titles = []

    # Find the README file in the project root directory
    current_dir = Path(__file__).parent
    project_root = current_dir.parent
    readme_paths = [
        project_root / "README.md",
        project_root / "Readme.md",
        project_root / "readme.md",
    ]

    readme_content = ""
    for path in readme_paths:
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                readme_content = f.read()
                break

    if not readme_content:
        readme_content = "README file not found in the project directory."

    # If specific sections are requested, extract them
    if section_titles and readme_content:
        # Pattern to match headers (# Header, ## Header, etc.)
        header_pattern = r"^(#+)\s+(.+)$"

        # Split the content by lines
        lines = readme_content.split("\n")
        all_filtered_content = []

        # Process each requested section
        for section_title in section_titles:
            filtered_content = []
            in_target_section = False
            current_level = 0

            for line in lines:
                # Check if line is a header
                header_match = re.match(header_pattern, line, re.MULTILINE)

                if header_match:
                    header_level = len(header_match.group(1))  # Number of # symbols
                    header_text = header_match.group(2).strip()

                    # If we found the target section
                    if header_text.lower() == section_title.lower():
                        in_target_section = True
                        current_level = header_level
                        filtered_content.append(line)  # Include the section header
                    # If we're in the target section and find another header of same or higher level, exit section
                    elif in_target_section and header_level <= current_level:
                        in_target_section = False
                    # Include subsection headers if we're in the target section
                    elif in_target_section:
                        filtered_content.append(line)
                # Include content if we're in the target section
                elif in_target_section:
                    filtered_content.append(line)

            # Add this section's content to our combined content
            if filtered_content:
                all_filtered_content.extend(filtered_content)
                # Add a separator between sections unless it's the last section
                if section_title != section_titles[-1]:
                    all_filtered_content.append("\n---\n")
            else:
                all_filtered_content.append(
                    f"Section '{section_title}' not found in README.\n"
                )

        # Update the content to only the filtered sections
        if all_filtered_content:
            readme_content = "\n".join(all_filtered_content)
        else:
            readme_content = "None of the requested sections were found in README."

    # Create a header similar to other components
    if not section_titles:
        section_display = "Full Documentation"
    elif len(section_titles) == 1:
        section_display = f"{section_titles[0]} Section"
    else:
        section_display = "Sections: " + ", ".join(section_titles)

    # Display the README content using Markdown
    with gr.Blocks(elem_classes="readme-box"):
        readme_display = gr.Markdown(
            value=readme_content, elem_classes="readme-content"
        )

    return readme_display