"""Deterministic training-data generator for the email-risk classifier.

Every example originates from a list the engine already trusts:
  DISPOSABLE <- disposable_email_domains.blocklist (100k+ curated domains)
  ROLE       <- ROLE_BASED_PREFIXES (curated role local parts)
  NORMAL     -> synthetic name-like local parts on REAL MX-having domains -
                never blocklist entries
The model's job is to GENERALIZE these lists to domains/locals they do not
contain, so the training data must mirror exactly what the lists assert.
Labels: 0=NORMAL, 1=DISPOSABLE, 2=ROLE, 3=SPAM_TRAP.
"""
from __future__ import annotations

import random
from typing import Iterable

LABEL_NAMES = ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"]
NORMAL, DISPOSABLE, ROLE, SPAM_TRAP = 0, 1, 2, 3

# Honeypot-style local parts - tiny class, since real trap lists are private.
TRAP_LOCALS = ("spamtrap", "honeypot", "abuse", "blackhole")


def label_name(label: int) -> str:
    return LABEL_NAMES[label]


def _local_variants(seeds: Iterable[str], limit: int, rng: random.Random) -> list[str]:
    """Name-like local parts derived from seeds (dotted, joined, initial)."""
    out: list[str] = []
    seen: set[str] = set()
    for s in seeds:
        s = s.strip().lower()
        if not s:
            continue
        parts = s.split(".")
        variants = [s]
        if len(parts) > 1:
            variants.append(parts[0] + parts[-1])
            variants.append(f"{parts[0][0]}.{parts[-1]}")
        for v in variants:
            if v and v not in seen:
                seen.add(v)
                out.append(v)
        if len(out) >= limit:
            break
    rng.shuffle(out)
    return out[:limit]


def build_dataset(
    *,
    blocklist: set,
    free: set,
    role_prefixes: set,
    normal_domains: set,
    normal_locals: set,
    per_class_limit: int = 20000,
    seed: int = 42,
) -> list[dict]:
    """Build labeled rows ``{"text": "<local> [SEP] <domain>", "label": int}``.

    ``per_class_limit`` caps each class so the dataset stays balanced even
    though the disposable blocklist alone has 100k+ entries.
    """
    rng = random.Random(seed)
    rows: list[dict] = []

    def add(label: int, texts: Iterable[str]) -> None:
        texts = list(texts)
        rng.shuffle(texts)
        for t in texts[:per_class_limit]:
            rows.append({"text": t, "label": label})

    normal_domains = sorted(d for d in normal_domains if d and "." in d)
    free = sorted(f for f in free if f and "." in f)
    role_prefixes = sorted(p for p in role_prefixes if p)
    locals_ = _local_variants(normal_locals, 64, rng)
    if not locals_:
        locals_ = ["user"]

    # NORMAL: name-like locals on real MX-having domains.
    add(NORMAL, [f"{loc} [SEP] {d}" for d in normal_domains for loc in locals_])

    # DISPOSABLE: real blocklist domains behind name-like locals.
    disp = sorted(d for d in blocklist if d and "." in d)
    add(DISPOSABLE, [f"{rng.choice(locals_)} [SEP] {d}" for d in disp])

    # ROLE: curated role prefixes on normal + free domains.
    role_domains = normal_domains + free or ["example.test"]
    add(ROLE, [f"{p}{rng.choice(['', '1', '2024'])} [SEP] {d}"
               for p in role_prefixes for d in role_domains])

    # SPAM_TRAP: honeypot-style locals (tiny; real trap lists are private).
    add(SPAM_TRAP, [f"{t} [SEP] {d}" for t in TRAP_LOCALS for d in normal_domains])

    return rows


def write_csv(rows: list[dict], path: str) -> None:
    """Persist rows as a two-column CSV (text,label) for datasets.load_dataset."""
    import csv

    with open(path, "w", encoding="utf-8", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["text", "label"])
        for r in rows:
            writer.writerow([r["text"], r["label"]])