"""
StackNet Demo
"""

import gradio as gr

# Test import of our modules
try:
    from src.ui.tabs import create_all_tabs
    from src.ui.handlers import Handlers
    MODULES_OK = True
except Exception as e:
    MODULES_OK = False
    IMPORT_ERROR = str(e)


def create_demo():
    with gr.Blocks(title="StackNet Demo") as demo:
        gr.Markdown("# StackNet Demo 1:1 Preview")

        if not MODULES_OK:
            gr.Markdown(f"**Import Error:** {IMPORT_ERROR}")
            return demo

        with gr.Accordion("Settings", open=True):
            api_key = gr.Textbox(
                label="StackNet Key",
                placeholder="Enter your key to enable generation...",
                type="password",
                value=""
            )

        with gr.Tabs():
            tabs = create_all_tabs()

        # Collect all buttons to disable when no API key
        all_buttons = [
            tabs["text_to_music"]["generate_btn"],
            tabs["music_to_music"]["cover_btn"],
            tabs["music_to_music"]["stems_btn"],
            tabs["text_to_image"]["generate_btn"],
            tabs["image_to_image"]["edit_btn"],
            tabs["text_to_video"]["generate_btn"],
            tabs["image_to_video"]["animate_btn"],
        ]

        def update_buttons(key):
            enabled = bool(key and key.strip())
            return [gr.update(interactive=enabled) for _ in all_buttons]

        api_key.change(
            fn=update_buttons,
            inputs=[api_key],
            outputs=all_buttons
        )

        # Initially disable all buttons
        demo.load(
            fn=lambda: [gr.update(interactive=False) for _ in all_buttons],
            outputs=all_buttons
        )

        # Text to Music
        ttm = tabs["text_to_music"]
        ttm["generate_btn"].click(
            fn=Handlers.generate_music,
            inputs=[
                ttm["prompt"],
                ttm["tags"],
                ttm["instrumental"],
                ttm["lyrics"],
                ttm["title"],
                api_key
            ],
            outputs=[ttm["output_audio"], ttm["status"]]
        )

        # Music to Music - Cover
        mtm = tabs["music_to_music"]
        mtm["cover_btn"].click(
            fn=Handlers.create_cover,
            inputs=[
                mtm["cover_audio_input"],
                mtm["cover_style_prompt"],
                mtm["cover_tags"],
                mtm["cover_title"],
                api_key
            ],
            outputs=[mtm["cover_output"], mtm["cover_status"]]
        )

        # Music to Music - Stems
        mtm["stems_btn"].click(
            fn=Handlers.extract_stems,
            inputs=[mtm["stems_audio_input"], api_key],
            outputs=[
                mtm["vocals_output"],
                mtm["drums_output"],
                mtm["bass_output"],
                mtm["other_output"],
                mtm["stems_status"]
            ]
        )

        # Text to Image
        tti = tabs["text_to_image"]
        tti["generate_btn"].click(
            fn=Handlers.generate_image,
            inputs=[
                tti["prompt"],
                tti["style"],
                tti["aspect_ratio"],
                api_key
            ],
            outputs=[tti["output_image"], tti["status"]]
        )

        # Image to Image
        iti = tabs["image_to_image"]
        iti["edit_btn"].click(
            fn=Handlers.edit_image,
            inputs=[
                iti["input_image"],
                iti["edit_prompt"],
                iti["strength"],
                api_key
            ],
            outputs=[iti["output_image"], iti["status"]]
        )

        # Text to Video
        ttv = tabs["text_to_video"]
        ttv["generate_btn"].click(
            fn=Handlers.generate_video,
            inputs=[
                ttv["prompt"],
                ttv["duration"],
                ttv["style"],
                api_key
            ],
            outputs=[ttv["output_video"], ttv["status"]]
        )

        # Image to Video
        itv = tabs["image_to_video"]
        itv["animate_btn"].click(
            fn=Handlers.animate_image,
            inputs=[
                itv["input_image"],
                itv["motion_prompt"],
                itv["duration"],
                api_key
            ],
            outputs=[itv["output_video"], itv["status"]]
        )

    return demo


demo = create_demo()

if __name__ == "__main__":
    demo.launch()