# app.py
import gradio as gr
import asyncio
from AI_Agent.coordinator import Coordinator

# Initialize coordinator
coord = Coordinator()

# Async function to run the entire pipeline
async def run_brief(brief):
    result = await coord.run_task(brief)
    return (
        result["decomposition"]["tasks_text"],       # Decomposed tasks
        result["assignment"]["assigned_tasks_text"], # Assigned tasks (Frontend/Backend)
        "\n".join(result["retrieval"]["contexts"]),  # Retrieved contexts
        result["reasoning"]["reasoning"],            # Reasoning
        result["synthesis"]["answer"]                # Final synthesis
    )

# Wrapper to run async function synchronously for Gradio
def run_brief_sync(brief):
    return asyncio.run(run_brief(brief))

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Multi-Chain AI Project Coordinator (CPU-friendly Gemma 3n)")
    
    brief_input = gr.Textbox(
        label="Project Brief",
        placeholder="Enter your project brief here",
        lines=2
    )
    
    decomposition_output = gr.Textbox(label="Decomposed Tasks")
    assignment_output = gr.Textbox(label="Assigned Tasks (Frontend/Backend)")
    retrieval_output = gr.Textbox(label="Retrieved Contexts")
    reasoning_output = gr.Textbox(label="Reasoning")
    synthesis_output = gr.Textbox(label="Synthesis / Final Answer")

    run_button = gr.Button("Run Pipeline")
    run_button.click(
        run_brief_sync,
        inputs=[brief_input],
        outputs=[
            decomposition_output,
            assignment_output,
            retrieval_output,
            reasoning_output,
            synthesis_output
        ]
    )

demo.launch()