import gradio as gr

# ---- IMPORT BACKENDS ----
from image_backend import predict_image_pil
from audio_backend import predict_audio


# =========================
# IMAGE LOGIC (UNCHANGED)
# =========================
def analyze_image(image):
    label, confidence, heatmap = predict_image_pil(image)

    if label == "Fake":
        if confidence >= 90:
            risk = "🚨 High likelihood of Deepfake"
        elif confidence >= 60:
            risk = "⚠️ Possibly Deepfake"
        else:
            risk = "⚠️ Uncertain Deepfake"
    else:
        if confidence >= 90:
            risk = "✅ Likely Real"
        elif confidence >= 60:
            risk = "⚠️ Possibly Real"
        else:
            risk = "⚠️ Uncertain – Needs Review"

    return label, f"{confidence} %", risk, heatmap


# =========================
# AUDIO LOGIC (UNCHANGED)
# =========================
def analyze_audio(audio_path):
    label, confidence = predict_audio(audio_path)

    if label == "fake":
        if confidence >= 90:
            risk = "🚨 High likelihood of Deepfake"
        elif confidence >= 60:
            risk = "⚠️ Possibly Deepfake"
        else:
            risk = "⚠️ Uncertain – Needs Review"
    else:
        if confidence >= 90:
            risk = "✅ Likely Real"
        elif confidence >= 60:
            risk = "⚠️ Possibly Real"
        else:
            risk = "⚠️ Uncertain – Needs Review"

    return label.capitalize(), f"{confidence} %", risk


# =========================
# UI (EXACT SAME)
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Unified Deepfake Detection System")

    with gr.Tabs():

        # HOME TAB
        with gr.Tab("🏠 Home"):
            gr.Markdown(
                """
                ## Welcome 👋  
                Select the type of media you want to analyze:
                """
            )

            gr.Markdown("### 🔍 Choose Detection Mode")
            gr.Markdown("- 🖼 **Image Deepfake Detection**\n- 🎧 **Audio Deepfake Detection**")

            gr.Markdown(
                """
                👉 Use the tabs above to switch between Image and Audio detection.
                """
            )

        # IMAGE TAB
        with gr.Tab("🖼 Image Deepfake"):
            gr.Markdown("# 🖼 Deepfake Image Detection System")

            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        label="Upload Image",
                        type="pil",
                        height=280
                    )
                    img_submit = gr.Button("Submit")
                    img_clear = gr.Button("Clear")

                with gr.Column(scale=2):
                    img_pred = gr.Text(label="Prediction")
                    img_conf = gr.Text(label="Confidence")
                    img_risk = gr.Text(label="Risk Assessment")
                    img_heatmap = gr.Image(
                        label="Explainability Heatmap",
                        height=280
                    )

            img_submit.click(
                fn=analyze_image,
                inputs=image_input,
                outputs=[img_pred, img_conf, img_risk, img_heatmap]
            )

            img_clear.click(
                fn=lambda: (None, "", "", None),
                inputs=None,
                outputs=[image_input, img_pred, img_conf, img_risk]
            )

        # AUDIO TAB
        with gr.Tab("🎧 Audio Deepfake"):
            gr.Markdown("# 🎧 Deepfake Audio Detection System")

            with gr.Row():
                with gr.Column(scale=1):
                    audio_input = gr.Audio(
                        label="Upload Audio (.wav)",
                        type="filepath"
                    )
                    aud_submit = gr.Button("Submit")
                    aud_clear = gr.Button("Clear")

                with gr.Column(scale=2):
                    aud_pred = gr.Text(label="Prediction")
                    aud_conf = gr.Text(label="Confidence")
                    aud_risk = gr.Text(label="Risk Assessment")

            aud_submit.click(
                fn=analyze_audio,
                inputs=audio_input,
                outputs=[aud_pred, aud_conf, aud_risk]
            )

            aud_clear.click(
                fn=lambda: (None, "", ""),
                inputs=None,
                outputs=[audio_input, aud_pred, aud_conf]
            )

demo.launch()