app.py

import os
from typing import Dict, List, Optional, Tuple
import logging
import asyncio

import gradio as gr

from config import (
    AVAILABLE_MODELS, DEMO_LIST, HTML_SYSTEM_PROMPT,
)
from api_clients import generation_code, tavily_client
from chat_processing import (
    clear_history, history_to_chatbot_messages, update_image_input_visibility, update_submit_button,
    get_gradio_language, send_to_sandbox
)
from file_processing import create_multimodal_message
from web_extraction import enhance_query_with_search, perform_web_search
from ux_components import create_top_demo_cards

# --- Gradio App UI ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Update the Gradio theme to use modern styling
theme = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.blue,  # Keeping it consistent
    neutral_hue=gr.themes.colors.neutral,
    font=gr.themes.GoogleFont("Inter"),
    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
)

def demo_card_click(evt: gr.EventData):
    """Handles clicks on the demo cards to populate the input textbox."""
    # This function seems complex for its purpose. Gradio's event data can be tricky.
    # A simpler approach might be to directly pass the description.
    # For now, keeping the logic but adding comments.
    try:
        # Get the index from the event data
        if hasattr(e, '_data') and e._data:
            # Try different ways to get the index
            if 'index' in e._data:
                index = e._data['index']
            elif 'component' in e._data and 'index' in e._data['component']:
                index = e._data['component']['index']
            elif 'target' in e._data and 'index' in e._data['target']:
                index = e._data['target']['index']
            else:
                # If we can't get the index, try to extract it from the card data
                index = 0
        else:
            index = 0
        
        # Ensure index is within bounds
        if index >= len(DEMO_LIST):
            index = 0
            
        return DEMO_LIST[index]['description']
    except (KeyError, IndexError, AttributeError) as e:
        # Return the first demo description as fallback
        return DEMO_LIST[0]['description']


with gr.Blocks(
    theme=gr.themes.Base(
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.blue,
        neutral_hue=gr.themes.colors.neutral,
        font=gr.themes.GoogleFont("Inter"),
        font_mono=gr.themes.GoogleFont("JetBrains Mono"),
        text_size=gr.themes.sizes.text_lg,
        spacing_size=gr.themes.sizes.spacing_lg,
        radius_size=gr.themes.sizes.radius_lg,
    ),
    title="AnyCoder - AI Code Generator",
    css="""
    .col-header { text-transform: uppercase; font-size: 0.9em; letter-spacing: 0.05em; color: #555; border-bottom: 1px solid #eee; padding-bottom: 0.5em; margin-bottom: 1em; }
    .col-header + div { margin-top: 0; } /* Remove extra spacing */
    .gradio-container .gr-button-primary { background: linear-gradient(to right, #42a5f5, #4791db); color: white; } /* Blue gradient */
    .gradio-container .gr-button-secondary { background-color: #e3f2fd; color: #1e88e5; } /* Lighter blue */
    .gradio-container .gr-button-small { font-size: 0.875rem; padding: 0.5rem 1rem; }  /* Smaller buttons*/
    .gradio-container .gr-textbox, .gradio-container .gr-dropdown { border-radius: 6px; }
    .gradio-container .tabs { border-bottom: 2px solid #e0f2f7; } /* Tabs container line */
    .gradio-container .tabitem { padding: 1rem; }  /* Tab content padding*/
    .gradio-container .gr-code { border-radius: 6px; background-color: #f5f5f5; } /* Code background */
    """
) as demo:
    gr.HTML("<h1 align='center'>Shasha - AI Code Generator</h1>")
    history = gr.State([])
    setting = gr.State({
        "system": HTML_SYSTEM_PROMPT,
    })
    current_model = gr.State(AVAILABLE_MODELS[0])  # Moonshot Kimi-K2
    open_panel = gr.State(None)
    last_login_state = gr.State(None)

    with gr.Row(equal_height=False):  # Ensure columns can have different heights
        with gr.Column(scale=1):
            gr.Markdown("## Controls", elem_classes=["col-header"])
            input = gr.Textbox(
                label="What would you like to build?",
                placeholder="Describe your application...",
                lines=4,
                interactive=True
            )
            # Language dropdown for code generation
            language_choices = [
                "python", "c", "cpp", "markdown", "latex", "json", "html", "css", "javascript", "jinja2", "typescript", "yaml", "dockerfile", "shell", "r", "sql"
            ]
            language_dropdown = gr.Dropdown(
                choices=language_choices,
                value="html",
                label="Code Language"
            )
            website_url_input = gr.Textbox(
                label="Website URL for redesign",
                placeholder="https://example.com",
                lines=1
            )
            file_input = gr.File(
                label="Reference file",
                file_types=[".pdf", ".txt", ".md", ".csv", ".docx", ".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".gif", ".webp"]
            )
            image_input = gr.Image(
                label="UI design image",
                visible=False  # Hidden by default; shown only for ERNIE-VL or GLM-VL
            )
            with gr.Row(equal_height=True):  #Keep buttons vertically aligned
                btn = gr.Button("Generate", variant="primary", size="lg", scale=2, interactive=False,)
                clear_btn = gr.Button("Clear", variant="secondary", size="sm", scale=1)
            
            gr.Markdown("---",)
            gr.Markdown("### Quick Examples", elem_classes=["col-header"])
            quick_examples_col = create_top_demo_cards(input)
            
            gr.Markdown("---",)
            gr.Markdown("### Settings", elem_classes=["col-header"])
            search_toggle = gr.Checkbox(
                label="🔍 Web search",
                value=False,
            )
            if not tavily_client:
                gr.Markdown("⚠️ Web search unavailable", visible=True)
            
            model_dropdown = gr.Dropdown(
                choices=[model['name'] for model in AVAILABLE_MODELS],
                value=AVAILABLE_MODELS[0]['name'],  # Moonshot Kimi-K2
                label="Model",
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                system_prompt_input = gr.Textbox(
                    value=HTML_SYSTEM_PROMPT,
                    label="System Prompt",
                    lines=5
                )
                save_prompt_btn = gr.Button("Save Prompt", variant="secondary", size="sm", min_width=100, elem_classes="gr-button-small")

        
        with gr.Column(scale=3):
            model_display = gr.Markdown(f"**Model:** {AVAILABLE_MODELS[0]['name']}", visible=True)
            with gr.Tabs():
                with gr.Tab("Code"):
                    code_output = gr.Code(
                        language="html", 
                        lines=28, 
                        interactive=False,
                        label="Generated Code"
                    )
                with gr.Tab("Preview"):
                    sandbox = gr.HTML(label="Live Preview")
                with gr.Tab("History"):
                    history_output = gr.Chatbot(show_label=False, height=600, type="messages")
    
    # --- Event Handlers ---
    def on_model_change(model_name):
        for m in AVAILABLE_MODELS:
            if m['name'] == model_name:
                return m, f"**Model:** {m['name']}", update_image_input_visibility(m)
        return AVAILABLE_MODELS[0], f"**Model:** {AVAILABLE_MODELS[0]['name']}", update_image_input_visibility(AVAILABLE_MODELS[0])

    def save_prompt(prompt_text):
        return {setting: {"system": prompt_text}}

    def update_code_language(language):
        return gr.update(language=get_gradio_language(language))

    def preview_logic(code, language):
        if language == "html":
            return send_to_sandbox(code)
        else:
            return "<div style='padding:1em;color:#888;text-align:center;'>Preview is only available for HTML.</div>"

    async def submit_query(*args):
        """Handles the main code generation logic asynchronously."""
        async for update in generation_code(*args):
            yield update

    btn.click(fn=submit_query,
              inputs=[input, image_input, file_input, website_url_input, setting, history, current_model, search_toggle, language_dropdown],
              outputs=[code_output, history, sandbox, history_output])

    input.change(update_submit_button, inputs=input, outputs=btn)

    model_dropdown.change(
        on_model_change,
        inputs=model_dropdown,
        outputs=[current_model, model_display, image_input]
    )
    save_prompt_btn.click(save_prompt, inputs=system_prompt_input, outputs=setting)
    # Update preview when code or language changes
    language_dropdown.change(update_code_language, inputs=language_dropdown, outputs=code_output)
    code_output.change(preview_logic, inputs=[code_output, language_dropdown], outputs=sandbox)
    clear_btn.click(clear_history, outputs=[history, history_output, file_input, website_url_input])

if __name__ == "__main__":
    demo.queue(api_open=False, default_concurrency_limit=20).launch(ssr_mode=True, mcp_server=False, show_api=False)