import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
from torchvision import transforms
import pickle
import json
import math


# ── MLP Model Definition ──────────────────────────────────────────────────────
class MalwareClassifier(nn.Module):
    def __init__(self, input_dim=2381):
        super(MalwareClassifier, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.network(x).squeeze(1)


# ── ViT Model Definition ──────────────────────────────────────────────────────
class PatchEmbedding(nn.Module):
    def __init__(self, img_size, patch_size, in_channels, embed_dim):
        super().__init__()
        self.projection = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        x = self.projection(x)
        x = x.flatten(2)
        x = x.transpose(1, 2)
        return x


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads
        self.scale     = self.head_dim ** -0.5
        self.qkv       = nn.Linear(embed_dim, embed_dim * 3)
        self.proj      = nn.Linear(embed_dim, embed_dim)
        self.dropout   = nn.Dropout(dropout)

    def forward(self, x):
        B, N, D = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.dropout(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, D)
        return self.proj(x)


class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, mlp_dim, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn  = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.ffn   = nn.Sequential(
            nn.Linear(embed_dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, embed_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x


class VisionTransformer(nn.Module):
    def __init__(self, img_size, patch_size, in_channels, num_classes,
                 embed_dim, depth, num_heads, mlp_dim, dropout=0.1):
        super().__init__()
        num_patches      = (img_size // patch_size) ** 2
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        self.cls_token   = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed   = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.dropout     = nn.Dropout(dropout)
        self.transformer = nn.Sequential(*[
            TransformerBlock(embed_dim, num_heads, mlp_dim, dropout)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        B   = x.shape[0]
        x   = self.patch_embed(x)
        cls = self.cls_token.expand(B, -1, -1)
        x   = torch.cat([cls, x], dim=1)
        x   = x + self.pos_embed
        x   = self.dropout(x)
        x   = self.transformer(x)
        x   = self.norm(x[:, 0])
        return self.head(x)


# ── Load models ───────────────────────────────────────────────────────────────
device = torch.device("cpu")

# MLP
mlp_model = MalwareClassifier(input_dim=2381)
mlp_model.load_state_dict(torch.load("malware_classifier.pth", map_location=device))
mlp_model.eval()

with open("scaler.pkl", "rb") as f:
    scaler = pickle.load(f)

# ViT
with open("vit_class_names.json", "r") as f:
    class_names = json.load(f)

vit_model = VisionTransformer(
    img_size=64, patch_size=8, in_channels=1,
    num_classes=len(class_names),
    embed_dim=256, depth=6, num_heads=8, mlp_dim=512, dropout=0.1
)
vit_model.load_state_dict(torch.load("vit_malware_final.pth", map_location=device))
vit_model.eval()

# EMBER feature extractor
from ember.features import PEFeatureExtractor
extractor = PEFeatureExtractor(feature_version=2)

# ViT image transforms
vit_transforms = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((64, 64)),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])


# ── Helper: raw bytes → grayscale image ──────────────────────────────────────
def bytes_to_image(file_bytes):
    arr  = np.frombuffer(file_bytes, dtype=np.uint8)
    size = int(np.sqrt(len(arr)))
    if size < 8:
        return None
    arr = arr[:size * size].reshape(size, size)
    return Image.fromarray(arr, mode="L")


# ── Two-stage pipeline ────────────────────────────────────────────────────────
def analyze_file(file):
    if file is None:
        return (
            "No file uploaded.",
            "—",
            "—",
            "—",
            None,
            ""
        )

    try:
        with open(file, "rb") as f:
            file_bytes = f.read()

        # ── Stage 1: MLP malware detection ───────────────────────────────────
        features = extractor.feature_vector(file_bytes)
        features = np.array(features, dtype=np.float32).reshape(1, -1)
        features = scaler.transform(features)

        tensor      = torch.tensor(features, dtype=torch.float32).to(device)
        with torch.no_grad():
            mlp_prob = mlp_model(tensor).item()

        mlp_verdict    = "MALWARE" if mlp_prob >= 0.5 else "BENIGN"
        mlp_confidence = f"{mlp_prob:.2%}"

        # ── Stage 2: ViT family classification (only if MLP says malware) ────
        if mlp_prob < 0.5:
            return (
                f"✅  BENIGN",
                mlp_confidence,
                "Skipped — file classified as benign by Stage 1",
                "—",
                None,
                (
                    f"Stage 1 (MLP) malware probability: {mlp_prob:.2%}\n"
                    f"Below threshold of 50% — classified as benign.\n"
                    f"Stage 2 (ViT) was not run."
                )
            )

        # Convert bytes to grayscale image for ViT
        img = bytes_to_image(file_bytes)
        if img is None:
            return (
                f"⚠️  MALWARE (Stage 2 failed)",
                mlp_confidence,
                "Could not convert file to image",
                "—",
                None,
                f"Stage 1 probability: {mlp_prob:.2%}\nFile too small for ViT conversion."
            )

        img_tensor = vit_transforms(img).unsqueeze(0).to(device)

        with torch.no_grad():
            logits   = vit_model(img_tensor)
            probs    = F.softmax(logits, dim=1).squeeze()
            conf, idx = probs.max(0)

        family     = class_names[idx.item()]
        vit_conf   = f"{conf.item():.2%}"

        # Top 5 predictions
        top5_vals, top5_idx = probs.topk(5)
        top5_text = "\n".join([
            f"  {i+1}. {class_names[i2.item()]:<22} {v.item():.2%}"
            for i, (v, i2) in enumerate(zip(top5_vals, top5_idx))
        ])

        details = (
            f"Stage 1 — MLP (EMBER features)\n"
            f"  Malware probability : {mlp_prob:.2%}\n"
            f"  Verdict             : MALWARE — passed to Stage 2\n\n"
            f"Stage 2 — ViT (grayscale image)\n"
            f"  Image size          : {img.size[0]}x{img.size[1]} px from {len(file_bytes):,} bytes\n"
            f"  Predicted family    : {family}\n"
            f"  Confidence          : {conf.item():.2%}\n\n"
            f"Top 5 family predictions:\n{top5_text}"
        )

        # Display the grayscale image
        img_display = img.resize((256, 256), Image.NEAREST)

        return (
            f"🔴  MALWARE",
            mlp_confidence,
            family,
            vit_conf,
            img_display,
            details
        )

    except Exception as e:
        return (
            "ERROR",
            "—",
            "—",
            "—",
            None,
            f"Could not process file: {str(e)}"
        )


# ── Gradio UI ─────────────────────────────────────────────────────────────────
with gr.Blocks(title="Malware Detection Pipeline") as demo:

    gr.Markdown("# Two-Stage Malware Detection Pipeline")
    gr.Markdown(
        "Upload a Windows executable (`.exe` or `.dll`) and the pipeline will:\n\n"
        "**Stage 1 — MLP:** Analyze 2381 static PE features to determine if the file is malware\n\n"
        "**Stage 2 — ViT:** If malware is detected, convert the binary to a grayscale image "
        "and classify it into one of 25 known malware families using a Vision Transformer\n\n"
        "No file is executed at any point."
    )

    with gr.Row():
        file_input = gr.File(label="Upload PE file (.exe or .dll)")

    with gr.Row():
        analyze_btn = gr.Button("Analyze", variant="primary", size="lg")

    gr.Markdown("### Stage 1 — MLP Malware Detection")
    with gr.Row():
        verdict_out  = gr.Textbox(label="Verdict",             interactive=False)
        mlp_conf_out = gr.Textbox(label="Malware probability", interactive=False)

    gr.Markdown("### Stage 2 — ViT Family Classification")
    with gr.Row():
        family_out   = gr.Textbox(label="Malware family",  interactive=False)
        vit_conf_out = gr.Textbox(label="ViT confidence",  interactive=False)

    with gr.Row():
        image_out = gr.Image(
            label="Grayscale image fed to ViT (what the transformer sees)",
            type="pil"
        )

    details_out = gr.Textbox(label="Full analysis details", interactive=False, lines=12)

    analyze_btn.click(
        fn=analyze_file,
        inputs=file_input,
        outputs=[verdict_out, mlp_conf_out, family_out, vit_conf_out, image_out, details_out]
    )

    gr.Markdown(
        "---\n"
        "**Stage 1 — MLP:** Trained on EMBER 2018 · 800k samples · 95% accuracy · ROC-AUC 0.9878\n\n"
        "**Stage 2 — ViT:** Trained on MalImg · 9,339 images · 25 families · 98% accuracy · Built from scratch"
    )


app = demo.app