import gradio as gr
from utils.generation_streaming import generate_response_stream

with gr.Blocks(title="Policy Assistant") as demo:
    gr.Markdown("### ⚡ Kenya Policy QA – Verbatim, Sentiment, and Coherence")

    with gr.Row():
        input_box = gr.Textbox(
            label="Enter your policy question",
            lines=2,
            placeholder="e.g., What are the objectives of Kenya’s energy policies?"
        )

    with gr.Row():
        sentiment_toggle = gr.Checkbox(label="Enable Sentiment Analysis", value=True)
        coherence_toggle = gr.Checkbox(label="Enable Coherence Check", value=True)

    output_box = gr.Textbox(label="LLM Response", lines=25, interactive=False)
    # output_box = gr.Textbox(label="LLM Response", lines=25, interactive=False, stream=True)

    run_btn = gr.Button("🔍 Generate")

    run_btn.click(
        fn=generate_response_stream,
        inputs=[input_box, sentiment_toggle, coherence_toggle],
        outputs=output_box
    )

    input_box.submit(
        fn=generate_response_stream,
        inputs=[input_box, sentiment_toggle, coherence_toggle],
        outputs=output_box
    )

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