"""
Gradio UI Application for the Dual-Engine Malware Analysis Pipeline.

This module provides the frontend interface for uploading executables or simulated profiles,
routing them to the static structural engine (Engine A) and the visual byte-map engine (Engine B),
and rendering their ensemble metrics securely in the browser.
"""

import gradio as gr
from src.engine_a.inference import EngineAInfer
from src.engine_b.inference import EngineBInfer
import os
from datetime import datetime
from zoneinfo import ZoneInfo
import torch
from fpdf import FPDF
import tempfile

print("Loading ML models...")
engine_a = EngineAInfer()
engine_b = EngineBInfer()
print("Models loaded successfully.")


def analyze_malware(file_path):
    """
    Main processing pipeline triggered by the Gradio 'Analyze File' button.

    Args:
        file_path (str): The path to the uploaded file.

    Returns:
        tuple: (output_a_text, image_b, prob_table, ensemble_verdict) for the UI components.
    """
    if not file_path:
        return "No file provided.", None, {}, "N/A"

    try:
        # Engine A Analysis
        result_a = engine_a.predict(file_path)
        prob_a = result_a["malware_prob"] * 100
        status_a = "Malicious" if result_a["is_malware"] else "Benign"

        text_a = {
            "Malicious": result_a["malware_prob"],
            "Benign": 1.0 - result_a["malware_prob"],
        }

        # Explainability 
        explain_dict = {}
        total_abs_shap = sum(abs(f[1]) for f in result_a["top_features"])
        if total_abs_shap == 0:
            total_abs_shap = 1

        for feature, shap_val in result_a["top_features"]:
            impact_dir = "Malicious" if shap_val > 0 else "Benign"
            label_name = f"{feature} [{impact_dir}]"
            # Normalize to 0-1 so gr.Label can render the horizontal bars properly
            normalized_impact = float(abs(shap_val) / total_abs_shap)
            explain_dict[label_name] = normalized_impact

        # Engine B Analysis
        result_b = engine_b.predict(file_path)
        all_probs = result_b["all_probabilities"]
        image_b = result_b["image"]

        # Ensemble Logic
        ensemble_score = (prob_a + result_b["confidence"] * 100) / 2
        final_verdict_text = "MALICIOUS" if ensemble_score > 50 else "BENIGN"
        ensemble_verdict = f"Final Verdict: {final_verdict_text} (Risk Score: {ensemble_score:.1f}/100)"

        top_3_probs = sorted(all_probs.items(), key=lambda x: x[1], reverse=True)[:3]

        # Generate PDF Report
        pdf_path = generate_pdf_report(
            file_path=file_path,
            status_a=status_a,
            prob_a=prob_a,
            top_features=result_a["top_features"],
            family_b=result_b["family"],
            prob_b=result_b["confidence"] * 100,
            image_b=image_b,
            ensemble_score=ensemble_score,
            final_verdict=final_verdict_text,
            top_3_probs=top_3_probs,
        )

        return (
            text_a,
            explain_dict,
            image_b,
            all_probs,
            ensemble_verdict,
            gr.update(value=pdf_path, visible=True),
        )

    except Exception as e:
        return (
            f"Error processing file in Engine A: {e}",
            {},
            None,
            {},
            "Error",
            gr.update(visible=False),
        )


def generate_pdf_report(
    file_path,
    status_a,
    prob_a,
    top_features,
    family_b,
    prob_b,
    image_b,
    ensemble_score,
    final_verdict,
    top_3_probs,
):

    class PDFReport(FPDF):
        def header(self):
            # Header Bar
            self.set_fill_color(40, 40, 40)
            self.rect(0, 0, 210, 30, "F")
            self.set_y(10)
            self.set_font("Helvetica", style="B", size=24)
            self.set_text_color(255, 255, 255)
            self.cell(0, 10, "MALWARE ANALYSIS REPORT", border=0, align="C")
            self.ln(15)

    pdf = PDFReport()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.set_y(40)  # Start below header

    # 1. Overview Section
    pdf.set_fill_color(230, 230, 230)
    pdf.set_text_color(0, 0, 0)
    pdf.set_font("Helvetica", style="B", size=14)
    pdf.cell(
        0,
        10,
        " 1. Execution Overview",
        border=0,
        new_x="LMARGIN",
        new_y="NEXT",
        fill=True,
    )
    pdf.ln(2)

    pdf.set_font("Helvetica", size=11)
    pdf.cell(40, 8, "Target File:", border=0)
    pdf.set_font("Helvetica", style="B", size=11)
    pdf.cell(
        0, 8, f"{os.path.basename(file_path)}", border=0, new_x="LMARGIN", new_y="NEXT"
    )

    pdf.set_font("Helvetica", size=11)
    pdf.cell(40, 8, "Timestamp:", border=0)
    pdf.set_font("Helvetica", style="B", size=11)
    pdf.cell(
        0,
        8,
        f"{datetime.now(ZoneInfo('Asia/Kolkata')).strftime('%Y-%m-%d %H:%M:%S')} IST",
        border=0,
        new_x="LMARGIN",
        new_y="NEXT",
    )

    pdf.set_font("Helvetica", size=11)
    pdf.cell(40, 8, "Final Verdict:", border=0)
    pdf.set_font("Helvetica", style="B", size=12)

    pdf.cell(
        0,
        8,
        f"{final_verdict} (Risk Score: {ensemble_score:.1f} / 100)",
        border=0,
        new_x="LMARGIN",
        new_y="NEXT",
    )
    pdf.ln(5)

    # 2. Engine A
    pdf.set_font("Helvetica", style="B", size=14)
    pdf.cell(
        0,
        10,
        " 2. Structural Engine (EMBER)",
        border=0,
        new_x="LMARGIN",
        new_y="NEXT",
        fill=True,
    )
    pdf.ln(2)

    pdf.set_font("Helvetica", size=11)
    pdf.cell(
        0,
        8,
        f"Verdict: {status_a.upper()} (Confidence: {prob_a:.1f}%)",
        new_x="LMARGIN",
        new_y="NEXT",
    )
    pdf.ln(2)

    pdf.set_font("Helvetica", style="B", size=10)
    pdf.cell(140, 8, "Top Contributing Feature (SHAP)", border=1, fill=True)
    pdf.cell(
        50,
        8,
        "Impact Direction",
        border=1,
        new_x="LMARGIN",
        new_y="NEXT",
        fill=True,
        align="C",
    )

    pdf.set_font("Helvetica", size=10)
    for feature, shap_val in top_features:
        impact_dir = "MALICIOUS" if shap_val > 0 else "BENIGN"
        pdf.cell(140, 8, f" {feature}", border=1)
        pdf.cell(
            50,
            8,
            f"{shap_val:+.2f} ({impact_dir})",
            border=1,
            new_x="LMARGIN",
            new_y="NEXT",
            align="C",
        )
    pdf.ln(8)

    # 3. Engine B
    pdf.set_font("Helvetica", style="B", size=14)
    pdf.cell(
        0,
        10,
        " 3. Visual Engine (Malimg)",
        border=0,
        new_x="LMARGIN",
        new_y="NEXT",
        fill=True,
    )
    pdf.ln(2)

    pdf.set_font("Helvetica", size=11)
    pdf.cell(
        0,
        8,
        f"Predicted Malware Family: {family_b.upper()} (Confidence: {prob_b:.1f}%)",
        new_x="LMARGIN",
        new_y="NEXT",
    )
    pdf.ln(2)

    # Embed Image
    temp_dir = tempfile.gettempdir()
    img_path = os.path.join(temp_dir, "byte_map.png")
    image_b.save(img_path)

    pdf.set_font("Helvetica", style="B", size=10)
    pdf.cell(0, 8, "Grayscale Byte Map Render:", new_x="LMARGIN", new_y="NEXT")

    # Draw image with a border
    x_pos = pdf.get_x()
    y_pos = pdf.get_y()
    pdf.rect(x_pos, y_pos, 70, 70)
    pdf.image(img_path, x=x_pos, y=y_pos, w=70, h=70)

    # Draw Top 3 Probabilities Table next to the image
    pdf.set_y(y_pos)

    # Table Header
    pdf.set_x(x_pos + 75)
    pdf.set_font("Helvetica", style="B", size=10)
    pdf.cell(70, 8, "Predicted Family", border=1, fill=True)
    pdf.cell(
        40,
        8,
        "Confidence",
        border=1,
        new_x="LMARGIN",
        new_y="NEXT",
        fill=True,
        align="C",
    )

    # Table Rows
    pdf.set_font("Helvetica", size=10)
    for family, prob in top_3_probs:
        pdf.set_x(x_pos + 75)
        pdf.cell(70, 8, f" {family}", border=1)
        pdf.cell(
            40,
            8,
            f"{prob * 100:.2f}%",
            border=1,
            new_x="LMARGIN",
            new_y="NEXT",
            align="C",
        )

    report_path = os.path.join(
        temp_dir, f"Analysis_Report_{os.path.basename(file_path)}.pdf"
    )
    pdf.output(report_path)
    return report_path


def get_system_status():
    """
    Fetches real-time system metrics for the UI header.

    Returns:
        str: Formatted markdown string containing time, compute device, and engine status.
    """
    current_time = datetime.now(ZoneInfo("Asia/Kolkata")).strftime("%Y-%m-%d %H:%M:%S")
    device = "CUDA" if torch.cuda.is_available() else "CPU"
    return f"**System Time:** {current_time} IST | **Compute Device:** {device} | **Engines:** 2 (Static Structural & Visual Byte-Map)"


# Build the Gradio UI
with gr.Blocks() as app:
    # Use HTML to make the title massively larger than the buttons
    gr.HTML(
        "<h1 style='text-align: center; font-size: 3.5rem; margin-bottom: 0.2rem; font-weight: bold;'>Dual-Engine Malware Analysis Pipeline</h1>"
    )
    gr.HTML(
        "<p style='text-align: center; font-size: 1.2rem; color: gray;'>Upload a Windows Executable (.exe, .dll) or a Simulated Profile (.json)</p>"
    )

    status_bar = gr.Markdown(get_system_status())

    # Compact inputs
    with gr.Row():
        file_input = gr.File(label="Upload File", type="filepath", scale=4)
        with gr.Column(scale=1):
            analyze_btn = gr.Button("Analyze File", variant="primary", size="lg")

    download_btn = gr.DownloadButton("Download Report (PDF)", visible=False, size="lg")

    gr.Markdown("---")

    # Single column layout for maximum horizontal space
    ensemble_output = gr.Textbox(label="Final Verdict", interactive=False, lines=1)
    output_a = gr.Label(
        label="Engine A: Structural (EMBER)",
    )
    explain_plot = gr.Label(
        label="Engine A: Explainability (Feature Impact)", num_top_classes=3
    )
    image_output = gr.Image(label="Engine B: Grayscale Byte Map", type="pil")
    output_b = gr.Label(label="Engine B: Family Probabilities", num_top_classes=24)

    analyze_btn.click(
        analyze_malware,
        inputs=[file_input],
        outputs=[
            output_a,
            explain_plot,
            image_output,
            output_b,
            ensemble_output,
            download_btn,
        ],
    )

    # Live update the status bar
    timer = gr.Timer(1)
    timer.tick(get_system_status, inputs=None, outputs=status_bar)

    # Load immediately on page open
    app.load(get_system_status, inputs=None, outputs=status_bar)