import gradio as gr
from PIL import Image
import numpy as np
import pandas as pd
import tempfile
import folium
from reportlab.pdfgen import canvas
from ultralytics import YOLO
from transformers import pipeline
import base64
import zipfile
import os

# ===========================================
# LOAD MODELS
# ===========================================
detection_model = YOLO("yolov5s.pt")
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# ===========================================
# GLOBAL STORAGE
# ===========================================
history = []
incident_counter = 1


# ===========================================
# HELPERS
# ===========================================
def generate_incident_id():
    global incident_counter
    code = f"IG-2025-{incident_counter:05d}"
    incident_counter += 1
    return code


def severity_badge(sev):
    colors = {
        "High": "#ff3b3b",
        "Medium": "#ffa500",
        "Low": "#00c853"
    }
    return f"<span style='padding:6px 12px;border-radius:6px;background:{colors[sev]};color:white;font-weight:bold'>{sev}</span>"


def generate_map(lat, lon):
    m = folium.Map(location=[lat, lon], zoom_start=15)
    folium.Marker([lat, lon], popup="Incident Location").add_to(m)

    f = tempfile.NamedTemporaryFile(delete=False, suffix=".html")
    m.save(f.name)
    f.close()
    return f"""<iframe src="{f.name}" width="100%" height="350" style="border:none;border-radius:10px;"></iframe>"""


def create_pdf(objects, severity, summary, lat, lon, incident_id):
    f = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
    c = canvas.Canvas(f.name)

    c.setFont("Helvetica-Bold", 18)
    c.drawString(100, 800, f"InfraGuard Incident Report")

    c.setFont("Helvetica", 12)
    c.drawString(100, 770, f"Incident ID: {incident_id}")
    c.drawString(100, 750, f"Objects: {objects}")
    c.drawString(100, 730, f"Severity: {severity}")
    c.drawString(100, 710, f"AI Summary: {summary}")
    c.drawString(100, 690, f"Location: {lat}, {lon}")

    c.save()
    return f.name


# ===========================================
# MAIN PROCESS FUNCTION
# ===========================================
def process(img, lat, lon):
    global history

    if img is None:
        return None, "Upload an image first", "N/A", None, None, None, pd.DataFrame(history)

    img_np = np.array(img)
    results = detection_model.predict(img_np, imgsz=640)

    # Object detection
    objs = [detection_model.names[int(box.cls)] for box in results[0].boxes]
    objects_text = ", ".join(objs) if objs else "None"

    # Severity
    severity = "Low"
    if "fire" in objs:
        severity = "High"
    elif any(x in objs for x in ["car", "truck", "bus"]):
        severity = "Medium"

    # Summary
    story = f"Detected objects: {objects_text}. Severity: {severity}."
    try:
        summary = summarizer(story, max_length=40, min_length=10)[0]["summary_text"]
    except:
        summary = story

    # Annotated image
    annotated = Image.fromarray(results[0].plot())

    # Incident ID
    incident_id = generate_incident_id()

    # PDF
    pdf_file = create_pdf(objects_text, severity, summary, lat, lon, incident_id)

    # Map
    map_html = generate_map(lat, lon)

    # Save to history
    thumb = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    img.save(thumb.name)

    entry = {
        "ID": incident_id,
        "Image": thumb.name,
        "Objects": objects_text,
        "Severity": severity,
        "Summary": summary,
        "Latitude": lat,
        "Longitude": lon,
        "PDF": pdf_file
    }
    history.append(entry)

    df = pd.DataFrame(history)

    # Timeline card
    card = f"""
    <div style='background:rgba(0,60,90,0.4);border:1px solid #00eaff;padding:18px;margin-top:10px;border-radius:10px;backdrop-filter:blur(10px);'>
        <h3 style='color:#00eaff'>New Incident Recorded</h3>
        <p><b>ID:</b> {incident_id}</p>
        <p><b>Objects:</b> {objects_text}</p>
        <p><b>Severity:</b> {severity_badge(severity)}</p>
        <p><b>Location:</b> {lat}, {lon}</p>
    </div>
    """

    return annotated, summary, severity, pdf_file, map_html, card, df


# ===========================================
# EXPORT ALL PDFS AS ZIP
# ===========================================
def download_all_pdfs():
    if not history:
        return None

    zip_path = tempfile.NamedTemporaryFile(delete=False, suffix=".zip").name
    with zipfile.ZipFile(zip_path, "w") as z:
        for h in history:
            z.write(h["PDF"], os.path.basename(h["PDF"]))
    return zip_path


# ===========================================
# DELETE ROW
# ===========================================
def delete_row(idx):
    global history
    if 0 <= idx < len(history):
        history.pop(idx)
    return pd.DataFrame(history)


# ===========================================
# CUSTOM CSS (Cyber Security Pro Theme)
# ===========================================
css = """
body { background:#03121f !important; }

.gradio-container {
    background:#03121f !important;
    color:white !important;
}

h1 {
    color:#00eaff;
    text-align:center;
    font-family:'Arial Black';
}

.gr-button {
    background:#00eaff !important;
    color:black !important;
    border-radius:10px !important;
    border:1px solid #00ffff !important;
    font-weight:bold;
    transition:0.2s;
}

.gr-button:hover {
    box-shadow:0 0 10px #00eaff;
    transform:scale(1.03);
}
"""


# ===========================================
# UI
# ===========================================
with gr.Blocks(css=css) as demo:
    gr.HTML("<h1>🚨 INFRA GUARD – Cyber Security AI Dashboard</h1>")

    with gr.Tabs():
        # ---------------- ANALYZE TAB ----------------
        with gr.Tab("Analyze Incident"):
            with gr.Row():
                with gr.Column(scale=1):
                    img_in = gr.Image(type="pil", label="Upload Image")
                    lat = gr.Number(value=24.8607, label="Latitude")
                    lon = gr.Number(value=67.0011, label="Longitude")
                    btn = gr.Button("Analyze", variant="primary")

                with gr.Column(scale=1):
                    img_out = gr.Image(label="Annotated Result")
                    summary = gr.Textbox(label="AI Summary")
                    severity = gr.Textbox(label="Severity")
                    pdf = gr.File(label="PDF Report")
                    map_html = gr.HTML(label="Incident Map")

            gr.Markdown("### 🕒 Timeline Feed")
            timeline = gr.HTML()

        # ---------------- HISTORY TAB ----------------
        with gr.Tab("Incident History"):
            hist_table = gr.Dataframe(headers=["ID", "Image", "Objects", "Severity", "Summary", "Latitude", "Longitude", "PDF"], value=[])
            delete_index = gr.Number(label="Delete Row Index")
            delete_btn = gr.Button("Delete Selected")
            delete_btn.click(delete_row, inputs=[delete_index], outputs=[hist_table])

            export_btn = gr.Button("Download All PDFs (ZIP)")
            export_zip = gr.File()
            export_btn.click(download_all_pdfs, outputs=export_zip)

    btn.click(
        process,
        inputs=[img_in, lat, lon],
        outputs=[img_out, summary, severity, pdf, map_html, timeline, hist_table]
    )


demo.launch()