import streamlit as st
import lightgbm as lgb
import numpy as np
import ember
import tempfile
import os
import sys
import shap
import lief
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import pandas as pd
import hashlib
import requests
import re
import yara
import networkx as nx
from collections import Counter
from datetime import datetime
from fpdf import FPDF
import urllib.request

# Fix ember compatibility with newer lief
import ember.features as _ember_features

def _patched_raw_features(self, bytez):
    import hashlib
    import lief as _lief
    try:
        lief_binary = _lief.PE.parse(list(bytez))
    except Exception:
        lief_binary = None
    features = {"sha256": hashlib.sha256(bytez).hexdigest()}
    features.update({fe.name: fe.raw_features(bytez, lief_binary) for fe in self.features})
    return features

from sklearn.feature_extraction import FeatureHasher as _FeatureHasher
_original_transform = _FeatureHasher.transform

def _patched_transform(self, X):
    if self.input_type == "string":
        X = [[x] if isinstance(x, str) else x for x in X]
    return _original_transform(self, X)

_FeatureHasher.transform = _patched_transform
_ember_features.PEFeatureExtractor.raw_features = _patched_raw_features

if not hasattr(np, 'int'):
    np.int = int

MODEL_PATH = "data/ember_model_2018.txt"
MODEL_URL = "https://huggingface.co/SamSJ13/ember-model-2018/resolve/main/ember_model_2018.txt"

if not os.path.exists(MODEL_PATH):
    os.makedirs("data", exist_ok=True)
    with st.spinner("Downloading model... this may take a minute."):
        urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)

sys.path.insert(0, "src")

VT_API_KEY = "2c4d35875f3461ecf225e4b2c7b2d7c1d9c60627eac2b5fa90a9f966c50c29ca"

st.set_page_config(page_title="Malware Classifier", page_icon="🛡️", layout="wide")

if "dark_mode" not in st.session_state:
    st.session_state.dark_mode = True
if "session_history" not in st.session_state:
    st.session_state.session_history = []

col_title, col_theme = st.columns([6, 1])
with col_title:
    st.title("PE Malware Classifier")
    st.write("Upload a Windows executable (.exe, .dll, .sys) to analyze it using the EMBER 2018 model trained on 1.1M samples.")
with col_theme:
    if st.button("Toggle Theme", key="theme_btn"):
        st.session_state.dark_mode = not st.session_state.dark_mode

dark = st.session_state.dark_mode
text_color = "white" if dark else "black"
plot_bg = "#0e1117" if dark else "#f0f0f0"

mode = st.radio("Mode", ["Single File Analysis", "Compare Two Files", "Batch Analysis", "Session History"], horizontal=True)

@st.cache_resource
def load_model():
    return lgb.Booster(model_file=MODEL_PATH)

model = load_model()

def get_virustotal_report(file_hash):
    url = f"https://www.virustotal.com/api/v3/files/{file_hash}"
    headers = {"x-apikey": VT_API_KEY}
    response = requests.get(url, headers=headers)
    if response.status_code == 200:
        return response.json()
    return None

def extract_strings(bytez):
    pattern = re.compile(b'[\x20-\x7f]{5,}')
    strings = [s.decode('ascii', errors='ignore') for s in pattern.findall(bytez)]
    urls = [s for s in strings if s.startswith('http://') or s.startswith('https://')]
    registry = [s for s in strings if s.startswith('HKEY_')]
    paths = [s for s in strings if re.match(r'[A-Za-z]:\\', s)]
    mz = [s for s in strings if 'MZ' in s]
    return {
        "total_strings": len(strings),
        "urls": urls[:20],
        "registry_keys": registry[:20],
        "file_paths": paths[:20],
        "mz_headers": mz[:10],
    }

def run_yara_scan(bytez):
    try:
        rules = yara.compile(filepath="rules.yar")
        matches = rules.match(data=bytez)
        results = []
        for match in matches:
            severity = match.meta.get("severity", "medium")
            description = match.meta.get("description", match.rule)
            results.append({"rule": match.rule, "description": description, "severity": severity})
        return results
    except Exception:
        return []

def detect_packer(binary, bytez):
    packers = []
    if binary:
        section_names = [s.name.lower() for s in binary.sections]
        if any('upx' in n for n in section_names):
            packers.append("UPX")
        if any('aspack' in n for n in section_names):
            packers.append("ASPack")
        if any('themida' in n for n in section_names):
            packers.append("Themida")
        if any(s.entropy > 7.5 for s in binary.sections):
            packers.append("Unknown packer (high entropy)")
    if b'UPX!' in bytez and "UPX" not in packers:
        packers.append("UPX (signature found)")
    return packers

def detect_malware_family(indicators, strings, binary):
    hints = []
    all_imports = []
    if binary and binary.has_imports:
        for lib in binary.imports:
            for entry in lib.entries:
                if not entry.is_ordinal:
                    all_imports.append(entry.name.lower())
    if any(api in all_imports for api in ['cryptencrypt', 'cryptdecrypt', 'findfiles']):
        hints.append(("Ransomware", "Encryption APIs + file enumeration detected"))
    if any(api in all_imports for api in ['setwindowshookex', 'keybd_event', 'getasynckeystate']):
        hints.append(("Keylogger", "Keyboard hooking APIs detected"))
    if any(api in all_imports for api in ['createremotethread', 'writeprocessmemory', 'virtualallocex']):
        hints.append(("RAT / Injector", "Process injection APIs detected"))
    if strings["mz_headers"]:
        hints.append(("Dropper", "Embedded executable found in binary"))
    if strings["urls"]:
        hints.append(("Downloader / C2", f"{len(strings['urls'])} embedded URLs found"))
    if any(api in all_imports for api in ['internetopen', 'httpsendrequesta', 'urldownloadtofile']):
        hints.append(("Downloader", "Network download APIs detected"))
    return hints

def get_threat_indicators(binary, score, strings):
    indicators = []
    if binary:
        if not binary.has_signatures:
            indicators.append(("No digital signature", "high"))
        if binary.has_tls:
            indicators.append(("Has TLS callbacks (common in malware)", "medium"))
        for s in binary.sections:
            if s.entropy > 7.0:
                indicators.append((f"High entropy section: {s.name} ({s.entropy:.2f}) — may be packed/encrypted", "high"))
        if binary.has_imports:
            all_imports = []
            for lib in binary.imports:
                for entry in lib.entries:
                    if not entry.is_ordinal:
                        all_imports.append(entry.name.lower())
            suspicious_apis = ["virtualalloc", "writeprocessmemory", "createremotethread",
                               "shellexecute", "winexec", "createprocess", "loadlibrary",
                               "getprocaddress", "setwindowshookex", "keybd_event"]
            for api in suspicious_apis:
                if api in all_imports:
                    indicators.append((f"Suspicious API: {api}", "medium"))
    if strings["urls"]:
        indicators.append((f"{len(strings['urls'])} embedded URL(s) found", "medium"))
    if strings["mz_headers"]:
        indicators.append(("Embedded MZ header — possible dropper", "high"))
    if score > 0.8:
        indicators.append((f"Model confidence very high: {score:.2%}", "high"))
    elif score > 0.5:
        indicators.append((f"Model flagged as malware: {score:.2%}", "medium"))
    return indicators

def analyze_file(uploaded_file):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".exe") as tmp:
        tmp.write(uploaded_file.read())
        tmp_path = tmp.name
    try:
        bytez = open(tmp_path, "rb").read()
        file_hash = hashlib.sha256(bytez).hexdigest()
        md5_hash = hashlib.md5(bytez).hexdigest()
        extractor = ember.PEFeatureExtractor(2, print_feature_warning=False)
        features = np.array(extractor.feature_vector(bytez), dtype=np.float32)
        score = model.predict([features])[0]
        verdict = "MALWARE" if score > 0.5 else "BENIGN"
        binary = lief.parse(tmp_path)
        strings = extract_strings(bytez)
        indicators = get_threat_indicators(binary, score, strings)
        packers = detect_packer(binary, bytez)
        yara_matches = run_yara_scan(bytez)
        family_hints = detect_malware_family(indicators, strings, binary)
        return {
            "filename": uploaded_file.name,
            "bytez": bytez,
            "file_hash": file_hash,
            "md5_hash": md5_hash,
            "features": features,
            "score": score,
            "verdict": verdict,
            "binary": binary,
            "strings": strings,
            "indicators": indicators,
            "packers": packers,
            "family_hints": family_hints,
            "yara_matches": yara_matches,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
    finally:
        os.unlink(tmp_path)

def build_import_graph(binary):
    G = nx.DiGraph()
    if binary and binary.has_imports:
        G.add_node("target", node_type="file")
        for lib in binary.imports:
            G.add_node(lib.name, node_type="dll")
            G.add_edge("target", lib.name, weight=len(lib.entries))
    return G

def render_analysis(result, show_vt=True, show_shap=True):
    score = result["score"]
    verdict = result["verdict"]
    binary = result["binary"]
    strings = result["strings"]
    indicators = result["indicators"]
    bytez = result["bytez"]
    features = result["features"]
    confidence = score * 100 if score > 0.5 else (1 - score) * 100

    if verdict == "MALWARE":
        st.error("Verdict: MALWARE")
    else:
        st.success("Verdict: BENIGN")

    col1, col2, col3 = st.columns(3)
    col1.metric("Malware Probability", f"{score:.4f}")
    col2.metric("Confidence", f"{confidence:.1f}%")
    col3.metric("File Size", f"{len(bytez) / 1024:.1f} KB")
    st.progress(float(score), text=f"Malware score: {score:.4f}")
    st.divider()

    st.subheader("File Hashes")
    c1, c2 = st.columns(2)
    c1.code(f"SHA256: {result['file_hash']}")
    c2.code(f"MD5:    {result['md5_hash']}")
    st.divider()

    st.subheader("Packer Detection")
    if result["packers"]:
        for p in result["packers"]:
            st.error(f"Packer detected: {p}")
    else:
        st.success("No known packers detected.")
    st.divider()

    st.subheader("Malware Family Hints")
    if result["family_hints"]:
        for family, reason in result["family_hints"]:
            st.warning(f"{family}: {reason}")
    else:
        st.success("No malware family patterns detected.")
    st.divider()

    st.subheader("YARA Rule Matches")
    if result.get("yara_matches"):
        for match in result["yara_matches"]:
            if match["severity"] == "high":
                st.error(f"HIGH: [{match['rule']}] {match['description']}")
            elif match["severity"] == "medium":
                st.warning(f"MEDIUM: [{match['rule']}] {match['description']}")
            else:
                st.info(f"LOW: [{match['rule']}] {match['description']}")
    else:
        st.success("No YARA rules matched.")
    st.divider()

    st.subheader("Threat Indicators")
    if indicators:
        for text, level in indicators:
            if level == "high":
                st.error(f"HIGH: {text}")
            elif level == "medium":
                st.warning(f"MEDIUM: {text}")
            else:
                st.info(f"LOW: {text}")
    else:
        st.success("No threat indicators found.")
    st.divider()

    st.subheader("PE File Structure")
    if binary and binary.sections:
        section_data = []
        for s in binary.sections:
            section_data.append({
                "Name": s.name if s.name else "(unnamed)",
                "Size (KB)": round(s.size / 1024, 2),
                "Entropy": round(s.entropy, 3),
                "Virtual Size (KB)": round(s.virtual_size / 1024, 2),
            })
        st.dataframe(pd.DataFrame(section_data), use_container_width=True)
        fig, ax = plt.subplots(figsize=(5, 2.5))
        colors_bar = ["#ff4b4b" if s["Entropy"] > 7.0 else "#4b8fff" if s["Entropy"] > 6.0 else "#4bff91" for s in section_data]
        ax.bar([s["Name"] for s in section_data], [s["Entropy"] for s in section_data], color=colors_bar)
        ax.axhline(7.0, color="#ff4b4b", linestyle="--", linewidth=0.8, label="High entropy threshold")
        ax.set_ylabel("Entropy", color=text_color)
        ax.set_title("Section Entropy (>7.0 suspicious)", color=text_color)
        ax.set_ylim(0, 8.5)
        fig.patch.set_facecolor(plot_bg)
        ax.set_facecolor(plot_bg)
        ax.tick_params(colors=text_color)
        ax.yaxis.label.set_color(text_color)
        ax.legend(facecolor=plot_bg, labelcolor=text_color, fontsize=7)
        st.pyplot(fig, use_container_width=False)
        plt.close()
    st.divider()

    if binary and binary.has_imports:
        st.subheader("Imported Libraries")
        import_data = [{"Library": lib.name, "Functions Imported": len(lib.entries)} for lib in binary.imports]
        df_imports = pd.DataFrame(import_data).sort_values("Functions Imported", ascending=False)
        st.dataframe(df_imports, use_container_width=True)

        st.subheader("Import Graph")
        G = build_import_graph(binary)
        fig, ax = plt.subplots(figsize=(8, 5))
        pos = nx.spring_layout(G, seed=42)
        node_colors = ["#ff4b4b" if G.nodes[n].get("node_type") == "file" else "#4b8fff" for n in G.nodes]
        nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=800, ax=ax)
        nx.draw_networkx_labels(G, pos, font_size=6, font_color="white", ax=ax)
        nx.draw_networkx_edges(G, pos, edge_color="#888888", arrows=True, ax=ax)
        ax.set_title("DLL Import Graph (red = file, blue = DLL)", color=text_color)
        fig.patch.set_facecolor(plot_bg)
        ax.set_facecolor(plot_bg)
        ax.axis("off")
        st.pyplot(fig, use_container_width=False)
        plt.close()
    st.divider()

    st.subheader("Hex Viewer (first 256 bytes)")
    hex_bytes = bytez[:256]
    hex_lines = []
    for i in range(0, len(hex_bytes), 16):
        chunk = hex_bytes[i:i+16]
        hex_part = " ".join(f"{b:02x}" for b in chunk)
        ascii_part = "".join(chr(b) if 32 <= b < 127 else "." for b in chunk)
        hex_lines.append(f"{i:04x}  {hex_part:<47}  {ascii_part}")
    st.code("\n".join(hex_lines), language="text")
    st.divider()

    st.subheader("Extracted Strings")
    tabs = st.tabs(["URLs", "Registry Keys", "File Paths", "MZ Headers"])
    with tabs[0]:
        if strings["urls"]:
            for u in strings["urls"]:
                st.code(u)
        else:
            st.write("None found.")
    with tabs[1]:
        if strings["registry_keys"]:
            for r in strings["registry_keys"]:
                st.code(r)
        else:
            st.write("None found.")
    with tabs[2]:
        if strings["file_paths"]:
            for p in strings["file_paths"]:
                st.code(p)
        else:
            st.write("None found.")
    with tabs[3]:
        if strings["mz_headers"]:
            for m in strings["mz_headers"]:
                st.code(m)
        else:
            st.write("None found.")
    st.divider()

    if show_vt:
        st.subheader("VirusTotal Report")
        with st.spinner("Querying VirusTotal..."):
            vt_report = get_virustotal_report(result["file_hash"])
            if vt_report:
                stats = vt_report["data"]["attributes"]["last_analysis_stats"]
                malicious = stats.get("malicious", 0)
                total = sum(stats.values())
                st.metric("AV Engines Flagging as Malicious", f"{malicious} / {total}")
                if malicious > 0:
                    st.error(f"{malicious} AV engines detected this file as malicious.")
                    results = vt_report["data"]["attributes"]["last_analysis_results"]
                    flagged = {k: v for k, v in results.items() if v["category"] == "malicious"}
                    vt_data = [{"Engine": k, "Result": v["result"]} for k, v in flagged.items()]
                    st.dataframe(pd.DataFrame(vt_data), use_container_width=True)
                else:
                    st.success("No AV engines flagged this file.")
            else:
                st.info("File not found in VirusTotal database.")
        st.divider()

    if show_shap:
        st.subheader("Why did the model give this verdict?")
        with st.spinner("Computing SHAP explanation..."):
            explainer = shap.TreeExplainer(model)
            shap_values = explainer.shap_values(features.reshape(1, -1))
            shap_flat = shap_values[1][0] if isinstance(shap_values, list) else shap_values[0]
            top_idx = np.argsort(np.abs(shap_flat))[-15:][::-1]
            top_shap = shap_flat[top_idx]
            group_dims = [256, 256, 104, 10, 62, 255, 1280, 128, 30]
            group_names = ["byte_hist", "byte_entropy", "strings", "general",
                           "header", "section", "imports", "exports", "datadirs"]
            boundaries = np.cumsum(group_dims)

            def get_group(idx):
                for i, b in enumerate(boundaries):
                    if idx < b:
                        return group_names[i]
                return "unknown"

            labels = [f"{get_group(i)}[{i}]" for i in top_idx]
            fig, ax = plt.subplots(figsize=(5, 3))
            colors = ["#ff4b4b" if v > 0 else "#4b8fff" for v in top_shap]
            ax.barh(labels[::-1], top_shap[::-1], color=colors[::-1])
            ax.axvline(0, color=text_color, linewidth=0.8)
            ax.set_xlabel("SHAP value", color=text_color)
            ax.set_title("Top 15 features influencing prediction", color=text_color)
            fig.patch.set_facecolor(plot_bg)
            ax.set_facecolor(plot_bg)
            ax.tick_params(colors=text_color)
            ax.xaxis.label.set_color(text_color)
            st.pyplot(fig, use_container_width=False)
            plt.close()

# ── SINGLE FILE MODE ──────────────────────────────────────────────
if mode == "Single File Analysis":
    uploaded_file = st.file_uploader("Choose a PE file", type=["exe", "dll", "sys"])
    if uploaded_file is not None:
        with st.spinner("Analyzing..."):
            try:
                result = analyze_file(uploaded_file)
                st.session_state.last_result = result
                st.session_state.session_history.append({
                    "File": result["filename"],
                    "Verdict": result["verdict"],
                    "Score": round(result["score"], 4),
                    "Indicators": len(result["indicators"]),
                    "Packers": ", ".join(result["packers"]) if result["packers"] else "None",
                    "Time": result["timestamp"],
                })
                st.divider()
                render_analysis(result)
                st.divider()
                st.subheader("Export Report")
                if st.button("Generate PDF Report", key="pdf_btn"):
                    result = st.session_state.get("last_result", result)
                    pdf = FPDF()
                    pdf.add_page()
                    pdf.set_font("Helvetica", size=16)
                    pdf.cell(200, 10, "Malware Analysis Report", ln=True, align="C")
                    pdf.set_font("Helvetica", size=11)
                    pdf.ln(5)
                    pdf.cell(200, 8, f"File: {result['filename']}", ln=True)
                    pdf.cell(200, 8, f"Verdict: {result['verdict']}", ln=True)
                    pdf.cell(200, 8, f"Malware Score: {result['score']:.4f}", ln=True)
                    pdf.cell(200, 8, f"SHA256: {result['file_hash']}", ln=True)
                    pdf.cell(200, 8, f"MD5: {result['md5_hash']}", ln=True)
                    pdf.cell(200, 8, f"Analyzed: {result['timestamp']}", ln=True)
                    pdf.ln(5)
                    pdf.set_font("Helvetica", size=13)
                    pdf.cell(200, 8, "Threat Indicators:", ln=True)
                    pdf.set_font("Helvetica", size=10)
                    for text, level in result["indicators"]:
                        clean_text = text.replace("\u2014", "-").replace("\u2013", "-")
                        pdf.cell(200, 7, f"  [{level.upper()}] {clean_text}", ln=True)
                    pdf.ln(3)
                    pdf.set_font("Helvetica", size=13)
                    pdf.cell(200, 8, "Packers:", ln=True)
                    pdf.set_font("Helvetica", size=10)
                    for p in result["packers"]:
                        pdf.cell(200, 7, f"  {p}", ln=True)
                    if not result["packers"]:
                        pdf.cell(200, 7, "  None detected", ln=True)
                    pdf.ln(3)
                    pdf.set_font("Helvetica", size=13)
                    pdf.cell(200, 8, "Malware Family Hints:", ln=True)
                    pdf.set_font("Helvetica", size=10)
                    for family, reason in result["family_hints"]:
                        pdf.cell(200, 7, f"  {family}: {reason}", ln=True)
                    if not result["family_hints"]:
                        pdf.cell(200, 7, "  None detected", ln=True)
                    pdf_bytes = bytes(pdf.output())
                    st.download_button("Download PDF", data=pdf_bytes, file_name=f"malware_report_{result['filename']}.pdf", mime="application/pdf")
            except Exception as e:
                st.error(f"Error: {e}")
                st.exception(e)

# ── COMPARE MODE ──────────────────────────────────────────────────
elif mode == "Compare Two Files":
    col1, col2 = st.columns(2)
    with col1:
        file1 = st.file_uploader("File 1", type=["exe", "dll", "sys"], key="f1")
    with col2:
        file2 = st.file_uploader("File 2", type=["exe", "dll", "sys"], key="f2")
    if file1 and file2:
        with st.spinner("Analyzing both files..."):
            try:
                r1 = analyze_file(file1)
                r2 = analyze_file(file2)
                st.divider()
                st.subheader("Side-by-Side Comparison")
                fig, ax = plt.subplots(figsize=(5, 2.5))
                ax.bar([r1["filename"][:20], r2["filename"][:20]],
                       [r1["score"], r2["score"]],
                       color=["#ff4b4b" if r1["score"] > 0.5 else "#4bff91",
                              "#ff4b4b" if r2["score"] > 0.5 else "#4bff91"])
                ax.axhline(0.5, color="yellow", linestyle="--", linewidth=1, label="Threshold")
                ax.set_ylim(0, 1)
                ax.set_ylabel("Malware Score", color=text_color)
                ax.set_title("Malware Score Comparison", color=text_color)
                fig.patch.set_facecolor(plot_bg)
                ax.set_facecolor(plot_bg)
                ax.tick_params(colors=text_color)
                ax.yaxis.label.set_color(text_color)
                ax.legend(facecolor=plot_bg, labelcolor=text_color)
                st.pyplot(fig, use_container_width=False)
                plt.close()
                compare_data = {
                    "Metric": ["Verdict", "Malware Score", "File Size", "Sections",
                               "Imports", "Has Signature", "Packers", "Threat Indicators", "Family Hints"],
                    r1["filename"][:25]: [
                        r1["verdict"], f"{r1['score']:.4f}",
                        f"{len(r1['bytez'])/1024:.1f} KB",
                        len(r1["binary"].sections) if r1["binary"] else "N/A",
                        sum(len(l.entries) for l in r1["binary"].imports) if r1["binary"] and r1["binary"].has_imports else 0,
                        "Yes" if r1["binary"] and r1["binary"].has_signatures else "No",
                        ", ".join(r1["packers"]) if r1["packers"] else "None",
                        len(r1["indicators"]),
                        ", ".join([f[0] for f in r1["family_hints"]]) if r1["family_hints"] else "None",
                    ],
                    r2["filename"][:25]: [
                        r2["verdict"], f"{r2['score']:.4f}",
                        f"{len(r2['bytez'])/1024:.1f} KB",
                        len(r2["binary"].sections) if r2["binary"] else "N/A",
                        sum(len(l.entries) for l in r2["binary"].imports) if r2["binary"] and r2["binary"].has_imports else 0,
                        "Yes" if r2["binary"] and r2["binary"].has_signatures else "No",
                        ", ".join(r2["packers"]) if r2["packers"] else "None",
                        len(r2["indicators"]),
                        ", ".join([f[0] for f in r2["family_hints"]]) if r2["family_hints"] else "None",
                    ],
                }
                st.dataframe(pd.DataFrame(compare_data), use_container_width=True)
                st.divider()
                c1, c2 = st.columns(2)
                with c1:
                    st.subheader(f"File 1: {r1['filename']}")
                    render_analysis(r1, show_vt=False, show_shap=False)
                with c2:
                    st.subheader(f"File 2: {r2['filename']}")
                    render_analysis(r2, show_vt=False, show_shap=False)
            except Exception as e:
                st.error(f"Error: {e}")
                st.exception(e)

# ── BATCH ANALYSIS ───────────────────────────────────────────────
elif mode == "Batch Analysis":
    st.subheader("Batch File Analysis")
    st.write("Upload multiple PE files at once to analyze and compare them all.")
    uploaded_files = st.file_uploader("Choose PE files", type=["exe", "dll", "sys"], accept_multiple_files=True, key="batch_upload")
    if uploaded_files:
        if st.button("Analyze All", key="batch_btn"):
            batch_results = []
            progress = st.progress(0)
            for i, f in enumerate(uploaded_files):
                with st.spinner(f"Analyzing {f.name}..."):
                    try:
                        r = analyze_file(f)
                        batch_results.append({
                            "File": r["filename"],
                            "Verdict": r["verdict"],
                            "Score": round(r["score"], 4),
                            "Packers": ", ".join(r["packers"]) if r["packers"] else "None",
                            "YARA Hits": len(r["yara_matches"]),
                            "Threat Indicators": len(r["indicators"]),
                            "Family Hints": ", ".join([x[0] for x in r["family_hints"]]) if r["family_hints"] else "None",
                            "SHA256": r["file_hash"],
                        })
                        st.session_state.session_history.append({
                            "File": r["filename"],
                            "Verdict": r["verdict"],
                            "Score": round(r["score"], 4),
                            "Indicators": len(r["indicators"]),
                            "Packers": ", ".join(r["packers"]) if r["packers"] else "None",
                            "Time": r["timestamp"],
                        })
                    except Exception as e:
                        batch_results.append({
                            "File": f.name, "Verdict": "ERROR", "Score": -1,
                            "Packers": "N/A", "YARA Hits": 0, "Threat Indicators": 0,
                            "Family Hints": str(e), "SHA256": "N/A",
                        })
                progress.progress((i + 1) / len(uploaded_files))
            st.divider()
            st.subheader("Batch Results")
            df_batch = pd.DataFrame(batch_results)

            def color_verdict(val):
                if val == "MALWARE":
                    return "background-color: #ff4b4b; color: white"
                elif val == "BENIGN":
                    return "background-color: #1a4a1a; color: #4bff91"
                return ""

            st.dataframe(df_batch.style.applymap(color_verdict, subset=["Verdict"]), use_container_width=True)
            fig, ax = plt.subplots(figsize=(max(6, len(batch_results)), 3))
            colors = ["#ff4b4b" if r["Verdict"] == "MALWARE" else "#4bff91" for r in batch_results]
            ax.bar([r["File"][:15] for r in batch_results], [r["Score"] for r in batch_results], color=colors)
            ax.axhline(0.5, color="yellow", linestyle="--", linewidth=1, label="Malware threshold")
            ax.set_ylim(0, 1)
            ax.set_ylabel("Malware Score", color=text_color)
            ax.set_title("Batch Risk Scores", color=text_color)
            fig.patch.set_facecolor(plot_bg)
            ax.set_facecolor(plot_bg)
            ax.tick_params(colors=text_color)
            ax.yaxis.label.set_color(text_color)
            ax.legend(facecolor=plot_bg, labelcolor=text_color)
            plt.xticks(rotation=45, ha="right")
            st.pyplot(fig, use_container_width=False)
            plt.close()

# ── SESSION HISTORY ───────────────────────────────────────────────
elif mode == "Session History":
    st.subheader("Files Analyzed This Session")
    if st.session_state.session_history:
        df_hist = pd.DataFrame(st.session_state.session_history)
        st.dataframe(df_hist, use_container_width=True)
        st.subheader("Risk Score Timeline")
        fig, ax = plt.subplots(figsize=(7, 3))
        scores = [row["Score"] for row in st.session_state.session_history]
        files = [row["File"][:15] for row in st.session_state.session_history]
        colors = ["#ff4b4b" if s > 0.5 else "#4bff91" for s in scores]
        ax.bar(files, scores, color=colors)
        ax.axhline(0.5, color="yellow", linestyle="--", linewidth=1, label="Malware threshold")
        ax.set_ylim(0, 1)
        ax.set_ylabel("Malware Score", color=text_color)
        ax.set_title("Risk Score Across Session", color=text_color)
        fig.patch.set_facecolor(plot_bg)
        ax.set_facecolor(plot_bg)
        ax.tick_params(colors=text_color)
        ax.yaxis.label.set_color(text_color)
        ax.legend(facecolor=plot_bg, labelcolor=text_color)
        st.pyplot(fig, use_container_width=False)
        plt.close()
        if st.button("Clear History", key="clear_btn"):
            st.session_state.session_history = []
            st.rerun()
    else:
        st.info("No files analyzed yet. Go to Single File Analysis to get started.")