import gradio as gr
import os
import sys
import threading

print("--- Desk Agent Starting Up ---")
print("This application requires the following Hugging Face Space secrets to be set:")
print("- JULES_API_KEY")
print("- HF_TOKEN")
print("- BLABLADOR_API_KEY")
print("-----------------------------")

# Ensure the multi_agent_system is in the Python path
sys.path.append(os.getcwd())
from multi_agent_system.orchestrator.agent import main_orchestration_workflow

# --- Streaming UI Updates ---
def stream_orchestrator(msg, history, jules_key, hf_token, blablador_key):
    """
    A generator that handles streaming UI updates from the orchestrator workflow.
    """
    history.append([msg, ""])
    yield history

    # Set API keys from UI or environment variables
    api_keys = {
        "JULES_API_KEY": os.environ.get("JULES_API_KEY", jules_key),
        "HF_TOKEN": os.environ.get("HF_TOKEN", hf_token),
        "BLABLADOR_API_KEY": os.environ.get("BLABLADOR_API_KEY", blablador_key),
    }
    for key, value in api_keys.items():
        if not value:
            history[-1][1] = f"Error: {key} is not set. Please provide it as a Space secret or in the UI."
            yield history
            return
        os.environ[key] = value

    workflow_generator = main_orchestration_workflow()
    final_bot_message = ""

    # Stream updates from the workflow
    for update in workflow_generator:
        final_bot_message += update
        history[-1][1] = final_bot_message
        yield history

    # Clean up environment variables
    for key in api_keys:
        if key in os.environ:
            del os.environ[key]

def disable_ui():
    return gr.Textbox(interactive=False), gr.Button(interactive=False)

def enable_ui():
    return gr.Textbox(interactive=True, placeholder="Workflow complete. You can start a new one."), gr.Button(interactive=True)

# --- Build the Gradio UI ---
with gr.Blocks(title="Desk Agent") as demo:
    gr.Markdown("# Desk Agent Orchestrator")
    gr.Markdown("A conversational interface to a multi-agent workflow system. Press 'Start Workflow' to begin.")

    chatbot = gr.Chatbot(label="Agent Conversation", height=500)
    msg_input = gr.Textbox(label="Your Message", placeholder="Press 'Start Workflow' to begin...")

    with gr.Accordion("API Key Overrides (Optional)", open=False):
        gr.Markdown("API keys will be automatically loaded from Space secrets if available.")
        jules_key_input = gr.Textbox(label="Jules API Key", type="password")
        hf_token_input = gr.Textbox(label="Hugging Face Token", type="password")
        blablador_key_input = gr.Textbox(label="Blablador API Key", type="password")

    start_button = gr.Button("Start Workflow", variant="primary")
    clear_button = gr.ClearButton([msg_input, chatbot])

    # --- Event Handling ---
    start_button.click(
        fn=disable_ui, outputs=[msg_input, start_button]
    ).then(
        fn=stream_orchestrator,
        inputs=[msg_input, chatbot, jules_key_input, hf_token_input, blablador_key_input],
        outputs=[chatbot]
    ).then(
        fn=enable_ui, outputs=[msg_input, start_button]
    )

    msg_input.submit(
        fn=disable_ui, outputs=[msg_input, start_button]
    ).then(
        fn=stream_orchestrator,
        inputs=[msg_input, chatbot, jules_key_input, hf_token_input, blablador_key_input],
        outputs=[chatbot]
    ).then(
        fn=enable_ui, outputs=[msg_input, start_button]
    )


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