import gradio as gr
from gradio_client import Client, handle_file
import os


# Define your Hugging Face token (make sure to set it as an environment variable)
HF_TOKEN = os.getenv("HF_TOKEN")  # Replace with your actual token if not using an environment variable

# Initialize the Gradio Client for the specified API
client = Client("mangoesai/Elections_Comparison_Agent_V4", hf_token=HF_TOKEN)

# client_name = ['2016 Election','2024 Election', 'Comparison two years']



def stream_chat_with_rag(
    message: str,
    # history: list,
    client_name: str
):
    # print(f"Message: {message}")
    #answer = client.predict(question=question, api_name="/run_graph")
    answer, fig = client.predict(
    	query= message,
		election_year=client_name,
		api_name="/process_query"
    )

    # Debugging: Print the raw response
    print("Raw answer from API:")
    print(answer)


    return answer, fig


    

def heatmap(top_n):
    fig = client.predict(
    	query= message,
		election_year=client_name,
		api_name="/heatmap"
    )
    return fig

# Create Gradio interface
with gr.Blocks(title="Reddit Election Analysis") as demo:
    gr.Markdown("# Reddit Public sentiment & Social topic distribution ")
    with gr.Row():
        with gr.Column():
            with gr.Row():
                top_n = gr.Dropdown(choices= [1,2,3,4,5,6,7,8,9,10])
            with gr.Row():
                fresh_btn = gr.Button("Refresh Heatmap")
        with gr.Column():
            output_heatmap = gr.Plot(
                label="Top Public sentiment & Social topic Heatmap",
                container=True,  # Ensures the plot is contained within its area
                elem_classes="heatmap-plot"  # Add a custom class for styling
            )

    gr.Markdown("# Reddit Election Posts/Comments Analysis")
    gr.Markdown("Ask questions about election-related comments and posts")

    with gr.Row():
        with gr.Column():
            year_selector = gr.Radio(
                choices=["2016 Election", "2024 Election", "Comparison two years"],
                label="Select Election Year",
                value="2016 Election"
            )

            query_input = gr.Textbox(
                label="Your Question",
                placeholder="Ask about election comments or posts..."
            )

            submit_btn = gr.Button("Submit")

            gr.Markdown("""
            ## Example Questions:
            - Is there any comments don't like the election results
            - Summarize the main discussions about voting process
            - What are the common opinions about candidates?
            """)
        with gr.Column():
            output_text = gr.Textbox(
                label="Response",
                lines=20
            )

    with gr.Row():
        output_plot = gr.Plot(
            label="Topic Distribution",
            container=True,  # Ensures the plot is contained within its area
            elem_classes="topic-plot"  # Add a custom class for styling
        )

    # Add custom CSS to ensure proper plot sizing
    gr.HTML("""
        <style>
            .topic-plot {
                min-height: 600px;
                width: 100%;
                margin: auto;
            }
            .heatmap-plot {
                min-height: 400px;
                width: 100%;
                margin: auto;
            }
        </style>
    """)
    fresh_btn.click(
        fn=heatmap,
        inputs=top_n,
        outputs=output_heatmap
    )

    # Update both outputs when submit is clicked
    submit_btn.click(
        fn=stream_chat_with_rag,
        inputs=[query_input, year_selector],
        outputs=[output_text, output_plot]
    )


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