import gradio as gr
from agent import CodingAgent

agent = CodingAgent()

def handle_query(message, history):
    response = agent.answer(message)
    history.append((message, response))
    return history, ""

def upload_file(file):
    return agent.ingest_file(file.name)

def clear_memory():
    return agent.clear_context()

def get_info():
    return agent.get_stats()

with gr.Blocks(title="LLaMA-3 Coding Agent") as demo:
    gr.Markdown("# 🦙 TinyLlama Coding Agent\nSupports code Q&A + PDF / code file RAG")

    with gr.Tab("Chat"):
        chatbot = gr.Chatbot()
        with gr.Row():
            msg = gr.Textbox(placeholder="Ask a coding question")
            send = gr.Button("Send")
        send.click(handle_query, [msg, chatbot], [chatbot, msg])
        msg.submit(handle_query, [msg, chatbot], [chatbot, msg])

    with gr.Tab("Upload PDF /code file"):
        file_input = gr.File(label="Upload PDF or Python File", file_types=[".pdf", ".py"])
        upload_btn = gr.Button("Upload")
        output = gr.Textbox()
        upload_btn.click(upload_file, file_input, output)

    with gr.Tab("System"):
        info_btn = gr.Button("Get Info")
        clear_btn = gr.Button("Clear Memory")
        info_box = gr.Textbox()
        status_box = gr.Textbox()
        info_btn.click(get_info, outputs=info_box)
        clear_btn.click(clear_memory, outputs=status_box)

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