import gradio as gr
import gen
import psychohistory

with gr.Blocks(title="PSYCHOHISTORY") as app:
    with gr.Tab("Search"):
        with gr.Row():
            txt_search = gr.Textbox(value="Peace in Eurasia in 2027", label="Search Term", scale=5)
            btn_search = gr.Button("Look", scale=1)
        with gr.Row():
            mem_results = gr.JSON(label="Results")
            btn_search.click(
                gen.generate,
                inputs=[txt_search],
                outputs=mem_results
            )

        with gr.Row():
            img_output = gr.Image(label="Graph Visualization", type="filepath")  # Add an Image component

            # Trigger graph generation after JSON is generated
            mem_results.change(
                psychohistory.main,
                inputs=[mem_results],
                outputs=img_output
            )

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