"""
VulnScout — Gradio app (Hugging Face Space front-end).

Paste a Solidity contract/function → VulnScout encodes it on a **Modal GPU** with the
SCAR sparse retriever and surfaces the most similar *known-vulnerable* code, each with the
real auditor finding behind it. Every audit is live Modal inference (the judged artifact).

Requires (as Space secrets): MODAL_TOKEN_ID, MODAL_TOKEN_SECRET  → so this Space can call
the deployed `vulnscout` Modal app. Without them, the UI loads in a clearly-labelled demo
fallback instead of crashing.
"""

import os
import gradio as gr

MODAL_APP = "vulnscout"
MODAL_CLS = "VulnScoutEngine"

# --- Lazy Modal engine handle (cached) ---------------------------------------
_engine = None
_engine_err = None


def _get_engine():
    global _engine, _engine_err
    if _engine is not None or _engine_err is not None:
        return _engine
    try:
        import modal
        Engine = modal.Cls.from_name(MODAL_APP, MODAL_CLS)
        _engine = Engine()
    except Exception as e:  # missing tokens, app not deployed, etc. — degrade, don't crash
        _engine_err = str(e)
    return _engine


SEVERITY_COLORS = {
    "critical": "#b3001b", "high": "#d7263d", "medium": "#e08a00",
    "low": "#3a7ca5", "informational": "#6c757d", "info": "#6c757d",
}


def _severity_of(finding: str) -> str:
    head = (finding or "").strip().lower()
    for sev in ("critical", "high", "medium", "low", "informational", "info"):
        if head.startswith(sev) or f"severity: {sev}" in head:
            return sev
    return "info"


def _render(out: dict) -> str:
    if not out or "error" in out:
        return f"<p style='color:#b3001b'>⚠️ {out.get('error','no response') if out else 'no response'}</p>"

    m = (f"<div style='font-family:ui-monospace,monospace;font-size:13px;"
         f"background:#0f1117;color:#cfe;padding:10px 14px;border-radius:8px;margin-bottom:12px'>"
         f"⚡ <b>live on Modal {out['gpu']}</b> · encode {out['encode_ms']}ms · "
         f"retrieve {out['retrieve_ms']}ms · "
         f"<b>this audit ≈ ${out['est_cost_usd']:.6f}</b> · "
         f"index = {out['n_index']:,} known vulns · "
         f"{out['n_active']} sparse features active</div>")

    cards = []
    for i, r in enumerate(out.get("results", []), 1):
        sev = _severity_of(r["finding"])
        color = SEVERITY_COLORS.get(sev, "#6c757d")
        code = r["code"] if len(r["code"]) <= 1100 else r["code"][:1100] + "\n… (truncated)"
        cards.append(
            f"<div style='border:1px solid #2a2f3a;border-left:5px solid {color};"
            f"border-radius:8px;padding:12px 14px;margin-bottom:12px;background:#161922'>"
            f"<div style='display:flex;justify-content:space-between;align-items:center'>"
            f"<span style='color:{color};font-weight:700;text-transform:uppercase;font-size:12px'>"
            f"{sev}</span>"
            f"<span style='color:#8a93a6;font-size:12px'>match #{i} · "
            f"similarity {r['score']:.3f} · source: {r['source']}</span></div>"
            f"<div style='margin:8px 0;color:#dfe6f0'>{_esc(r['finding'])}</div>"
            f"<pre style='background:#0d0f15;color:#c9d4e5;padding:10px;border-radius:6px;"
            f"overflow-x:auto;font-size:12px;margin:0'><code>{_esc(code)}</code></pre>"
            f"</div>"
        )
    if not cards:
        cards = ["<p>No similar known vulnerabilities found.</p>"]

    feats = out.get("active_features", [])
    feat_html = ""
    if feats:
        chips = " ".join(
            f"<span style='display:inline-block;background:#1d2330;color:#9ecbff;"
            f"border:1px solid #2a3550;border-radius:12px;padding:2px 9px;margin:2px;font-size:11px;"
            f"font-family:ui-monospace,monospace'>#{f['feature']} · {f['weight']:.2f}</span>"
            for f in feats)
        feat_html = (f"<details style='margin-top:6px'><summary style='cursor:pointer;color:#8a93a6;"
                     f"font-size:12px'>🔍 Why these matches — top sparse SAE features that fired "
                     f"(interpretable signal)</summary><div style='margin-top:8px'>{chips}</div></details>")

    return m + "".join(cards) + feat_html


def _esc(s: str) -> str:
    return (s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;"))


WAKING = ("<div style='font-family:ui-monospace,monospace;font-size:13px;"
          "background:#0f1117;color:#9ecbff;padding:14px 16px;border-radius:8px;"
          "border:1px solid #1e2737'>⏳ <b>Waking the Modal GPU…</b> the first audit after idle "
          "spins up an A10G from scale-to-zero (~30s). After that, audits are ~0.5s. "
          "Encoding your contract live on Modal…</div>")


def audit(contract: str, top_k: int):
    engine = _get_engine()
    if engine is None:
        yield ("<div style='padding:14px;border:1px solid #d7263d;border-radius:8px'>"
               "⚠️ <b>Modal backend not connected.</b> Set <code>MODAL_TOKEN_ID</code> and "
               "<code>MODAL_TOKEN_SECRET</code> as Space secrets and deploy "
               "<code>modal_app.py</code>. See the demo video in the README for a live run."
               f"<br><span style='color:#8a93a6;font-size:12px'>({_engine_err})</span></div>")
        return
    if not (contract or "").strip():
        yield "<p style='color:#e08a00'>Paste a Solidity contract or function to audit.</p>"
        return
    yield WAKING
    try:
        out = engine.audit.remote(contract, int(top_k))
    except Exception as e:
        yield f"<p style='color:#b3001b'>⚠️ Modal call failed: {_esc(str(e))}</p>"
        return
    yield _render(out)


EXAMPLES = [
    ["""// Arbitrary delegatecall: attacker-controlled target runs in this contract's context
function execute(address target, bytes calldata data) external {
    (bool ok, ) = target.delegatecall(data);   // <-- no access control / allowlist
    require(ok, "call failed");
}""", 5],
    ["""// Unchecked return value of a low-level transfer
function pay(address payable to, uint256 amt) external {
    to.send(amt);                     // <-- return value ignored
}""", 5],
    ["""// Reentrancy: external call before state update
function withdraw(uint256 amount) public {
    require(balances[msg.sender] >= amount, "insufficient");
    (bool ok, ) = msg.sender.call{value: amount}("");
    require(ok, "transfer failed");
    balances[msg.sender] -= amount;   // <-- updated AFTER the external call
}""", 5],
]

HERO = """
<div style="text-align:center;padding:30px 18px 6px">
  <div style="font-size:12px;letter-spacing:.30em;color:#22d3ee;font-weight:700;text-transform:uppercase">
    Build Small &nbsp;·&nbsp; powered by Modal</div>
  <div style="font-size:48px;font-weight:800;margin:8px 0 2px;line-height:1.05;
    background:linear-gradient(90deg,#e2e8f0 10%,#22d3ee 60%,#6366f1 100%);
    -webkit-background-clip:text;background-clip:text;color:transparent">🛡️ VulnScout</div>
  <div style="font-size:17px;color:#94a3b8;max-width:700px;margin:8px auto 0;line-height:1.5">
    Sparse, interpretable vulnerability <b style="color:#cbd5e1">precedent search</b> for Solidity.
    Paste a contract — get the closest <b style="color:#cbd5e1">known-exploited</b> code and the
    auditor's finding behind it, encoded <b style="color:#22d3ee">live on a Modal GPU</b>.</div>
</div>
"""

STATS = """
<div style="display:flex;gap:12px;justify-content:center;flex-wrap:wrap;margin:14px 0 6px">
  <div class="vs-stat"><div class="vs-num">7,552</div><div class="vs-lbl">known vulns indexed</div></div>
  <div class="vs-stat"><div class="vs-num">1.5B</div><div class="vs-lbl">params · laptop-class</div></div>
  <div class="vs-stat"><div class="vs-num">~10ms</div><div class="vs-lbl">sparse retrieval</div></div>
  <div class="vs-stat"><div class="vs-num">$0.0002</div><div class="vs-lbl">per live audit</div></div>
</div>
"""

MODAL_CARD = """
<div style="border:1px solid #1e2737;border-radius:12px;padding:14px 16px;
  background:linear-gradient(180deg,#0e1420,#0b0f17);margin-top:12px">
  <div style="font-size:11px;letter-spacing:.18em;text-transform:uppercase;color:#22d3ee;
    font-weight:700;margin-bottom:8px">⚡ Modal, three ways</div>
  <div style="color:#cbd5e1;font-size:13.5px;line-height:1.95">
    🏋️ <b>Trained</b> on Modal H100s · ~$280<br>
    📦 <b>Indexed</b> the corpus on Modal · Batch<br>
    ⚡ <b>This audit</b> runs on Modal · live</div>
</div>
"""

EMPTY = """
<div style="text-align:center;padding:40px 16px;border:1px dashed #1e2737;
  border-radius:14px;background:#0b0f17;color:#64748b">
  <div style="font-size:34px">🔍</div>
  <div style="margin-top:8px;font-size:15px">Paste a Solidity contract or pick an example, then
    <b style="color:#94a3b8">Scout for vulnerabilities</b>.</div>
</div>
"""

CSS = """
.gradio-container {max-width: 1100px !important; margin: 0 auto !important;}
footer {display: none !important;}
.vs-stat {background:#0e1420; border:1px solid #1e2737; border-radius:12px;
  padding:10px 18px; min-width:130px; text-align:center;}
.vs-num {font-size:21px; font-weight:800; color:#e2e8f0;
  font-family:'JetBrains Mono',ui-monospace,monospace;}
.vs-lbl {font-size:11px; color:#7c8aa0; margin-top:3px; letter-spacing:.02em;}
.scout-btn, .scout-btn button {width:100% !important; font-size:16px !important;
  font-weight:700 !important; letter-spacing:.01em; padding:14px 16px !important;
  box-shadow:0 8px 26px rgba(34,211,238,.22) !important;}
.code-card .cm-editor, .code-card textarea {font-family:'JetBrains Mono',ui-monospace,monospace !important;
  font-size:13px !important;}
.gradio-container thead th {background:#0e1420 !important; color:#cbd5e1 !important; border-color:#1e2737 !important;}
.gradio-container tbody td {background:#0c1018 !important; color:#c0ccdb !important; border-color:#1b2230 !important;}
.gradio-container tbody tr:hover td {background:#11161f !important;}
"""

theme = gr.themes.Base(
    primary_hue=gr.themes.colors.cyan,
    secondary_hue=gr.themes.colors.indigo,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="#080b11",
    body_text_color="#cbd5e1",
    block_background_fill="#0c1018",
    block_border_color="#1b2230",
    block_label_text_color="#8a97ad",
    input_background_fill="#0a0e15",
    button_primary_background_fill="linear-gradient(90deg,#22d3ee,#6366f1)",
    button_primary_background_fill_hover="linear-gradient(90deg,#3edcf5,#7c7ff5)",
    button_primary_text_color="#06121a",
    slider_color="#22d3ee",
)

with gr.Blocks(title="VulnScout — Solidity vulnerability scout", css=CSS, theme=theme) as demo:
    gr.HTML(HERO)
    gr.HTML(STATS)
    with gr.Row():
        with gr.Column(scale=3):
            inp = gr.Code(label="Solidity contract / function", language=None,
                          lines=16, elem_classes="code-card")
        with gr.Column(scale=2):
            top_k = gr.Slider(1, 10, value=5, step=1, label="Matches to return")
            btn = gr.Button("🔎  Scout for vulnerabilities", variant="primary",
                            elem_classes="scout-btn")
            gr.HTML(MODAL_CARD)
    out = gr.HTML(EMPTY)
    gr.HTML("<div style='color:#7c8aa0;font-size:13px;font-weight:600;margin:18px 0 -4px'>⚡ Try an example</div>")
    gr.Examples(examples=EXAMPLES, inputs=[inp, top_k])
    gr.HTML("<div style='text-align:center;color:#475569;font-size:12px;margin:22px 0 6px'>"
            "VulnScout finds precedent — it does not replace a professional audit. "
            "· Built on SCAR · Compute by Modal</div>")
    btn.click(audit, inputs=[inp, top_k], outputs=out)

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