src/data_loader.py

"""
Load UCI SMS Spam (HuggingFace), phishing URL feeds (OpenPhish + PhishTank JSON),
manual digital_arrest_labels.csv — merge, clean, save data/processed/combined.csv.
"""

from __future__ import annotations

import json
import logging
import re
import urllib.request
from pathlib import Path

import pandas as pd

from datasets import load_dataset

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

# ── Paths ────────────────────────────────────────────────────────────────────
ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = ROOT / "data" / "raw"
PROC_DIR = ROOT / "data" / "processed"
PROC_DIR.mkdir(parents=True, exist_ok=True)
RAW_DIR.mkdir(parents=True, exist_ok=True)

MANUAL_CSV = RAW_DIR / "digital_arrest_labels.csv"
OUTPUT_CSV = PROC_DIR / "combined.csv"

OPENPHISH_FEED = "https://openphish.com/feed.txt"
# PhishTank periodically changes or restricts this endpoint; try fallbacks.
PHISHTANK_JSON_URLS = (
    "http://data.phishtank.com/data/online-valid.json",
    "https://data.phishtank.com/data/online-valid.json",
)

_HTTP_HEADERS = {
    "User-Agent": "Mozilla/5.0 (compatible; VES-scam-research/1.0; +https://github.com/)",
}

# Optional: normalize category strings (documentation only; any category allowed)
CATEGORY_MAP = {
    "safe": "safe",
    "digital_arrest": "digital_arrest",
    "otp_fraud": "otp_fraud",
    "courier_scam": "courier_scam",
    "phishing": "phishing",
    "spam": "spam",
}

_URL_RE = re.compile(r"https?://\S+|www\.\S+")
_PHONE_RE = re.compile(r"(\+91[\-\s]?)?[6-9]\d{9}|\b\d{10}\b|\b\d{5}[\s\-]\d{5}\b")
_EXTRA_RE = re.compile(r"\s{2,}")


def _http_get(url: str, timeout: int = 60) -> bytes:
    req = urllib.request.Request(url, headers=_HTTP_HEADERS)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read()


# ── 1. HuggingFace UCI SMS Spam (hub id used in examples: ucirvine/sms_spam) ──
def load_sms_spam() -> pd.DataFrame:
    """Load UCI SMS Spam Collection from Hugging Face (sms / label → text / label, 0=safe 1=scam)."""
    log.info("Loading SMS Spam (UCI) from HuggingFace: ucirvine/sms_spam …")
    try:
        ds = load_dataset("ucirvine/sms_spam", split="train")
        df = ds.to_pandas()
        if "sms" not in df.columns:
            raise ValueError(f"Expected column 'sms', got: {list(df.columns)}")

        df = df.rename(columns={"sms": "text"})
        # label: 0 ham (safe), 1 spam (scam); handle string labels if present
        if df["label"].dtype == object:
            lam = str.lower
            df["label"] = df["label"].map(lambda x: 0 if lam(str(x)) in ("ham", "0", "safe") else 1)
        else:
            df["label"] = df["label"].astype(int)

        df["category"] = df["label"].map({0: "safe", 1: "spam"})
        df = df[["text", "label", "category"]]
        scam = int(df["label"].sum())
        log.info("  SMS Spam: %d rows | scam=%d, safe=%d", len(df), scam, len(df) - scam)
        return df
    except Exception as e:
        log.error("  Failed to load SMS Spam dataset: %s", e)
        return pd.DataFrame(columns=["text", "label", "category"])


# ── 2a. OpenPhish — public text feed (one URL per line, not JSON) ────────────
def _collect_openphish_urls(limit: int) -> list[str]:
    cache = RAW_DIR / "openphish_feed.txt"
    try:
        if not cache.exists():
            log.info("  Downloading OpenPhish feed …")
            data = _http_get(OPENPHISH_FEED)
            cache.write_bytes(data)
        text = cache.read_text(encoding="utf-8", errors="ignore")
        urls = [ln.strip() for ln in text.splitlines() if ln.strip().startswith("http")]
        return urls[:limit]
    except Exception as e:
        log.warning("  OpenPhish feed failed (%s).", e)
        return []


# ── 2b. PhishTank — public JSON feed (hosted at phishtank.com, not openphish.com)
def _collect_phishtank_urls(limit: int) -> list[str]:
    cache = RAW_DIR / "phishtank_online_valid.json"
    try:
        if not cache.exists():
            last_err: Exception | None = None
            data: bytes | None = None
            for url in PHISHTANK_JSON_URLS:
                try:
                    log.info("  Downloading PhishTank JSON: %s …", url)
                    data = _http_get(url, timeout=120)
                    break
                except Exception as e:
                    last_err = e
                    log.warning("  PhishTank URL failed (%s): %s", url, e)
            if data is None:
                raise last_err or RuntimeError("PhishTank download failed")
            cache.write_bytes(data)
        raw = json.loads(cache.read_text(encoding="utf-8", errors="ignore"))
        urls: list[str] = []
        if isinstance(raw, list):
            for row in raw:
                if isinstance(row, dict) and row.get("url"):
                    urls.append(str(row["url"]).strip())
        elif isinstance(raw, dict):
            for row in raw.get("phish", []) or raw.get("data", []) or []:
                if isinstance(row, dict) and row.get("url"):
                    urls.append(str(row["url"]).strip())
        return urls[:limit]
    except Exception as e:
        log.warning("  PhishTank JSON failed (%s).", e)
        return []


def load_phishing_templates(
    openphish_max: int = 2000,
    phishtank_max: int = 800,
) -> pd.DataFrame:
    """
    Build synthetic phishing *messages* from malicious URLs (OpenPhish + PhishTank).
    Note: OpenPhish provides a plain-text URL list; PhishTank provides JSON.
    """
    log.info("Loading phishing URL feeds (OpenPhish + PhishTank) …")
    seen: set[str] = set()
    urls: list[str] = []
    for u in _collect_openphish_urls(openphish_max) + _collect_phishtank_urls(phishtank_max):
        if u not in seen:
            seen.add(u)
            urls.append(u)

    templates = [
        "Your account has been compromised. Verify immediately: {url}",
        "Action required: your KYC is pending. Click here: {url}",
        "You have a pending refund of ₹4,500. Claim now: {url}",
        "URGENT: Your SBI account will be blocked. Update: {url}",
        "Your parcel is on hold. Pay ₹299 customs: {url}",
    ]

    rows = []
    for i, url in enumerate(urls[:2000]):
        template = templates[i % len(templates)]
        rows.append({"text": template.format(url=url), "label": 1, "category": "phishing"})

    df = pd.DataFrame(rows)
    log.info("  Phishing templates: %d rows (from %d unique URLs)", len(df), len(urls))
    return df


# ── 3. Manual CSV ─────────────────────────────────────────────────────────────
def load_manual_labels() -> pd.DataFrame:
    """Load data/raw/digital_arrest_labels.csv (text, label int, category)."""
    if not MANUAL_CSV.exists():
        log.warning("  %s not found — creating starter template.", MANUAL_CSV)
        starter = pd.DataFrame(
            [
                {
                    "text": (
                        "This is CBI officer Rajesh Sharma. You are under digital arrest for money laundering. "
                        "Do not disconnect or we will send police to your address."
                    ),
                    "label": 1,
                    "category": "digital_arrest",
                },
                {
                    "text": (
                        "Narcotics Control Bureau has registered a case against your Aadhaar. "
                        "You must stay on this call. Disconnecting is a criminal offence."
                    ),
                    "label": 1,
                    "category": "digital_arrest",
                },
                {
                    "text": (
                        "Your FedEx parcel from China has been seized at customs. "
                        "Pay ₹2,400 to release it immediately. Call back on this number."
                    ),
                    "label": 1,
                    "category": "courier_scam",
                },
                {
                    "text": (
                        "OTP has been sent to your number for KYC verification. "
                        "Please share the OTP with me to complete your bank verification."
                    ),
                    "label": 1,
                    "category": "otp_fraud",
                },
                {"text": "Hi Priya, are you coming to college tomorrow? Let me know!", "label": 0, "category": "safe"},
                {
                    "text": "Your Amazon order #408-2938 has been shipped and will arrive by Friday.",
                    "label": 0,
                    "category": "safe",
                },
            ]
        )
        starter.to_csv(MANUAL_CSV, index=False)
        log.info("  Starter template written to %s", MANUAL_CSV)

    log.info("Loading manual labels from %s …", MANUAL_CSV)
    try:
        df = pd.read_csv(MANUAL_CSV)
        required = {"text", "label", "category"}
        missing = required - set(df.columns)
        if missing:
            raise ValueError(f"Manual CSV missing columns: {missing}")
        df["label"] = df["label"].astype(int)
        df["category"] = df["category"].astype(str).str.strip().str.lower()
        df = df[df["text"].astype(str).str.strip().astype(bool)].reset_index(drop=True)
        log.info(
            "  Manual labels: %d rows | categories: %s",
            len(df),
            df["category"].value_counts().to_dict(),
        )
        return df[["text", "label", "category"]]
    except Exception as e:
        log.error("  Failed to load manual labels: %s", e)
        return pd.DataFrame(columns=["text", "label", "category"])


def clean_text(text: str) -> str:
    if not isinstance(text, str):
        return ""
    text = _URL_RE.sub("[URL]", text)
    text = _PHONE_RE.sub("[PHONE]", text)
    text = _EXTRA_RE.sub(" ", text)
    return text.strip()


def build_dataset(
    save: bool = True,
    undersample_safe: bool = True,
    max_safe_ratio: float = 2.5,
) -> pd.DataFrame:
    frames = [
        load_sms_spam(),
        load_phishing_templates(),
        load_manual_labels(),
    ]
    df = pd.concat([f for f in frames if not f.empty], ignore_index=True)

    log.info("Cleaning text …")
    df["text"] = df["text"].astype(str).apply(clean_text)
    df = df[df["text"].str.len() >= 10].reset_index(drop=True)

    before = len(df)
    df = df.drop_duplicates(subset=["text"]).reset_index(drop=True)
    log.info("Deduplication: %d → %d rows", before, len(df))

    if undersample_safe:
        scam_count = int(df["label"].sum())
        safe_cap = int(scam_count * max_safe_ratio)
        safe_df = df[df["label"] == 0]
        scam_df = df[df["label"] == 1]
        if len(safe_df) > safe_cap:
            safe_df = safe_df.sample(n=safe_cap, random_state=42)
            df = pd.concat([scam_df, safe_df], ignore_index=True)
            df = df.sample(frac=1, random_state=42).reset_index(drop=True)
            log.info("Undersampled safe class to %d (max ratio %.1f:1)", safe_cap, max_safe_ratio)

    return df if not save else _save_and_log(df)


def _save_and_log(df: pd.DataFrame) -> pd.DataFrame:
    total = len(df)
    scam_n = int(df["label"].sum())
    safe_n = total - scam_n
    scam_pct = 100.0 * scam_n / total if total else 0.0
    safe_pct = 100.0 * safe_n / total if total else 0.0

    log.info("─" * 50)
    log.info("Final dataset: %d samples", total)
    log.info("  Scam : %d (%.1f%%)", scam_n, scam_pct)
    log.info("  Safe : %d (%.1f%%)", safe_n, safe_pct)
    log.info("Categories:\n%s", df["category"].value_counts().to_string())
    log.info("─" * 50)

    df.to_csv(OUTPUT_CSV, index=False)
    log.info("Saved → %s", OUTPUT_CSV)
    return df


if __name__ == "__main__":
    out = build_dataset(save=True)
    total = len(out)
    scam_n = int(out["label"].sum())
    safe_n = total - scam_n
    scam_pct = 100.0 * scam_n / total if total else 0.0
    safe_pct = 100.0 * safe_n / total if total else 0.0
    print(f"Loaded {total} samples | Scam: {scam_pct:.0f}% | Safe: {safe_pct:.0f}%")

    if total > 0:
        print("\nSample scam rows:")
        scam_df = out[out["label"] == 1]
        if len(scam_df) >= 1:
            n = min(3, len(scam_df))
            print(scam_df[["text", "category"]].sample(n, random_state=0).to_string(index=False))
        print("\nSample safe rows:")
        safe_df = out[out["label"] == 0]
        if len(safe_df) >= 1:
            n = min(3, len(safe_df))
            print(safe_df[["text", "category"]].sample(n, random_state=0).to_string(index=False))