# app.py

import gradio as gr

# IMPORT ALL YOUR TOOLS HERE
# This is essential so the @tool decorators run and populate the registry
from tools import echo_text, get_current_time, get_current_weather
# Import other tools as you create them: import tools.another_tool

from utils.tool_manager import get_tool_registry, format_docstring_as_markdown

# Get the registry after all tool files have been imported
# This registry initially contains metadata (name, func, ui_builder)
TOOL_REGISTRY = get_tool_registry()

# Prepare dropdown choices using (display_name, internal_name) tuples
# This makes it easy to get the internal name when an item is selected
TOOL_DROPDOWN_CHOICES = [(tool_info['name'], internal_name) for internal_name, tool_info in TOOL_REGISTRY.items()]
DEFAULT_TOOL_INTERNAL_NAME = TOOL_DROPDOWN_CHOICES[0][1] if TOOL_DROPDOWN_CHOICES else None

# --- Gradio App Layout ---
with gr.Blocks(theme="JohnSmith9982/small_and_pretty") as app: # Apply custom CSS here
    gr.Markdown("# LeafCat MCP Multi Tools") # Added a class for styling

    # Tool Selection Dropdown
    tool_dropdown = gr.Dropdown(
        choices=TOOL_DROPDOWN_CHOICES, # Use (display_name, internal_name) tuples
        value=DEFAULT_TOOL_INTERNAL_NAME, # Default value is the internal name
        label="Select a tool",
        interactive=True,
        #elem_id="component-dropdown" # Added ID for potential specific styling
    )

    # Markdown component to display tool docstring
    tool_docstring = gr.Markdown("Tool's docs will be here.", min_height=200)

    # Iterate through the registered tools and build their UIs
    # and store the built components back in the registry
    # All UIs are built initially, but their containing gr.Group is hidden by the ui_builder
    for internal_name, tool_info in TOOL_REGISTRY.items():
        ui_builder = tool_info['ui_builder']
        tool_func = tool_info['func']

        # Build the UI components using the tool's ui_builder function
        # The ui_builder MUST return (ui_group, inputs, outputs, button)
        ui_group, inputs, outputs, button = ui_builder()

        # Store references to the built components directly in the registry entry
        # This makes them accessible for wiring events and visibility updates
        tool_info['built_ui'] = {
            'group': ui_group,
            'inputs': inputs,
            'outputs': outputs,
            'button': button
        }

        # Wire up the button click event for this tool
        # Use the actual tool function as the handler
        # Use api_name corresponding to the tool function's name
        button.click(
            fn=tool_func,
            inputs=inputs,
            outputs=outputs,
            api_name=tool_func.__name__ # Required by the user
        )

    # Helper to get the list of all built UI gr.Group components
    # This is needed as the outputs for the change/load events
    def get_built_ui_groups():
         # The order here must match the order of components provided to the outputs parameter
         # Iterating through TOOL_REGISTRY.values() should give consistent order
         return [tool_info['built_ui']['group'] for tool_info in TOOL_REGISTRY.values()]

    # Function to handle tool selection change
    # It receives the internal_name of the selected tool
    def on_tool_select(selected_tool_internal_name):
        """
        Handles the tool selection dropdown change event.
        Returns a list of gr.update objects to control visibility of tool UIs
        and update the docstring display.
        """
        updates = []
        selected_tool_info = TOOL_REGISTRY[selected_tool_internal_name]

        # Update visibility of tool UI groups
        for internal_name_in_loop, tool_info_in_loop in TOOL_REGISTRY.items():
            ui_group = tool_info_in_loop['built_ui']['group']
            # Set visible to True only for the selected tool's group
            updates.append(gr.update(visible=(internal_name_in_loop == selected_tool_internal_name)))

        # Update the docstring display
        docstring_content = format_docstring_as_markdown(selected_tool_info['func'].__doc__)
        updates.append(gr.update(value=docstring_content))

        return updates

    # Link the dropdown change event to the function that updates UI visibility and docstring
    # The outputs list now includes all tool groups AND the docstring component
    outputs_to_update = get_built_ui_groups() + [tool_docstring]

    tool_dropdown.change(
        fn=on_tool_select,
        inputs=tool_dropdown,
        outputs=outputs_to_update, # Update all groups and the docstring
        api_name=False # Do not expose this function as an API
    )

    # Initial UI setup: Trigger the on_tool_select for the default selected tool
    # This makes the first tool's UI and its docstring visible on load
    app.load(
        fn=on_tool_select,
        inputs=tool_dropdown, # Pass the initial value of the dropdown (internal name)
        outputs=outputs_to_update, # Update all groups and the docstring
        api_name=False # Do not expose this function as an API
    )


# Run the app
if __name__ == "__main__":
    app.launch(mcp_server=True, server_name="0.0.0.0", server_port=7860)