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

import gradio as gr
import requests

BACKEND_URL = "https://adarshds-thesisbackend.hf.space"

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):
    try:
        r = requests.post(
            RAG_URL,
            json={"question": message},
            timeout=120,
        )
        r.raise_for_status()
        data = r.json()

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

        history.append({"role": "user", "content": message})
        history.append({
            "role": "assistant",
            "content": data.get("answer", "⚠️ No answer returned.") + images
        })

        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
# =============================
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, type="messages")
        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
        )


# ✅ Required for Hugging Face Spaces
demo.queue().launch()