import os
import onnx
import gradio as gr

# ---------------- FIX ONNX LOCALE ISSUE ---------------- #

import onnx

def fix_onnx_locale_issue(model_path):
    fixed_path = model_path.replace(".onnx", "_fixed.onnx")

    if os.path.exists(fixed_path):
        return fixed_path

    model = onnx.load(model_path)

    nodes = model.graph.node
    new_nodes = []

    replace_map = {}

    # Step 1: find StringNormalizer
    for node in nodes:
        if node.op_type == "StringNormalizer":
            input_name = node.input[0]
            output_name = node.output[0]

            # map normalized → original input
            replace_map[output_name] = input_name
        else:
            new_nodes.append(node)

    # Step 2: fix inputs of remaining nodes
    for node in new_nodes:
        for i, inp in enumerate(node.input):
            if inp in replace_map:
                node.input[i] = replace_map[inp]

    # Step 3: replace graph nodes
    model.graph.ClearField("node")
    model.graph.node.extend(new_nodes)

    onnx.save(model, fixed_path)
    return fixed_path


# ---------------- PATCH MODEL BEFORE IMPORT ---------------- #

MODEL_DIR = "models"

binary_raw = os.path.join(MODEL_DIR, "binary_model.onnx")
category_raw = os.path.join(MODEL_DIR, "category_model.onnx")

# Fix models BEFORE model.py loads them
binary_fixed = fix_onnx_locale_issue(binary_raw)
category_fixed = fix_onnx_locale_issue(category_raw)

# Trick: overwrite original paths so model.py uses fixed ones
os.replace(binary_fixed, binary_raw)
os.replace(category_fixed, category_raw)


# ---------------- IMPORT MODEL ---------------- #

from model import run_model


# ---------------- UI LOGIC ---------------- #

CATEGORIES = {
    "phishing": "Phishing",
    "job_scam": "Job Scam",
    "crypto": "Cryptocurrency Scam",
    "adult": "Adult Content",
    "giveaway": "Giveaway Scam",
    "marketing": "Marketing",
    "spam": "Spam",
    "normal": "Normal",
}

CATEGORY_ICONS = {
    "phishing":  "⚠",
    "job_scam":  "⚠",
    "crypto":    "⚠",
    "adult":     "⚠",
    "giveaway":  "⚠",
    "marketing": "◈",
    "spam":      "⚠",
    "normal":    "✓",
}


def detect_spam(text):
    if not text or not text.strip():
        return "—", "—", 0.0, {}, "Awaiting input"

    try:
        output = run_model(text)

        if not output["ok"]:
            return output["error"], "—", 0.0, {}, "Error"

        pred = output["result"]

        is_spam = pred["is_spam"]
        confidence = pred["confidence"]
        category = pred["category"]

        result_label = "SPAM DETECTED" if is_spam else "CLEAN"
        category_label = CATEGORIES.get(category, category)

        details = {
            "category": category,
            "confidence": confidence,
            "threshold": pred.get("threshold_used"),
        }

        summary = f"{'SPAM' if is_spam else 'HAM'}  ·  {category_label}  ·  {confidence:.4f} confidence"

        return result_label, category_label, confidence, details, summary

    except Exception as e:
        return f"Error: {str(e)}", "—", 0.0, {"error": str(e)}, "Crash"


# ================================================================
#  CUSTOM CSS — full dark security dashboard aesthetic
#  Fonts: JetBrains Mono (data) + DM Sans (UI)
# ================================================================

CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;700&family=DM+Sans:wght@300;400;500;600&display=swap');

/* ── DESIGN TOKENS ─────────────────────────────────────────── */
:root {
  --ss-bg:          #080C12;
  --ss-surface-1:   #0D1320;
  --ss-surface-2:   #111926;
  --ss-border:      rgba(255,255,255,0.055);
  --ss-border-em:   rgba(255,255,255,0.12);

  --ss-safe:        #0CF2A0;
  --ss-safe-dim:    rgba(12,242,160,0.12);
  --ss-safe-glow:   rgba(12,242,160,0.18);

  --ss-danger:      #FF3D5A;
  --ss-danger-dim:  rgba(255,61,90,0.12);

  --ss-info:        #4D9FFF;
  --ss-info-dim:    rgba(77,159,255,0.10);

  --ss-warn:        #FFAA00;

  --ss-text-1:      #EAF0FF;
  --ss-text-2:      #8090AA;
  --ss-text-3:      #4A5568;

  --ss-mono:        'JetBrains Mono', 'Fira Code', monospace;
  --ss-sans:        'DM Sans', system-ui, sans-serif;

  --ss-radius-sm:   6px;
  --ss-radius-md:   10px;
  --ss-radius-lg:   14px;
  --ss-transition:  0.18s cubic-bezier(0.4, 0, 0.2, 1);
}

/* ── GLOBAL RESETS ─────────────────────────────────────────── */
*, *::before, *::after { box-sizing: border-box; }

body,
.gradio-container,
gradio-app {
  background: var(--ss-bg) !important;
  font-family: var(--ss-sans) !important;
  color: var(--ss-text-1) !important;
  min-height: 100vh;
}

.gradio-container {
  max-width: 960px !important;
  margin: 0 auto !important;
  padding: 0 20px 80px !important;
  background:
    radial-gradient(ellipse 70% 40% at 50% -5%,  rgba(77,159,255,0.07) 0%, transparent 65%),
    radial-gradient(ellipse 50% 30% at 90% 80%,  rgba(12,242,160,0.05) 0%, transparent 55%),
    var(--ss-bg) !important;
}

footer,
.footer,
.built-with { display: none !important; }

/* Strip all default Gradio card/panel chrome */
.gr-box,
.gr-form,
.gap,
.contain,
.panel,
.tabs,
.tabitem,
.block {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  padding: 0 !important;
}

/* ── CUSTOM HEADER ─────────────────────────────────────────── */
#ss-header {
  position: relative;
  text-align: center;
  padding: 52px 0 44px;
  overflow: hidden;
}

#ss-header::after {
  content: '';
  position: absolute;
  bottom: 0; left: 50%;
  transform: translateX(-50%);
  width: 100%; height: 1px;
  background: linear-gradient(90deg,
    transparent 0%,
    var(--ss-border-em) 25%,
    var(--ss-safe) 50%,
    var(--ss-border-em) 75%,
    transparent 100%
  );
  opacity: 0.6;
}

.ss-wordmark {
  display: inline-flex;
  align-items: center;
  gap: 14px;
  margin-bottom: 12px;
}

.ss-shield {
  width: 48px; height: 48px;
  display: flex; align-items: center; justify-content: center;
  background: var(--ss-safe-dim);
  border: 1px solid rgba(12,242,160,0.28);
  border-radius: 12px;
  font-size: 22px;
  animation: ss-pulse 3.5s ease-in-out infinite;
}

@keyframes ss-pulse {
  0%, 100% { box-shadow: 0 0 0 0 rgba(12,242,160,0); }
  50%       { box-shadow: 0 0 0 8px rgba(12,242,160,0.06); }
}

.ss-title {
  font-family: var(--ss-sans);
  font-size: 36px;
  font-weight: 600;
  letter-spacing: -0.5px;
  color: var(--ss-text-1);
  margin: 0;
  line-height: 1;
}

.ss-title span {
  color: var(--ss-safe);
}

.ss-tagline {
  font-family: var(--ss-mono);
  font-size: 11.5px;
  letter-spacing: 0.18em;
  text-transform: uppercase;
  color: var(--ss-text-2);
  margin: 10px 0 18px;
}

.ss-badges {
  display: inline-flex;
  gap: 10px;
  align-items: center;
  justify-content: center;
  flex-wrap: wrap;
}

.ss-badge {
  display: inline-flex;
  align-items: center;
  gap: 6px;
  padding: 5px 12px;
  background: var(--ss-surface-2);
  border: 1px solid var(--ss-border-em);
  border-radius: 100px;
  font-family: var(--ss-mono);
  font-size: 10.5px;
  letter-spacing: 0.08em;
  color: var(--ss-text-2);
}

.ss-badge.live::before {
  content: '';
  display: inline-block;
  width: 6px; height: 6px;
  border-radius: 50%;
  background: var(--ss-safe);
  animation: ss-blink 2s ease-in-out infinite;
}

@keyframes ss-blink {
  0%, 100% { opacity: 1; }
  50%       { opacity: 0.25; }
}

/* ── SECTION LABELS ────────────────────────────────────────── */
.ss-section-label {
  font-family: var(--ss-mono) !important;
  font-size: 10px !important;
  letter-spacing: 0.18em !important;
  text-transform: uppercase !important;
  color: var(--ss-text-3) !important;
  margin: 28px 0 10px !important;
  display: flex !important;
  align-items: center !important;
  gap: 10px !important;
}

.ss-section-label::after {
  content: '';
  flex: 1;
  height: 1px;
  background: var(--ss-border);
}

/* ── INPUT CARD ────────────────────────────────────────────── */
#ss-input-card {
  background: var(--ss-surface-1) !important;
  border: 1px solid var(--ss-border) !important;
  border-radius: var(--ss-radius-lg) !important;
  padding: 20px !important;
  transition: border-color var(--ss-transition) !important;
}

#ss-input-card:focus-within {
  border-color: rgba(77,159,255,0.3) !important;
}

/* ── TEXTBOX ───────────────────────────────────────────────── */
.ss-textarea textarea,
.ss-textarea input[type="text"] {
  background: rgba(0, 8, 18, 0.7) !important;
  border: 1px solid rgba(77,159,255,0.14) !important;
  border-radius: var(--ss-radius-md) !important;
  color: var(--ss-text-1) !important;
  font-family: var(--ss-mono) !important;
  font-size: 13px !important;
  line-height: 1.65 !important;
  padding: 14px 16px !important;
  resize: none !important;
  transition: border-color var(--ss-transition), box-shadow var(--ss-transition) !important;
  caret-color: var(--ss-safe) !important;
}

.ss-textarea textarea:focus,
.ss-textarea input[type="text"]:focus {
  border-color: rgba(77,159,255,0.45) !important;
  box-shadow: 0 0 0 3px rgba(77,159,255,0.07) !important;
  outline: none !important;
}

.ss-textarea textarea::placeholder {
  color: var(--ss-text-3) !important;
  font-style: italic;
}

.ss-textarea label,
.ss-textarea label span {
  font-family: var(--ss-mono) !important;
  font-size: 10px !important;
  letter-spacing: 0.16em !important;
  text-transform: uppercase !important;
  color: var(--ss-text-3) !important;
  margin-bottom: 8px !important;
}

/* ── BUTTONS ───────────────────────────────────────────────── */
#ss-analyze-btn button {
  background: var(--ss-safe-dim) !important;
  border: 1px solid rgba(12,242,160,0.35) !important;
  border-radius: var(--ss-radius-md) !important;
  color: var(--ss-safe) !important;
  font-family: var(--ss-mono) !important;
  font-size: 12px !important;
  font-weight: 700 !important;
  letter-spacing: 0.18em !important;
  text-transform: uppercase !important;
  padding: 13px 28px !important;
  width: 100% !important;
  cursor: pointer !important;
  transition: all var(--ss-transition) !important;
}

#ss-analyze-btn button:hover {
  background: rgba(12,242,160,0.2) !important;
  box-shadow: 0 0 20px rgba(12,242,160,0.14) !important;
  transform: translateY(-1px) !important;
}

#ss-analyze-btn button:active {
  transform: translateY(0) !important;
}

#ss-clear-btn button {
  background: transparent !important;
  border: 1px solid var(--ss-border-em) !important;
  border-radius: var(--ss-radius-md) !important;
  color: var(--ss-text-3) !important;
  font-family: var(--ss-mono) !important;
  font-size: 10px !important;
  letter-spacing: 0.14em !important;
  text-transform: uppercase !important;
  padding: 9px !important;
  width: 100% !important;
  cursor: pointer !important;
  transition: all var(--ss-transition) !important;
}

#ss-clear-btn button:hover {
  border-color: rgba(255,61,90,0.35) !important;
  color: var(--ss-danger) !important;
}

/* ── OUTPUT CARDS ──────────────────────────────────────────── */
.ss-result-card,
.ss-output-card {
  background: var(--ss-surface-1) !important;
  border: 1px solid var(--ss-border) !important;
  border-radius: var(--ss-radius-lg) !important;
  padding: 20px !important;
  transition: border-color var(--ss-transition) !important;
}

.ss-result-card:hover,
.ss-output-card:hover {
  border-color: var(--ss-border-em) !important;
}

/* Output field text */
.ss-output-field input,
.ss-output-field textarea {
  background: rgba(0, 8, 18, 0.6) !important;
  border: 1px solid rgba(77,159,255,0.10) !important;
  border-radius: var(--ss-radius-sm) !important;
  color: var(--ss-info) !important;
  font-family: var(--ss-mono) !important;
  font-size: 14px !important;
  font-weight: 700 !important;
  padding: 12px 14px !important;
  transition: all var(--ss-transition) !important;
  cursor: default !important;
}

.ss-result-field input,
.ss-result-field textarea {
  background: rgba(0, 8, 18, 0.6) !important;
  border: 1px solid rgba(12,242,160,0.14) !important;
  border-radius: var(--ss-radius-sm) !important;
  color: var(--ss-safe) !important;
  font-family: var(--ss-mono) !important;
  font-size: 15px !important;
  font-weight: 700 !important;
  padding: 13px 14px !important;
  letter-spacing: 0.08em !important;
  cursor: default !important;
}

.ss-number-field input[type="number"] {
  background: rgba(0, 8, 18, 0.6) !important;
  border: 1px solid rgba(77,159,255,0.10) !important;
  border-radius: var(--ss-radius-sm) !important;
  color: var(--ss-info) !important;
  font-family: var(--ss-mono) !important;
  font-size: 22px !important;
  font-weight: 700 !important;
  padding: 12px 14px !important;
  text-align: center !important;
  cursor: default !important;
}

.ss-summary-field input,
.ss-summary-field textarea {
  background: var(--ss-surface-2) !important;
  border: 1px solid var(--ss-border-em) !important;
  border-radius: var(--ss-radius-sm) !important;
  color: var(--ss-text-2) !important;
  font-family: var(--ss-mono) !important;
  font-size: 12px !important;
  letter-spacing: 0.06em !important;
  padding: 10px 14px !important;
  cursor: default !important;
}

/* All output labels */
.ss-output-field label span,
.ss-result-field label span,
.ss-number-field label span,
.ss-summary-field label span,
.ss-json-field label span {
  font-family: var(--ss-mono) !important;
  font-size: 10px !important;
  letter-spacing: 0.16em !important;
  text-transform: uppercase !important;
  color: var(--ss-text-3) !important;
  margin-bottom: 6px !important;
}

/* ── JSON COMPONENT ────────────────────────────────────────── */
.ss-json-field .json-component,
.ss-json-field [class*="json"],
.ss-json-field pre {
  background: rgba(0, 8, 18, 0.7) !important;
  border: 1px solid rgba(77,159,255,0.10) !important;
  border-radius: var(--ss-radius-sm) !important;
  font-family: var(--ss-mono) !important;
  font-size: 12px !important;
  color: var(--ss-text-2) !important;
  padding: 12px 14px !important;
}

/* ── EXAMPLES ──────────────────────────────────────────────── */
.gr-samples,
.gr-samples-table {
  background: var(--ss-surface-1) !important;
  border: 1px solid var(--ss-border) !important;
  border-radius: var(--ss-radius-lg) !important;
  overflow: hidden !important;
}

.gr-samples thead {
  background: rgba(77,159,255,0.05) !important;
  border-bottom: 1px solid var(--ss-border-em) !important;
}

.gr-samples th,
.gr-samples-table th {
  font-family: var(--ss-mono) !important;
  font-size: 9.5px !important;
  letter-spacing: 0.18em !important;
  text-transform: uppercase !important;
  color: var(--ss-text-3) !important;
  padding: 10px 16px !important;
  font-weight: 400 !important;
}

.gr-samples td,
.gr-samples-table td {
  font-family: var(--ss-mono) !important;
  font-size: 12px !important;
  color: var(--ss-text-2) !important;
  padding: 10px 16px !important;
  border-bottom: 1px solid var(--ss-border) !important;
  cursor: pointer !important;
  transition: all var(--ss-transition) !important;
}

.gr-samples tr:hover td,
.gr-samples-table tr:hover td {
  color: var(--ss-text-1) !important;
  background: rgba(77,159,255,0.04) !important;
}

.gr-samples tr:last-child td,
.gr-samples-table tr:last-child td {
  border-bottom: none !important;
}

/* ── SCROLLBAR ─────────────────────────────────────────────── */
::-webkit-scrollbar            { width: 4px; height: 4px; }
::-webkit-scrollbar-track      { background: transparent; }
::-webkit-scrollbar-thumb      { background: rgba(77,159,255,0.18); border-radius: 4px; }
::-webkit-scrollbar-thumb:hover{ background: rgba(77,159,255,0.35); }

/* ── DIVIDER ───────────────────────────────────────────────── */
#ss-divider {
  border: none !important;
  border-top: 1px solid var(--ss-border) !important;
  margin: 12px 0 !important;
}

/* ── MISC CLEANUP ──────────────────────────────────────────── */
.hide-border input,
.hide-border textarea {
  border: none !important;
  box-shadow: none !important;
}

/* Gradio number spinner buttons */
input[type=number]::-webkit-inner-spin-button,
input[type=number]::-webkit-outer-spin-button {
  -webkit-appearance: none;
  margin: 0;
}
"""


# ================================================================
#  PAGE LOAD JS
# ================================================================

PAGE_JS = """
function() {
  // Watch result field and tint it spam-red vs safe-green
  const observer = new MutationObserver(() => {
    const resultInputs = document.querySelectorAll('.ss-result-field input, .ss-result-field textarea');
    resultInputs.forEach(el => {
      const v = (el.value || el.innerText || '').trim().toUpperCase();
      if (v.includes('SPAM')) {
        el.style.color = '#FF3D5A';
        el.style.borderColor = 'rgba(255,61,90,0.25)';
      } else if (v.includes('CLEAN')) {
        el.style.color = '#0CF2A0';
        el.style.borderColor = 'rgba(12,242,160,0.25)';
      } else {
        el.style.color = '';
        el.style.borderColor = '';
      }
    });
  });
  observer.observe(document.body, { subtree: true, childList: true, characterData: true });
}
"""


# ================================================================
#  HEADER HTML
# ================================================================

HEADER_HTML = """
<div id="ss-header">
  <div class="ss-wordmark">
    <div class="ss-shield">🛡</div>
    <h1 class="ss-title">Spam<span>Shield</span></h1>
  </div>
  <p class="ss-tagline">Multilingual threat detection &nbsp;·&nbsp; ONNX runtime</p>
  <div class="ss-badges">
    <span class="ss-badge live">Model online</span>
    <span class="ss-badge">PyTorch · ~3M params</span>
    <span class="ss-badge">EN · ES · DE · RU · HI</span>
  </div>
</div>
"""

EXAMPLES = [
    "Congratulations! You've won a free iPhone. Click here to claim!",
    "Hey! How are you doing? Let's meet tomorrow.",
    "Your account has been compromised. Verify now.",
    "Earn $5000/week from home. No experience needed.",
    "Limited offer! Buy now!",
    "Free Bitcoin airdrop!",
]


# ================================================================
#  BUILD UI
# ================================================================

with gr.Blocks(css=CSS, js=PAGE_JS, title="SpamShield") as demo:

    # ── Header ──────────────────────────────────────────────────
    gr.HTML(HEADER_HTML)

    # ── Input section ───────────────────────────────────────────
    gr.HTML('<p class="ss-section-label">Input</p>')

    with gr.Group(elem_id="ss-input-card"):
        text_input = gr.Textbox(
            label="Message",
            lines=5,
            placeholder="Paste or type a message to analyze…",
            elem_classes=["ss-textarea"],
            show_label=True,
        )

        with gr.Row():
            with gr.Column(scale=3):
                analyze_btn = gr.Button(
                    "⟢  Analyze",
                    variant="primary",
                    elem_id="ss-analyze-btn",
                )
            with gr.Column(scale=1):
                clear_btn = gr.Button(
                    "Clear",
                    variant="secondary",
                    elem_id="ss-clear-btn",
                )

    # ── Results section ─────────────────────────────────────────
    gr.HTML('<p class="ss-section-label">Detection result</p>')

    with gr.Row():
        with gr.Column(scale=1, elem_classes=["ss-result-card"]):
            result = gr.Textbox(
                label="Verdict",
                interactive=False,
                elem_classes=["ss-result-field"],
            )
            category = gr.Textbox(
                label="Category",
                interactive=False,
                elem_classes=["ss-output-field"],
            )

        with gr.Column(scale=1, elem_classes=["ss-output-card"]):
            confidence = gr.Number(
                label="Confidence score",
                interactive=False,
                elem_classes=["ss-number-field"],
            )
            summary = gr.Textbox(
                label="Summary",
                interactive=False,
                elem_classes=["ss-summary-field"],
            )

    details = gr.JSON(
        label="Raw output",
        elem_classes=["ss-json-field"],
    )

    # ── Examples ────────────────────────────────────────────────
    gr.HTML('<p class="ss-section-label">Quick examples</p>')

    gr.Examples(
        examples=EXAMPLES,
        inputs=text_input,
        outputs=[result, category, confidence, details, summary],
        fn=detect_spam,
    )

    # ── Event bindings ───────────────────────────────────────────
    analyze_btn.click(
        detect_spam,
        inputs=text_input,
        outputs=[result, category, confidence, details, summary],
    )

    clear_btn.click(
        lambda: ("—", "—", 0.0, {}, "Awaiting input"),
        None,
        [result, category, confidence, details, summary],
    ).then(
        lambda: "",
        None,
        text_input,
    )


# ================================================================
#  LAUNCH
# ================================================================

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Base(
            primary_hue="emerald",
            secondary_hue="blue",
            neutral_hue="slate",
            font=[gr.themes.GoogleFont("DM Sans"), "system-ui", "sans-serif"],
            font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
        ),
        ssr_mode=False,
    )