import gradio as gr
from main_agent import agent
import os

if not os.path.exists("uploads"):
    os.makedirs("uploads")


def chatbot_logic(user_input, pdf_file=None, image_file=None):
    context = ""

    if pdf_file:
        if isinstance(pdf_file, str):
            pdf_path = pdf_file
        else:  # Otherwise, it's a file-like object
            pdf_path = os.path.join("uploads", pdf_file.name)
            with open(pdf_path, "wb") as f:
                f.write(pdf_file.read())
        context += f"pdf_path: {pdf_path}\n"


    if image_file:
        if isinstance(image_file, str):
            image_path = image_file
        else:
            image_path = os.path.join("uploads", image_file.name)
            with open(image_path, "wb") as f:
                f.write(image_file.read())
        context += f"image_path: {image_path}\n"


    # Let the agent decide which tool to use
    agent_prompt = f"{context}\nUser: {user_input}"
    result = agent.run(agent_prompt)
    return result


with gr.Blocks(theme=gr.themes.Base(), css=""".gradio-container { background-color: black !important; color: white; }""") as demo:
    gr.Markdown("""
    # 🤖✨ Scientific Paper Assistant
    Upload a PDF or image, and ask *anything* — let the AI do the rest!
    """, elem_id="title")

    with gr.Row():
        chatbot = gr.Chatbot(label="Chat with your AI Paper Assistant!", height=500)

    with gr.Row():
        pdf_upload = gr.File(label="📄 Upload your PDF (optional)")
        image_upload = gr.File(label="🖼️ Upload an image (optional)")

    with gr.Row():
        user_input = gr.Textbox(label="Your message", placeholder="Ask me to do analysis, mind map, data viz, web search...")

    send_btn = gr.Button("🚀 Send")

    def process_chat(user_input, pdf_file, image_file, history):
        response = chatbot_logic(user_input, pdf_file, image_file)
        history.append((user_input, response))
        return history, ""

    send_btn.click(
        fn=process_chat,
        inputs=[user_input, pdf_upload, image_upload, chatbot],
        outputs=[chatbot, user_input]
    )


    gr.HTML("""
    <style>
    #title { text-align: center; font-size: 28px; color: #6a1b9a; font-weight: bold; margin-bottom: 20px; }
    .gradio-container { background: #f3f1f5; border-radius: 12px; padding: 16px; box-shadow: 0 0 10px rgba(0,0,0,0.1);}
    .gradio-container .gradio-row { margin-bottom: 16px; }
    .gr-button { background: #6a1b9a; color: #fff; }
    .gr-button:hover { background: #4a148c; }
    </style>
    """)

demo.launch()