import gradio as gr

# Import UI creation and handler functions from tab modules
from tab_chat import create_chat_tab, handle_chat
from tab_code import create_code_tab
from tab_search import create_search_tab, handle_web_search
from tab_workflow import create_workflow_tab, handle_workflow_chat

MERMAID_SCRIPT = """
<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
"""

# --- Main Gradio UI Definition ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue"), head=MERMAID_SCRIPT) as demo:
    # Global state for the workflow tab
    workflow_topic_state = gr.State("")
    workflow_steps_state = gr.State("*Waiting for task to start...*")

    # --- Header ---
    with gr.Row():
        gr.Markdown("""
        # Ling & Ring Playground(wip)
        ### Experience chat, coding, search, and workflow automation
        """)
    with gr.Row():
        gr.Markdown("""
        [Ling Model Card](https://huggingface.co/inclusionAI/Ling-1T) | [Ring Model Card](https://huggingface.co/inclusionAI/Ring-1T)  
        This application uses API services from [ZenMux](https://zenmux.ai/) and incorporates code generation ideas and prompt snippets from [AnyCoder](https://huggingface.co/spaces/akhaliq/anycoder).
        """)

    # --- Main UI Tabs ---
    with gr.Tabs() as main_ui:
        # Create tabs by calling functions from modules
        with gr.Tab("Chat"):
            chat_components = create_chat_tab()
        with gr.Tab("Code Generation"):
            create_code_tab() # The code tab now handles its own events
        with gr.Tab("Web Search"):
            search_components = create_search_tab()
        with gr.Tab("Workflow"):
            workflow_components = create_workflow_tab()

    # --- Event Handling Logic ---

    # Chat Tab Events
    chat_submit_event = chat_components["chat_input"].submit(
        fn=handle_chat,
        inputs=[
            chat_components["chat_input"],
            chat_components["chatbot"],
            chat_components["system_prompt"],
            chat_components["temperature_slider"],
            chat_components["model_selector"]
        ],
        outputs=[
            chat_components["chatbot"],
            chat_components["chat_input"]
        ]
    )
    chat_components["send_button"].click(
        fn=handle_chat,
        inputs=[
            chat_components["chat_input"],
            chat_components["chatbot"],
            chat_components["system_prompt"],
            chat_components["temperature_slider"],
            chat_components["model_selector"]
        ],
        outputs=[
            chat_components["chatbot"],
            chat_components["chat_input"]
        ]
    )

    # Web Search Tab Events
    search_components["search_button"].click(
        fn=handle_web_search,
        inputs=[search_components["search_input"]],
        outputs=[search_components["search_results_output"]]
    )

    # Workflow Tab Events (Stateful & Incremental with Mermaid)
    def workflow_event_handler(user_input, history, topic, workflow):
        # The handler calls the core logic and returns all the necessary updates
        # for the UI components and the state objects.
        for history, new_topic, new_workflow, mermaid_html, _ in handle_workflow_chat(user_input, history, topic, workflow):
            yield history, new_topic, new_workflow, mermaid_html, new_topic, new_workflow, ""

    workflow_inputs = [
        workflow_components["chat_input"],
        workflow_components["chatbot"],
        workflow_topic_state,
        workflow_steps_state
    ]
    
    workflow_outputs = [
        workflow_components["chatbot"],
        workflow_components["topic_output"],
        workflow_components["workflow_output"],
        workflow_components["mermaid_output"],
        workflow_topic_state,
        workflow_steps_state,
        workflow_components["chat_input"]
    ]

    workflow_js_trigger = """() => { setTimeout(() => { try { mermaid.run({ nodes: document.querySelectorAll('.mermaid') }); } catch (e) { console.error('Mermaid render error:', e); } }, 200); }"""

    workflow_components["send_button"].click(
        fn=workflow_event_handler,
        inputs=workflow_inputs,
        outputs=workflow_outputs,
        js=workflow_js_trigger
    )
    
    workflow_components["chat_input"].submit(
        fn=workflow_event_handler,
        inputs=workflow_inputs,
        outputs=workflow_outputs,
        js=workflow_js_trigger
    )

demo.queue(default_concurrency_limit=10).launch()