import os, re, time, json, urllib.parse
import gradio as gr
import torch
import torch.nn.functional as F

# Optional robust domain parsing; code falls back if missing.
try:
    import tldextract
except Exception:
    tldextract = None

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

URL_MODEL_ID = "CrabInHoney/urlbert-tiny-v4-malicious-url-classifier"

# Force readable labels regardless of model config
ID2LABEL = {0: "benign", 1: "defacement", 2: "malware", 3: "phishing"}

URL_RE = re.compile(r"""(?xi)\b(?:https?://|www\.)[^\s<>"'()]+""")

KEYWORDS = {
    "phish","login","verify","account","secure","update","bank","wallet",
    "password","invoice","pay","reset","support","unlock","confirm"
}
SUSPICIOUS_TLDS = {
    "zip","mov","lol","xyz","top","country","link","click","cam","help",
    "gq","cf","tk","work","rest","monster","quest","live","io","ly"
}
URL_SHORTENERS = {
    "bit.ly","tinyurl.com","t.co","goo.gl","is.gd","buff.ly","ow.ly","rebrand.ly","cutt.ly"
}

_tok = None
_mdl = None

# ---------- utils ----------
def _extract_urls(text: str):
    raw = [m.group(0).strip() for m in URL_RE.finditer(text or "")]
    cleaned = []
    for u in raw:
        u = u.rstrip(").,;:!?•]}>\"'")
        cleaned.append(u)
    return sorted(set(cleaned))

def _load_model():
    global _tok, _mdl
    if _tok is not None and _mdl is not None:
        return _tok, _mdl
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    _tok = AutoTokenizer.from_pretrained(URL_MODEL_ID)
    _mdl = AutoModelForSequenceClassification.from_pretrained(URL_MODEL_ID)
    _mdl.eval()
    return _tok, _mdl

def _softmax(logits: torch.Tensor):
    return F.softmax(logits, dim=-1).tolist()

def _results_table(rows):
    lines = [
        "| URL | Model | Model Prob (%) | Heuristic | Fused Risk | Decision | Reasons |",
        "|---|---|---:|---:|---:|:--:|---|",
    ]
    for r in rows:
        u, lbl, pct, h, fused, decision, reasons = r
        lines.append(
            f"| `{u}` | **{lbl}** | {pct:.2f} | {h:.2f} | {fused:.2f} | {decision} | {reasons} |"
        )
    return "\n".join(lines)

def _forensic_block(url, token_ids, tokens, scores_sorted, cls_vec, elapsed_s, truncated):
    toks_prev = ", ".join(tokens[:64]) + (" …" if len(tokens) > 64 else "")
    ids_prev  = ", ".join(map(str, token_ids[:64])) + (" …" if len(token_ids) > 64 else "")
    cls_dim = len(cls_vec)
    cls_prev = ", ".join(f"{v:.4f}" for v in cls_vec[:16]) + (" …" if cls_dim > 16 else "")
    l2 = (sum(v*v for v in cls_vec)) ** 0.5
    md = []
    md.append(f"### 🔍 Forensics for `{url}`\n")
    md.append(f"- tokens: **{len(tokens)}** • truncated: **{'yes' if truncated else 'no'}**")
    md.append(f"- inference time: **{elapsed_s:.2f}s**\n")
    md.append("**Top-k scores**")
    md.append("| Class | Prob (%) | Logit |\n|---|---:|---:|")
    for s in scores_sorted:
        md.append(f"| **{s['label']}** | {s['prob']*100:.2f} | {s['logit']:.3f} |")
    md.append("\n**Token IDs (preview)**")
    md.append("```txt\n" + ids_prev + "\n```")
    md.append("**Tokens (preview)**")
    md.append("```txt\n" + toks_prev + "\n```")
    md.append("**[CLS] embedding (preview)**")
    md.append(f"`dim={cls_dim}`, `L2={l2:.4f}`")
    md.append("```txt\n" + cls_prev + "\n```")
    return "\n".join(md)

# ---------- heuristics ----------
def _safe_parse(url: str):
    if not re.match(r"^https?://", url, re.I):
        url = "http://" + url
    return urllib.parse.urlparse(url)

def _split_reg_domain(host: str):
    parts = host.split(".")
    if len(parts) >= 2:
        return parts[-2] + "." + parts[-1]
    return host

def _domain_parts(host: str):
    if tldextract:
        ext = tldextract.extract(host)  # subdomain, domain, suffix
        regdom = f"{ext.domain}.{ext.suffix}" if ext.domain and ext.suffix else host
        sub = ext.subdomain or ""
        tld = ext.suffix or ""
        core = ext.domain or ""
    else:
        regdom = _split_reg_domain(host)
        tld = regdom.split(".")[-1] if "." in regdom else ""
        sub = host[:-len(regdom)].rstrip(".") if host.endswith(regdom) else ""
        core = regdom.split(".")[0] if "." in regdom else regdom
    return regdom, sub, core, tld

def heuristic_features(u: str):
    feats = {}
    try:
        p = _safe_parse(u)
        feats["host"] = p.hostname or ""
        feats["path"] = p.path or "/"
        feats["query"] = p.query or ""
        regdom, sub, core, tld = _domain_parts(feats["host"])
        feats["registered_domain"] = regdom
        feats["subdomain"] = sub
        feats["tld"] = tld
        feats["labels"] = feats["host"].count(".") + (1 if feats["host"] else 0)
        feats["has_at"] = "@" in u
        feats["has_port"] = bool(p.netloc and ":" in p.netloc.split("@")[-1])
        feats["has_punycode"] = "xn--" in feats["host"]
        feats["len_url"] = len(u)
        feats["hyphen_in_regdom"] = "-" in (core or "")
        low_host = feats["host"].lower()
        low_path = feats["path"].lower()
        feats["kw_in_path"] = int(any(k in low_path for k in KEYWORDS))
        feats["kw_in_host"] = int(any(k in low_host for k in KEYWORDS))
        feats["kw_in_subdomain_only"] = int(
            feats["kw_in_host"] and (core and not any(k in (core.lower()) for k in KEYWORDS))
        )
        feats["suspicious_tld"] = int((feats["tld"].split(".")[-1] or "") in SUSPICIOUS_TLDS)
        feats["is_shortener"] = int(regdom.lower() in URL_SHORTENERS)
        alnum = sum(c.isalnum() for c in feats["query"])
        feats["query_ratio_alnum"] = (alnum / max(1, len(feats["query"]))) if feats["query"] else 0.0
        feats["parse_error"] = False
    except Exception:
        feats = {"parse_error": True}
    return feats

def heuristic_score(feats: dict) -> float:
    if feats.get("parse_error"):
        return 0.80
    s = 0.0
    s += 0.28 * feats["kw_in_path"]
    s += 0.24 * feats["kw_in_subdomain_only"]
    s += 0.10 * feats["kw_in_host"]
    s += 0.12 * feats["hyphen_in_regdom"]
    s += 0.10 * (feats["labels"] >= 4)
    s += 0.10 * feats["has_punycode"]
    s += 0.12 * feats["suspicious_tld"]
    s += 0.10 * feats["is_shortener"]
    s += 0.05 * feats["has_at"]
    s += 0.05 * feats["has_port"]
    s += 0.10 * (feats["len_url"] >= 100)
    if feats.get("query") and len(feats.get("query", "")) >= 40 and feats.get("query_ratio_alnum", 0) > 0.9:
        s += 0.10
    return max(0.0, min(1.0, s))

def heuristic_reasons(feats: dict) -> str:
    if feats.get("parse_error"):
        return "parse error"
    rs = []
    if feats.get("is_shortener"): rs.append("URL shortener")
    if feats.get("kw_in_path"): rs.append("keyword in path")
    if feats.get("kw_in_subdomain_only"): rs.append("keyword in subdomain")
    if feats.get("kw_in_host") and not feats.get("kw_in_subdomain_only"): rs.append("keyword in host")
    if feats.get("hyphen_in_regdom"): rs.append("hyphen in registered domain")
    if feats.get("labels", 0) >= 4: rs.append("deep subdomain nesting")
    if feats.get("has_punycode"): rs.append("punycode host")
    if feats.get("suspicious_tld"): rs.append(f"suspicious TLD: {feats.get('tld')}")
    if feats.get("has_at"): rs.append("@ in URL")
    if feats.get("has_port"): rs.append("explicit port")
    if feats.get("len_url", 0) >= 100: rs.append("very long URL")   # ✅ fixed
    if feats.get("query") and len(feats.get("query", "")) >= 40 and feats.get("query_ratio_alnum", 0) > 0.9:
        rs.append("long query blob")
    return ", ".join(rs) if rs else "no heuristic triggers"

def heuristic_hard_flag(feats: dict) -> (bool, str):
    if feats.get("parse_error"):
        return True, "unparsable URL"
    if feats.get("kw_in_subdomain_only") and feats.get("kw_in_path"):
        return True, "keyword in subdomain + keyword in path"
    if feats.get("is_shortener") and (feats.get("kw_in_host") or feats.get("kw_in_path")):
        return True, "URL shortener + keyword"
    if feats.get("suspicious_tld") and (feats.get("kw_in_host") or feats.get("kw_in_path")):
        return True, "suspicious TLD + keyword"
    if feats.get("labels", 0) >= 4 and (feats.get("kw_in_host") or feats.get("kw_in_path")):
        return True, "deep subdomain nesting + keyword"
    return False, ""

# ---------- core ----------
def _parse_allowlist(s: str):
    items = re.split(r"[,\s]+", (s or "").strip())
    return {x.strip().lower() for x in items if x.strip()}

def analyze(
    text: str,
    forensic: bool,
    show_json: bool,
    threshold: float,
    allowlist_txt: str,
    allowlist_override: bool
):
    """
    One Markdown output:
      - verdict + table (model, heuristic, fused + decision + reasons)
      - optional forensic blocks
      - optional raw JSON
    """
    text = (text or "").strip()
    if not text:
        return "Paste an email body or a URL."

    urls = [text] if (text.lower().startswith(("http://","https://","www.")) and " " not in text) else _extract_urls(text)
    if not urls:
        return "No URLs detected in the text."

    allowset = _parse_allowlist(allowlist_txt)

    tok, mdl = _load_model()

    rows = []
    forensic_blocks = []
    export_data = {"model_id": URL_MODEL_ID, "items": []}
    any_unsafe = False

    for u in urls:
        # model forward
        max_len = min(512, getattr(mdl.config, "max_position_embeddings", 512) or 512)
        enc = tok(u, truncation=True, max_length=max_len, return_tensors="pt", return_attention_mask=True)
        token_ids = enc["input_ids"][0].tolist()
        tokens = tok.convert_ids_to_tokens(enc["input_ids"][0])
        truncated = enc["input_ids"].shape[1] >= max_len and len(tokens) >= max_len

        t0 = time.time()
        with torch.no_grad():
            out = mdl(**enc, output_hidden_states=True)
        elapsed = time.time() - t0

        logits = out.logits.squeeze(0)
        probs  = _softmax(logits)
        scores = [{"label": ID2LABEL[i], "prob": float(probs[i]), "logit": float(logits[i])}
                  for i in range(len(probs))]
        scores_sorted = sorted(scores, key=lambda x: x["prob"], reverse=True)
        top = scores_sorted[0]

        # heuristics
        feats = heuristic_features(u)
        regdom = feats.get("registered_domain", "").lower()
        h_flag, h_reason = heuristic_hard_flag(feats)
        h_score = heuristic_score(feats)
        mdl_phish_like = sum(s["prob"] for s in scores_sorted if s["label"] in {"phishing","malware","defacement"})
        fused = 0.50 * mdl_phish_like + 0.50 * h_score

        # allowlist override (domain-based)
        allow_hit = regdom in allowset if regdom else False
        decision = "🛑 UNSAFE"
        reasons = (h_reason + (", " if h_reason else "") + heuristic_reasons(feats)).strip(", ")

        if allow_hit and allowlist_override:
            decision = "✅ SAFE"
            reasons = f"allowlisted domain ({regdom})"
            fused = min(fused, 0.01)  # clamp down the risk for display
        else:
            decision = "🛑 UNSAFE" if (h_flag or fused >= float(threshold)) else "✅ SAFE"

        if decision.startswith("🛑"):
            any_unsafe = True

        rows.append([u, top["label"], top["prob"]*100.0, h_score, fused, decision, reasons])

        # export + forensics
        hidden_states = out.hidden_states
        cls_vec = hidden_states[-1][0, 0, :].cpu().tolist()
        export_data["items"].append({
            "url": u, "token_ids": token_ids, "tokens": tokens, "truncated": truncated,
            "logits": [float(x) for x in logits.cpu().tolist()], "probs": [float(p) for p in probs],
            "scores_sorted": scores_sorted, "cls_vector": cls_vec, "cls_dim": len(cls_vec),
            "elapsed_sec": elapsed, "heuristic": feats, "heuristic_score": h_score,
            "fused_risk": fused, "hard_flag": h_flag, "hard_reason": h_reason,
            "allowlisted": allow_hit
        })

        if forensic:
            forensic_blocks.append(
                _forensic_block(u, token_ids, tokens, scores_sorted, cls_vec, elapsed, truncated)
            )

    verdict = "🔴 **UNSAFE (at least one link flagged)**" if any_unsafe else "🟢 **SAFE (no link over threshold)**"
    body = verdict + "\n\n" + _results_table(rows)

    if forensic and forensic_blocks:
        body += "\n\n---\n\n" + "\n\n---\n\n".join(forensic_blocks)

    if show_json:
        pretty = json.dumps(export_data, ensure_ascii=False, indent=2)
        body += "\n\n---\n\n**Raw forensics JSON (copy & save):**\n"
        body += "```json\n" + pretty + "\n```"

    return body

# ---------- UI ----------
demo = gr.Interface(
    fn=analyze,
    inputs=[
        gr.Textbox(lines=10, label="Email or URL", placeholder="Paste a URL or a full email…"),
        gr.Checkbox(label="Forensic mode (tokens, logits, [CLS])", value=True),
        gr.Checkbox(label="Show raw JSON at the end (copy/paste)", value=False),
        gr.Slider(0.0, 1.0, value=0.40, step=0.01, label="Decision threshold (fused risk ≥ threshold → UNSAFE)"),
        gr.Textbox(lines=2, label="Allowlist (domains, comma/space/newline separated)",
                   placeholder="example.com, github.com  microsoft.com"),
        gr.Checkbox(label="Allowlist overrides (force SAFE if registered domain matches)", value=True),
    ],
    outputs=gr.Markdown(label="Results"),
    title="🛡️ PhishingMail — Model + Heuristics (HF Free CPU)",
    description=(
        "Extract links, score with a tiny HF URL model and transparent heuristics. "
        "Short-circuits for classic phishing patterns. Adjust the threshold, and allowlist trusted domains."
    ),
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)