"""
AI-Powered Disaster Triage from Images
----------------------------------------
Gradio front-end. Run locally with:  python app.py
Deploys as-is to a Hugging Face Space (ZeroGPU-compatible, see README).
"""
try:
    import spaces
    GPU_DECORATOR = spaces.GPU
except ImportError:
    def GPU_DECORATOR(fn):
        return fn

import json
import tempfile
import os

import gradio as gr



# ---------------------------------------------------------------------------
# Optional: Hugging Face ZeroGPU support.
# On a normal local NVIDIA GPU, `spaces` isn't installed and this no-ops.
# On HF Spaces (ZeroGPU hardware), this decorator grants a GPU burst for the
# duration of the function call — required for Spaces' free GPU tier.
# ---------------------------------------------------------------------------


@GPU_DECORATOR
def analyze_image(image_path):
    # Import AFTER spaces has initialized
    from src.pipeline import run_triage

    if image_path is None:
        return None, "Please upload an image first.", "{}"

    annotated_rgb, result = run_triage(image_path)
    report_md = result.to_markdown()
    report_json = json.dumps(result.__dict__, indent=2)

    return annotated_rgb, report_md, report_json

def download_report(image_path, report_md):
    if not report_md:
        return None
    fd, path = tempfile.mkstemp(suffix=".md", prefix="triage_report_")
    with os.fdopen(fd, "w") as f:
        f.write("# Disaster Triage Report\n\n" + report_md)
    return path


CSS = """
#risk-panel { font-size: 1.05rem; }
footer { visibility: hidden }
"""

with gr.Blocks(title="AI Disaster Triage", css=CSS, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🚨 AI-Powered Disaster Triage from Images
        Upload a photo from a citizen, drone, CCTV feed, or rescue team.
        The system detects objects (YOLO11), reasons about the scene like a
        triage officer (Qwen2.5-VL), and returns a risk score + recommended
        response — not just a list of detected objects.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="filepath", label="Disaster Image", height=380)
            analyze_btn = gr.Button("🔍 Analyze Scene", variant="primary")
            gr.Markdown("*Tip: drone/CCTV/citizen photos of floods, collapses, blocked roads work best.*")

        with gr.Column(scale=1):
            annotated_output = gr.Image(label="Annotated Scene (YOLO11 + Risk Badge)", height=380)

    with gr.Row():
        with gr.Column(scale=2):
            report_output = gr.Markdown(label="Triage Report", elem_id="risk-panel")
            download_btn = gr.Button("⬇️ Download Report (.md)")
            file_output = gr.File(label="Report file", visible=True)
        with gr.Column(scale=1):
            json_output = gr.Code(label="Raw structured output (JSON)", language="json")

    analyze_btn.click(
        fn=analyze_image,
        inputs=[image_input],
        outputs=[annotated_output, report_output, json_output],
    )
    download_btn.click(
        fn=download_report,
        inputs=[image_input, report_output],
        outputs=[file_output],
    )

if __name__ == "__main__":
    demo.launch(share=True)