import gradio as gr

# ---- IMPORT BACKENDS ----
from image_backend import predict_image_pil
from report_generator import generate_report


# =========================
# IMAGE LOGIC
# =========================
def analyze_image(image):
    if image is None:
        return "", "", "", None, '<div class="status">Status: Idle</div>'

    label, confidence, heatmap = predict_image_pil(image)

    # Risk classification
    if label == "Fake":
        if confidence >= 90:
            risk = "high"
            message = "High likelihood of deepfake"
        elif confidence >= 60:
            risk = "warning"
            message = "Possibly deepfake"
        else:
            risk = "neutral"
            message = "Uncertain deepfake"
    else:
        if confidence >= 90:
            risk = "real"
            message = "Likely real"
        elif confidence >= 60:
            risk = "warning"
            message = "Possibly real"
        else:
            risk = "neutral"
            message = "Uncertain - review needed"

    risk_html = f"""
    <div class="risk-card {risk}">
        <div class="risk-title">{label}</div>
        <div class="risk-msg">{message}</div>
    </div>
    """

    return (
        label,
        f"{confidence} %",
        risk_html,
        heatmap,
        '<div class="status">Status: Completed</div>'
    )


# =========================
# CSS (UPDATED FOR DOWNLOAD FIX)
# =========================
css = """
body { background-color: #0f172a; }

.header {
    text-align: center;
    padding: 12px;
    font-size: 28px;
    font-weight: 600;
    color: white;
}

.section {
    color: #cbd5f5;
    margin-bottom: 10px;
}

.gr-box {
    border-radius: 12px !important;
    background: #1e293b !important;
    padding: 15px !important;
}

.risk-card {
    padding: 15px;
    border-radius: 10px;
    color: white;
    font-weight: bold;
}

.risk-title {
    font-size: 18px;
    margin-bottom: 5px;
}

.risk-msg {
    font-size: 14px;
    opacity: 0.9;
}

.risk-card.real { background: #16a34a; }
.risk-card.high { background: #dc2626; }
.risk-card.warning { background: #f59e0b; }
.risk-card.neutral { background: #64748b; }

.status {
    padding: 6px 12px;
    border-radius: 20px;
    background: #334155;
    color: white;
    display: inline-block;
}


"""


# =========================
# UI
# =========================
with gr.Blocks(css=css) as demo:

    # HEADER
    gr.Markdown('<div class="header">AI Driven Deepfake Detection Dashboard</div>')

    # SYSTEM OVERVIEW
    gr.Markdown("""
    ### 🔍 System Overview
    This system detects whether an uploaded image is **real or AI-generated (deepfake)** 
    using deep learning-based image forensics techniques.
    """)

    # MODEL INFO
    gr.Markdown("""
    ### 🧠 Model Details
    - Vision Transformer based architecture  
    - Learns fine-grained facial artifacts  
    - Uses attention for explainability  
    """)

    # STATUS
    status = gr.HTML('<div class="status">Status: Idle</div>')

    # MAIN UI
    with gr.Row():

        # INPUT PANEL
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Image")
            image_input = gr.Image(type="pil", height=300)

            img_submit = gr.Button("Analyze", variant="primary")
            img_clear = gr.Button("Reset")

        # OUTPUT PANEL
        with gr.Column(scale=2):
            gr.Markdown("### 📊 Analysis Results")

            img_pred = gr.Text(label="Prediction")
            img_conf = gr.Text(label="Confidence")

            img_risk = gr.HTML()

            img_heatmap = gr.Image(
                label="Explainability Heatmap",
                height=300,
                interactive=False
            )

            # REPORT FEATURE
            generate_btn = gr.Button("Generate Report")
            report_file = gr.File(label="Download Report")

    # INTERPRETATION GUIDE
    gr.Markdown("""
    ### 📖 How to Interpret Results
    - **Prediction** → Final classification (Real / Fake)  
    - **Confidence** → Model certainty score  
    - **Heatmap** → Highlights regions influencing the decision  
    - **Risk Level**:
      - 🔴 High → Strong deepfake indication  
      - 🟡 Warning → Possible manipulation  
      - ⚪ Neutral → Uncertain (manual review required)  
      - 🟢 Real → Likely authentic image  
    """)

    # LIMITATIONS
    gr.Markdown("""
    ### ⚠️ Limitations
    - Performance may drop on **low-resolution or heavily compressed images**
    - May struggle with **high-quality GAN-generated content**
    - Works best on **face-centric images**
    - Not a replacement for human forensic analysis
    """)

    # PRIVACY
    gr.Markdown("""
    ### 🔐 Privacy & Usage
    - Images are processed temporarily and not stored
    - Intended for **educational and research purposes**
    - Should be used as a **decision-support tool only**
    """)

    # =========================
    # EVENTS
    # =========================

    # Analyze
    img_submit.click(
        lambda img: ("", "", "", None, '<div class="status">Status: Processing...</div>'),
        inputs=image_input,
        outputs=[img_pred, img_conf, img_risk, img_heatmap, status]
    ).then(
        analyze_image,
        inputs=image_input,
        outputs=[img_pred, img_conf, img_risk, img_heatmap, status]
    )

    # Reset
    img_clear.click(
        lambda: (None, "", "", "", None, '<div class="status">Status: Idle</div>', None),
        None,
        [image_input, img_pred, img_conf, img_risk, img_heatmap, status, report_file]
    )

        # ✅ Generate Report
    generate_btn.click(
        generate_report,
        inputs=[img_pred, img_conf, image_input, img_heatmap],
        outputs=report_file
    )


demo.launch()