import gradio as gr
from fastapi import FastAPI
from fastapi.responses import JSONResponse

def create_gradio_interface(app: FastAPI, conversational_rag_chain, agent):
    def qa_function(message, history, system):
        if system == "RAG":
            response = conversational_rag_chain.invoke(
                {"input": message},
                config={"configurable": {"session_id": "abc123"}}
            )
            return response["answer"]
        elif system == "Agent":
            response = agent.invoke(
                {"input": message},
                config={"configurable": {"session_id": "agent_session"}}
            )
            return response['output']

    gr_app = gr.Blocks()

    with gr_app:
        gr.Markdown("# NCERT Q&A System")
        gr.Markdown("Ask questions based on the NCERT Sound chapter or use the Agent for broader queries.")

        chatbot = gr.Chatbot()
        msg = gr.Textbox()
        clear = gr.Button("Clear")

        system_choice = gr.Radio(["RAG", "Agent"], label="Choose System", value="RAG")

        def user(user_message, history, system):
            return "", history + [[user_message, None]]

        def bot(history, system):
            user_message = history[-1][0]
            bot_message = qa_function(user_message, history, system)
            history[-1][1] = bot_message
            return history

        msg.submit(user, [msg, chatbot, system_choice], [msg, chatbot], queue=False).then(
            bot, [chatbot, system_choice], chatbot
        )
        clear.click(lambda: None, None, chatbot, queue=False)

    return gr_app