"""ColdScore — cold-email deliverability scoring on an open-weights model.

Hugging Face Space. Loads a fine-tuned ModernBERT from the Hub and runs it in
this container. No proprietary model API is called.

Set MODEL_ID to your Hub repo before deploying.
"""

import csv
import io
import os
import re

import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_ID = os.environ.get("MODEL_ID", "Ashsinha1/coldscore-modernbert")
MAX_LEN = 512

# Lazy-load: the Space is deployed before the model is pushed, and loading at
# import time would crash-loop the container until the repo exists. This way the
# UI boots immediately and picks the model up on the first request after it lands.
_M = {"tok": None, "model": None, "error": None}


def get_model():
    if _M["model"] is None:
        try:
            _M["tok"] = AutoTokenizer.from_pretrained(MODEL_ID)
            m = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
            m.eval()
            _M["model"], _M["error"] = m, None
        except Exception as e:  # repo missing, still uploading, or gated
            _M["error"] = (
                f"Model `{MODEL_ID}` isn't loadable yet.\n\n"
                f"Train and push it first (see the Colab one-liner), then hit **Score it** "
                f"again — no redeploy needed.\n\n<sub>{type(e).__name__}: {e}</sub>"
            )
    return _M

GRADES = [
    (0.15, "A", "Should land in the primary inbox."),
    (0.35, "B", "Likely to deliver. A couple of things worth tightening."),
    (0.60, "C", "Coin flip. The promotions tab is a real possibility."),
    (0.85, "D", "High spam risk. Rewrite before sending at volume."),
    (1.01, "F", "Very high spam risk. Sending this will hurt your domain reputation."),
]


def grade_for(risk):
    for cutoff, letter, verdict in GRADES:
        if risk < cutoff:
            return letter, verdict
    return "F", GRADES[-1][2]


@torch.inference_mode()
def score_many(pairs):
    """Score a list of (subject, body). Batched — one forward pass per chunk."""
    M = get_model()
    if M["model"] is None:
        raise RuntimeError(M["error"])
    tokenizer, model = M["tok"], M["model"]
    out = []
    for i in range(0, len(pairs), 16):
        chunk = pairs[i:i + 16]
        enc = tokenizer([s for s, _ in chunk], [b for _, b in chunk],
                        truncation=True, padding=True, max_length=MAX_LEN,
                        return_tensors="pt")
        probs = model(**enc).logits.float().softmax(-1)[:, 1]
        out.extend(float(p) for p in probs)
    return out


# --- fix suggestions -------------------------------------------------------
# These are deterministic checks against published deliverability guidance, NOT
# the model's reasoning. ModernBERT can't tell you why it scored something —
# claiming otherwise would be inventing an explanation. Kept separate on purpose.

URL_RE = re.compile(r"https?://\S+|www\.\S+")
MERGE_RE = re.compile(r"\{\{?\s*[a-z_]+\s*\}?\}", re.I)


def suggestions(subject, body):
    tips = []
    full = f"{subject}\n{body}"
    letters = [c for c in full if c.isalpha()]
    caps = sum(1 for c in letters if c.isupper()) / len(letters) if letters else 0

    if MERGE_RE.search(full):
        tips.append("**Unrendered merge tag** — something like `{{first_name}}` will send literally. "
                    "This is the single most damaging thing in the list.")
    if caps > 0.3:
        tips.append(f"**{caps:.0%} of letters are capitals.** Shouting is a classic filter trigger.")
    if re.search(r"\b[A-Z]{4,}\b", full):
        tips.append("**ALL-CAPS words.** Drop them to sentence case.")
    n_ex = full.count("!")
    if n_ex > 2:
        tips.append(f"**{n_ex} exclamation marks.** One is plenty; zero is usually better.")
    links = len(URL_RE.findall(body))
    if links > 2:
        tips.append(f"**{links} links.** Multiple links in a first-touch email reads as bulk mail. Cut to one.")
    if re.match(r"^\s*(re|fwd?)\s*:", subject, re.I):
        tips.append("**Fake `Re:`/`Fwd:` prefix.** It implies a conversation that never happened — "
                    "recipients mark it as spam, and that costs you domain reputation.")
    if len(body.split()) < 25:
        tips.append("**Very short body.** Too little content to look like a real 1:1 email.")
    if len(body.split()) > 220:
        tips.append(f"**{len(body.split())} words.** Long first-touch emails lose replies; aim for 90–130.")
    if len(subject) > 60:
        tips.append(f"**{len(subject)}-character subject.** It'll be truncated on mobile. Aim under 50.")
    if "unsubscribe" not in body.lower():
        tips.append("_No opt-out._ Not required for genuine 1:1 mail, but at volume its absence "
                    "raises complaint rates — which is what actually burns a domain.")

    if not tips:
        return "No mechanical issues found. Anything left is about the writing itself."
    return "\n\n".join("- " + t for t in tips)


def score_single(subject, body):
    subject = (subject or "").strip()
    body = (body or "").strip()
    if not subject and not body:
        return {}, "Enter a subject and a body.", ""

    try:
        risk = score_many([(subject, body)])[0]
    except RuntimeError as e:
        return {}, str(e), ""
    letter, verdict = grade_for(risk)
    headline = f"### Grade {letter} — {risk:.1%} spam risk\n{verdict}"
    return ({"spam-prone": risk, "deliverable": 1 - risk}, headline, suggestions(subject, body))


# --- batch -----------------------------------------------------------------

def score_batch(file, pasted):
    """CSV with subject,body columns — or one email per line as 'subject | body'."""
    pairs, errors = [], []

    if file is not None:
        try:
            with open(file.name, encoding="utf-8-sig", newline="") as fh:
                reader = csv.DictReader(fh)
                cols = {c.lower().strip(): c for c in (reader.fieldnames or [])}
                if "subject" not in cols or "body" not in cols:
                    return [], f"CSV needs `subject` and `body` columns. Found: {reader.fieldnames}"
                for row in reader:
                    pairs.append(((row[cols["subject"]] or "").strip(),
                                  (row[cols["body"]] or "").strip()))
        except Exception as e:
            return [], f"Could not read that CSV: {e}"

    if pasted and pasted.strip():
        for line in pasted.strip().splitlines():
            if not line.strip():
                continue
            subject, _, body = line.partition("|")
            pairs.append((subject.strip(), body.strip()))

    if not pairs:
        return [], "Upload a CSV or paste some emails first."

    pairs = pairs[:200]
    try:
        risks = score_many(pairs)
    except RuntimeError as e:
        return [], str(e)
    rows = []
    for (s, b), r in zip(pairs, risks):
        letter, _ = grade_for(r)
        rows.append([letter, f"{r:.1%}", s[:70], f"{len(b.split())}w"])
    rows.sort(key=lambda x: -float(x[1].rstrip("%")))

    risky = sum(1 for r in risks if r >= 0.6)
    note = (f"Scored **{len(pairs)}** emails. **{risky}** at grade C or worse "
            f"({risky / len(pairs):.0%}). Sorted worst first.")
    if len(pairs) == 200:
        note += "\n\n_Capped at 200 rows for this demo._"
    return rows, note


GOOD = ("quick question about Northwind's onboarding",
        "Hi Sarah,\n\nI noticed Northwind shipped a self-serve trial last month. Curious how "
        "you're handling the handoff from trial to sales conversation right now.\n\nWe built "
        "something for that and I'd rather hear how you do it before I assume anything.\n\n"
        "Worth 15 minutes next week?\n\nAshish")
BAD = ("ACT NOW!!! Northwind qualifies for FREE revenue growth!!!",
       "Dear Sir/Madam,\n\nCONGRATULATIONS!!! Northwind has been SELECTED for our EXCLUSIVE "
       "offer!\n\nGUARANTEED 10X ROI or your money back!!! This is a LIMITED TIME offer that "
       "EXPIRES at midnight!\n\nCLICK HERE NOW: http://bit.ly/xy12\n\nDON'T MISS OUT!!! ACT FAST!!!")
MERGE = ("{{first_name}} - I have an amazing offer for {{company}}",
         "Hi {{first_name}},\n\nWe are the best in the world at what we do and our results are "
         "guaranteed.\n\nSign up free today! Click here to buy now: http://example.com/signup\n\n"
         "Best,\nSales Team")

with gr.Blocks(title="ColdScore", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
        # 🧲 ColdScore
        ### Will this cold email reach the inbox — or the spam folder?

        A fine-tuned **[ModernBERT](https://huggingface.co/answerdotai/ModernBERT-base)**
        (Apache-2.0, 149M params) running on a free CPU. Open weights, MIT licence.
        **No proprietary model API is called** — the weights are public, and you can run
        this identical model inside your own network.
        """
    )

    with gr.Tab("Score an email"):
        with gr.Row():
            with gr.Column(scale=3):
                subject = gr.Textbox(label="Subject", placeholder="quick question about onboarding")
                body = gr.Textbox(label="Body", lines=12, placeholder="Hi Sarah,\n\nI noticed...")
                btn = gr.Button("Score it", variant="primary", size="lg")
            with gr.Column(scale=2):
                headline = gr.Markdown()
                label = gr.Label(label="Model output", num_top_classes=2)

        gr.Markdown("#### What to fix")
        tips = gr.Markdown()
        gr.Markdown(
            "_The fix list is deterministic checks against published deliverability guidance — "
            "not the model's reasoning. A 149M-parameter transformer can't tell you why it "
            "scored something, and inventing an explanation would be worse than not having one._"
        )

        gr.Examples(examples=[list(GOOD), list(BAD), list(MERGE)],
                    inputs=[subject, body], label="Try one")

        btn.click(score_single, [subject, body], [label, headline, tips])

    with gr.Tab("Score a campaign"):
        gr.Markdown(
            "Score a whole list at once — the actual job if you're sending at volume.\n\n"
            "Upload a **CSV with `subject` and `body` columns**, or paste one email per line "
            "as `subject | body`."
        )
        with gr.Row():
            up = gr.File(label="CSV", file_types=[".csv"])
            paste = gr.Textbox(label="Or paste", lines=6,
                               placeholder="quick question about onboarding | Hi Sarah, I noticed...")
        bbtn = gr.Button("Score all", variant="primary")
        bnote = gr.Markdown()
        table = gr.Dataframe(headers=["Grade", "Risk", "Subject", "Length"],
                             datatype=["str", "str", "str", "str"],
                             wrap=True, label="Worst first")
        bbtn.click(score_batch, [up, paste], [table, bnote])

    with gr.Tab("Use it yourself"):
        gr.Markdown(
            f"""
            The model is public. Run it wherever your data is allowed to live:

            ```python
            from transformers import pipeline

            clf = pipeline("text-classification", model="{MODEL_ID}")
            clf({{"text": "quick question about onboarding", "text_pair": "Hi Sarah, I noticed..."}})
            # [{{'label': 'deliverable', 'score': 0.98}}]
            ```

            ### Why an open model instead of a frontier API

            Scoring a cold email means handing over a prospect's name, their email address, and
            your pitch. That's customer data. Sending it to a third-party model provider is a
            decision that shows up in your customers' security review.

            This model doesn't force that decision: the weights are public, inference is a CPU
            process, and it runs inside your own network with the four lines above. A frontier
            model would score these emails better — but for the half of the pipeline that touches
            every contact in the database, "good and local" beats "excellent and elsewhere".

            ### Limitations — read before trusting a number

            1. **Content only.** Inbox placement is dominated by things the model cannot see:
               domain and IP reputation, SPF/DKIM/DMARC alignment, sending volume, list hygiene,
               complaint rates. A perfect email from a burned domain still lands in spam.
            2. **The benchmark is easy.** Trained on 100 hand-written emails whose spam examples
               are unsubtle. A keyword blocklist already scores ~0.94 F1 on that corpus, so
               held-out numbers are an optimistic upper bound, not real-world accuracy.
            3. **Labels are judgements, not outcomes.** Nobody sent these emails and measured
               where they landed.
            4. **English only.** Not a phishing detector — well-written phishing scores clean.

            Treat it as a content lint, not a deliverability guarantee.

            ### Privacy

            Text you enter is processed in this container and is not stored or logged. If that
            still isn't good enough for your data — which is a reasonable position, and rather
            the point of the project — clone the Space or run the model locally.
            """
        )

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