app.py

import os
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

import gradio as gr
import requests

# =============================
# BACKEND CONFIG
# =============================

BACKEND_URL = "http://127.0.0.1:8000"

RAG_URL = f"{BACKEND_URL}/rag/ask"
UPLOAD_URL = f"{BACKEND_URL}/upload-pdf/"
VISUALIZE_URL = f"{BACKEND_URL}/visualize/"
GRADE_URL = f"{BACKEND_URL}/grade-annotation/"


# =============================
# 💬 RAG CHAT
# =============================
def ask_question(message, history):

    if history is None:
        history = []

    try:

        r = requests.post(
            RAG_URL,
            json={"question": message},
            timeout=120,
        )

        r.raise_for_status()
        data = r.json()

        answer = data.get("answer", "⚠️ No answer returned.")

        # Render images returned from backend
        images_md = ""

        for img in data.get("images", []):
            images_md += f"\n\n![image]({img})"

        history.append({
            "role": "user",
            "content": message
        })

        history.append({
            "role": "assistant",
            "content": answer + images_md
        })

        return "", history

    except Exception as e:

        history.append({
            "role": "assistant",
            "content": f"❌ Error: {str(e)}"
        })

        return "", history


# =============================
# 📄 UPLOAD PDF
# =============================
def upload_pdf(file):

    if file is None:
        return "⚠️ Please upload a PDF."

    try:

        with open(file.name, "rb") as f:

            r = requests.post(
                UPLOAD_URL,
                files={"file": f},
                timeout=300,
            )

        r.raise_for_status()

        return r.json().get("message", "✅ Uploaded")

    except Exception as e:

        return f"❌ {str(e)}"


# =============================
# 🧠 VISUALIZE ANATOMY
# =============================
def visualize(image, question):

    if image is None:
        return "⚠️ Upload an image.", None

    try:

        with open(image, "rb") as f:

            r = requests.post(
                VISUALIZE_URL,
                files={"image": f},
                data={"question": question},
                timeout=120,
            )

        r.raise_for_status()

        data = r.json()

        output_image_url = None

        if data.get("output_image"):
            output_image_url = f"{BACKEND_URL}/{data['output_image']}"

        return data.get("answer", ""), output_image_url

    except Exception as e:

        return f"❌ {str(e)}", None


# =============================
# 📝 GRADE ANNOTATION
# =============================
def grade(image):

    if image is None:
        return "⚠️ Upload an image."

    try:

        with open(image, "rb") as f:

            r = requests.post(
                GRADE_URL,
                files={"file": f},
                timeout=120,
            )

        r.raise_for_status()

        return r.json().get("result", "No result returned.")

    except Exception as e:

        return f"❌ {str(e)}"


# ==================================================
# 🎨 UI
# ==================================================

with gr.Blocks() as demo:

    gr.Markdown("# 🧠 Multimodal RAG Anatomy Tutor")

    # ================= CHAT =================
    with gr.Tab("💬 Chat"):

        chatbot = gr.Chatbot(height=500)

        msg = gr.Textbox(
            placeholder="Ask a question...",
            show_label=False
        )

        msg.submit(
            ask_question,
            [msg, chatbot],
            [msg, chatbot],
            show_progress=True
        )

    # ================= UPLOAD =================
    with gr.Tab("📄 Upload PDF"):

        file = gr.File(file_types=[".pdf"])

        upload_btn = gr.Button("Upload & Ingest")

        upload_output = gr.Textbox()

        upload_btn.click(
            upload_pdf,
            file,
            upload_output,
            show_progress=True
        )

    # ================= VISUALIZE =================
    with gr.Tab("🧠 Visualize Anatomy"):

        img = gr.Image(type="filepath")

        q = gr.Textbox(label="Ask about the image")

        viz_btn = gr.Button("Visualize")

        viz_answer = gr.Textbox(label="Answer")

        viz_image = gr.Image(label="Annotated Image")

        viz_btn.click(
            visualize,
            [img, q],
            [viz_answer, viz_image],
            show_progress=True
        )

    # ================= GRADE =================
    with gr.Tab("📝 Annotation Grading"):

        grade_img = gr.Image(type="filepath")

        grade_btn = gr.Button("Grade")

        grade_result = gr.Textbox(lines=15)

        grade_btn.click(
            grade,
            grade_img,
            grade_result,
            show_progress=True
        )


# =============================
# 🚀 LAUNCH
# =============================

demo.queue().launch()